""" zone_observation.py =================== """ from __future__ import annotations import json import logging import math import random import zlib from dataclasses import dataclass, field, asdict from datetime import datetime, timedelta, timezone from enum import Enum, unique from typing import Any, ClassVar, Dict, List, Optional, Tuple logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Schema version # --------------------------------------------------------------------------- # v3 migration note: # - ZoneObs gained two optional per-zone fields: precip_satellite_mm, # soil_moisture_satellite_pct (both Optional[float], default None — # None means "no satellite pass available", NOT zero). # - New BasinContext dataclass added (episode-level, NOT per-zone): # enso_oni, iod_dmi, itcz_latitude_deg, mslp_regional_hpa. # EpisodeContext gained an optional `basin_context` field. # - DataSource gained SATELLITE_PRECIP, SATELLITE_SOIL, PUBLISHED_INDEX. # - BasinContext later gained optional helio / space-weather fields # (solar_wind_speed_kms, kp_index, goes_xray_flux, helio_regime) with # quiet-Sun defaults so old records deserialise without a schema bump. # - ForecastConfig gained require_real_basin_context for strict real-only # basin/helio ingest (enforced in era5_data_pipeline.fetch_basin_context). # Old v2 records deserialise fine via ZoneObs.from_dict/EpisodeContext.from_dict # as long as the caller bumps stored "_schema_version" to 3 first (the new # fields all have defaults, so no other migration is required). SCHEMA_VERSION: int = 3 # --------------------------------------------------------------------------- # Known extras keys # --------------------------------------------------------------------------- KNOWN_EXTRAS: Dict[str, str] = { "gdd_base_c": "float -- crop-specific GDD base temperature (degrees C)", "crop_substage": "str -- variety-specific growth substage", "export_grade_risk": "float -- pre-computed quality risk from field notes [0,1]", "edge_node_id": "str -- edge sensor network node that sourced this observation", "contract_volume_mt": "float -- contracted volume for this zone (metric tonnes)", "soil_type": "str -- FAO soil classification string", "irrigation_source": "str -- 'rainfed' | 'irrigated' | 'supplemental'", "sar_flood_date": "str -- ISO-8601 date of most recent SAR flood detection", } # --------------------------------------------------------------------------- # Enumerations # --------------------------------------------------------------------------- @unique class CropStage(Enum): """Generalised crop growth stage. Kept coarse deliberately — variety-specific substages can be added via ZoneObs.extras['crop_substage'] without a schema bump. """ UNKNOWN = "unknown" LAND_PREP = "land_prep" # tillage, flooding (rice), bed preparation PLANTING = "planting" # transplanting / direct seeding VEGETATIVE = "vegetative" # tillering (rice), canopy closure REPRODUCTIVE = "reproductive" # booting -> heading -> flowering GRAIN_FILLING = "grain_filling" # dough stage -- highest moisture risk MATURATION = "maturation" # drying down, harvest window opens HARVEST = "harvest" # active harvest, logistics pressure FALLOW = "fallow" # between seasons @unique class AlertLevel(Enum): """Risk alert level emitted by the RL policy's terminate action.""" NONE = "none" # no alert warranted WATCH = "watch" # monitor closely -- conditions developing ADVISORY = "advisory" # elevated risk -- recommend pre-emptive action WARNING = "warning" # high probability of supply/quality impact CRITICAL = "critical" # immediate action required def severity(self) -> int: """Integer severity: NONE=0, WATCH=1, ADVISORY=2, WARNING=3, CRITICAL=4.""" return {"none": 0, "watch": 1, "advisory": 2, "warning": 3, "critical": 4}[self.value] def __lt__(self, other: "AlertLevel") -> bool: # type: ignore[override] return self.severity() < other.severity() def __le__(self, other: "AlertLevel") -> bool: # type: ignore[override] return self.severity() <= other.severity() def __gt__(self, other: "AlertLevel") -> bool: # type: ignore[override] return self.severity() > other.severity() def __ge__(self, other: "AlertLevel") -> bool: # type: ignore[override] return self.severity() >= other.severity() @unique class DataSource(Enum): """Provenance tag so downstream code can weight or distrust readings. Migration note (schema v2): DataSource.EDGE_NODE serialises as "edge_node". Any records previously serialised as "gnus_node" must be migrated: d["source"] = "edge_node" # was "gnus_node" """ ERA5_REANALYSIS = "era5_reanalysis" # ECMWF ERA5 via CDS or Open-Meteo OPENMETEO_LIVE = "openmeteo_live" # Open-Meteo forecast API (free tier) BMKG_STATION = "bmkg_station" # Indonesian met agency station data SATELLITE_NDVI = "satellite_ndvi" # Sentinel-2 / Landsat NDVI tile SATELLITE_PRECIP = "satellite_precip" # IMERG / CHIRPS retrieval (schema v3+) SATELLITE_SOIL = "satellite_soil" # SMAP L3/L4 retrieval (schema v3+) PUBLISHED_INDEX = "published_index" # NOAA/BOM basin-scale index, e.g. ONI/DMI (v3+) EDGE_NODE = "edge_node" # Distributed edge sensor network node SYNTHETIC = "synthetic" # Generated by make_synthetic_* for training UNKNOWN = "unknown" def is_observational(self) -> bool: """True for real-world sources (not synthetic or unknown).""" return self not in (DataSource.SYNTHETIC, DataSource.UNKNOWN) # --------------------------------------------------------------------------- # Utilities # --------------------------------------------------------------------------- def _clip(value: float, lo: float, hi: float) -> float: """stdlib-only clip. Avoids numpy dependency at this layer.""" return max(lo, min(hi, value)) def _stable_seed(key: str) -> int: """Deterministic integer seed from a string key. Uses zlib.crc32 rather than hash() -- hash() is randomised per process by PYTHONHASHSEED and would produce different synthetic obs for the same zone_id across runs, breaking training reproducibility. """ return zlib.crc32(key.encode("utf-8")) & 0x7FFFFFFF def _copy_and_pop_schema(d: Dict[str, Any]) -> Tuple[Optional[int], Dict[str, Any]]: """Return (schema_version, clean_copy) without mutating the input dict. All from_dict() methods call this instead of d.pop() directly, which was the v1 caller-mutation bug. """ d_copy = dict(d) sv = d_copy.pop("_schema_version", None) return sv, d_copy def _check_schema(sv: Optional[int], class_name: str) -> None: """Raise ValueError loudly on schema version mismatch.""" if sv is not None and sv != SCHEMA_VERSION: raise ValueError( f"{class_name}.from_dict: schema version mismatch -- " f"stored={sv}, current={SCHEMA_VERSION}. " f"Run migration script or increment SCHEMA_VERSION." ) # --------------------------------------------------------------------------- # ForecastConfig # --------------------------------------------------------------------------- @dataclass class ForecastConfig: """ Validated configuration for the weather forecasting RL environment. Mirrors InspectionConfig from the MEMS env: every economic parameter is validated in __post_init__, and the rational termination threshold is computed and checked explicitly so misconfiguration is loud. Rational termination condition (mirrors MEMS belief_floor logic): Issue alert when: P(risk_event) > false_alert_penalty / (alert_value + false_alert_penalty) With defaults: P > 20 / (100 + 20) = 0.1667 belief_floor MUST be set below this threshold or early termination can never be EV-positive. """ # --- Spatial --- n_zones: int = 1 # number of sourcing zones per episode horizon_days: int = 30 # forecast horizon (days) max_steps: int = 200 # max env steps per episode # --- Belief map --- prior_belief: float = 0.12 # initial P(risk_event) per zone belief_floor: float = 0.005 # minimum belief after decay belief_update_radius: int = 2 # spatial propagation radius (zone cells) belief_increase_rate: float = 0.30 # update magnitude when event confirmed belief_decrease_rate: float = 0.05 # update magnitude when event absent # --- Economics --- alert_value: float = 100.0 # reward for correct advisory issuance false_alert_penalty: float = 20.0 # penalty for unnecessary advisory miss_penalty: float = 200.0 # penalty per missed true risk event inspection_cost: float = 1.0 # cost per zone-step (resource use) # --- Exploration shaping --- # FIX (premature termination): ep_len_mean stayed at ~1.3-1.6 for the # entire 'normal' (n_zones=2) curriculum phase across 1M steps -- not # slow convergence, the actual argmax of the reward as specified. # Root cause: belief_map is already ~30%-informed by the true signal at # reset() (see _per_zone_beliefs' 0.7*prior + 0.3*signal blend), so the # terminate action's EV barely depends on whether the agent inspected # anything. Meanwhile a single inspect step costs inspection_cost=1.0 # raw, while its info_gain term (5.0 * info_gain in # _compute_zone_refinement_reward) rarely exceeds ~0.2-0.4 raw for # realistic belief deltas, and the uncertainty-decay payoff at # termination is a mean over ALL zones, diluting the effect of # inspecting any single one. Net EV of inspecting was negative under # almost all conditions. # # These two fields close that gap directly, without touching the core # alert/false-alarm economics that crop_risk_scorer.py's threshold # logic depends on: # zone_visit_bonus: added to the raw inspection reward on a zone's # FIRST visit only (the existing revisit_penalty path is untouched, # so repeat visits are still never profitable). Set > inspection_cost # so a first visit is always non-negative in expectation, independent # of info_gain -- exploration is never structurally punished. # unvisited_zone_penalty: subtracted from the termination reward, # scaled by the fraction of active zones never visited this episode. # Terminating having inspected nothing costs the full penalty; # terminating after visiting every zone costs nothing. Uses the same # /alert_value normalization as every other termination-reward term, # so it composes on the same scale as gain/cost/unc. zone_visit_bonus: float = 1.5 # raw bonus, first visit to a zone only unvisited_zone_penalty: float = 40.0 # raw penalty * (unvisited/active) at terminate # --- Robustness training --- economic_randomization: bool = False clean_episode_ratio: float = 0.7 # Spatial correlation of synthetic hazard events across zones in one # episode. 0 = independent per-zone draws (legacy); 1 = all-or-nothing # regional event (El Niño-style). Real Java drought/flood is highly # correlated; default 0.85 so multi-zone allocation sees joint events. event_spatial_correlation: float = 0.85 # If True (default), training reset permutes zone slot order each episode # so action index i is not permanently tied to zone_i / a fixed name. # Prevents policies from learning "always inspect slot 1" by index alone. # Real-eval injection path does not shuffle (order is caller-defined). shuffle_zone_order: bool = True alert_value_range: Tuple[float, float] = (0.8, 1.2) miss_penalty_range: Tuple[float, float] = (0.9, 1.1) # --- Soft reset --- soft_reset: bool = True # --- Seeding --- seed: Optional[int] = None # --- Data pipeline (era5_data_pipeline.py) --- # These fields are consumed by fetch_zone_obs() to control source selection # and noise injection. They must live here so user-set values are respected; # era5_data_pipeline.py previously fell back to getattr() defaults, silently # ignoring any values set on a ForecastConfig instance. real_data_ratio: float = 0.7 # fraction of episodes that attempt real data 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 noise_scale: float = 0.05 # noise magnitude (fraction of field range) # --- Satellite sources (schema v3+) --- # Opt-in and default False: existing callers get byte-identical behaviour # unless they explicitly enable these. When enabled, era5_data_pipeline.py # tries the satellite source first (via Google Earth Engine) and falls # back through the existing ERA5 -> Open-Meteo -> synthetic chain on # failure, exactly like every other source in fetch_zone_obs(). use_satellite_precip: bool = False # prefer IMERG/CHIRPS over ERA5 precip use_satellite_soil: bool = False # prefer SMAP over ERA5 soil moisture include_basin_context: bool = False # attach ENSO/IOD/monsoon/helio context to EpisodeContext # Strict real-only basin context: when True, fetch_basin_context() refuses # synthetic / quiet-default fallbacks and raises RuntimeError if any of # RONI, DMI, or the required SWPC helio fields cannot be obtained from # live published sources. Default False preserves the soft-fallback path # used by training and offline runs. Implies include_basin_context in # spirit (callers should set both), but the fetch function itself is the # enforcement point. require_real_basin_context: bool = False # --- Forecast backend (timesfm_wrapper.py integration) --- # Which ForecastBackend fetch_episode_context() uses to build the # EpisodeContext's ForecastResult. Default 'synthetic' reproduces the # pre-integration behaviour exactly (deterministic synthetic forecast), # so existing callers see no change unless they opt in. # 'synthetic' -- deterministic make_synthetic_forecast_result (offline) # 'baseline' -- deterministic statistical forecast from ZoneObs fields # (timesfm_wrapper.BaselineBackend; offline, no deps) # 'openmeteo' -- real Open-Meteo API forecast (needs network; falls # back to synthetic on any failure) # 'timesfm' -- local TimesFM checkpoint (needs the checkpoint file + # timesfm package; see timesfm_wrapper.LocalTimesFMBackend) forecast_backend: str = "synthetic" # --- Climatology anomalies (climatology.py integration) --- # Opt-in, default False for byte-identical back-compat. When True, # era5_data_pipeline.fetch_zone_obs() post-processes every fetched # ZoneObs through climatology.apply_climatology_anomalies(), which # populates precip_anomaly_idx / temp_anomaly_idx / soil_moisture_anom # as z-scores against a per-zone day-of-year climatology. Without this, # ALL real fetchers leave those fields at 0.0, which zeroes out most of # ZoneObs.drought_signal()/flood_signal() -- i.e. real observations are # barely scoreable. Synthetic training episodes deliberately leave this # off (their anomaly fields are injected directly by the event flags). use_climatology_anomalies: bool = False climatology_years: int = 10 # years of history for the climatology def __post_init__(self) -> None: if self.alert_value <= 0: raise ValueError(f"ForecastConfig: alert_value={self.alert_value} must be > 0") if self.false_alert_penalty < 0: raise ValueError(f"ForecastConfig: false_alert_penalty must be >= 0") if self.miss_penalty <= 0: raise ValueError(f"ForecastConfig: miss_penalty must be > 0") if self.inspection_cost <= 0: raise ValueError(f"ForecastConfig: inspection_cost must be > 0") if self.zone_visit_bonus < 0: raise ValueError(f"ForecastConfig: zone_visit_bonus must be >= 0") if self.unvisited_zone_penalty < 0: raise ValueError(f"ForecastConfig: unvisited_zone_penalty must be >= 0") if self.horizon_days < 1: raise ValueError(f"ForecastConfig: horizon_days must be >= 1") if self.n_zones < 1: raise ValueError(f"ForecastConfig: n_zones must be >= 1") if self.max_steps < 1: raise ValueError(f"ForecastConfig: max_steps must be >= 1") if not (0.0 <= self.real_data_ratio <= 1.0): raise ValueError( f"ForecastConfig: real_data_ratio={self.real_data_ratio} must be in [0, 1]" ) if not (0.0 <= self.era5_ratio <= 1.0): raise ValueError( f"ForecastConfig: era5_ratio={self.era5_ratio} must be in [0, 1]" ) if self.noise_scale < 0.0: raise ValueError( f"ForecastConfig: noise_scale={self.noise_scale} must be >= 0" ) _VALID_BACKENDS = ("synthetic", "baseline", "openmeteo", "timesfm") if self.forecast_backend not in _VALID_BACKENDS: raise ValueError( f"ForecastConfig: forecast_backend={self.forecast_backend!r} " f"must be one of {_VALID_BACKENDS}" ) if not (1 <= self.climatology_years <= 30): raise ValueError( f"ForecastConfig: climatology_years={self.climatology_years} " f"must be in [1, 30]" ) self.alert_value = float(_clip(self.alert_value, 0.1, 10_000.0)) self.false_alert_penalty = float(_clip(self.false_alert_penalty, 0.0, 10_000.0)) self.miss_penalty = float(_clip(self.miss_penalty, 0.1, 100_000.0)) self.inspection_cost = float(_clip(self.inspection_cost, 0.01, 1_000.0)) self.zone_visit_bonus = float(_clip(self.zone_visit_bonus, 0.0, 1_000.0)) self.unvisited_zone_penalty = float(_clip(self.unvisited_zone_penalty, 0.0, 10_000.0)) self.prior_belief = float(_clip(self.prior_belief, 0.001, 0.999)) self.belief_floor = float(_clip(self.belief_floor, 0.001, 0.5)) self.clean_episode_ratio = float(_clip(self.clean_episode_ratio, 0.0, 1.0)) self.event_spatial_correlation = float( _clip(self.event_spatial_correlation, 0.0, 1.0) ) # Threshold is a 3-way EV break-even (alert vs. no-alert at believed # probability p), not just alert_value vs false_alert_penalty: # p > false_alert_penalty / (alert_value + false_alert_penalty + miss_penalty) # Single source of truth: crop_risk_scorer._alert_level and # weather_forecast_env._compute_termination_reward both read/derive # from this instead of keeping their own copy of the formula. rational = self.false_alert_penalty / max( self.alert_value + self.false_alert_penalty + self.miss_penalty, 1e-9 ) if self.belief_floor >= rational: logger.warning( f"ForecastConfig: belief_floor={self.belief_floor:.4f} >= " f"rational_termination_threshold={rational:.4f}. " f"Early termination will never be EV-positive. " f"Set belief_floor < {rational:.4f}." ) self._rational_threshold: float = rational @property def rational_termination_threshold(self) -> float: """P(event) above which issuing an alert has positive expected value, accounting for alert_value, false_alert_penalty, AND miss_penalty (3-way expected-value breakeven -- see the derivation in __post_init__). Single source of truth: crop_risk_scorer._alert_level and weather_forecast_env._compute_termination_reward both read this property rather than recomputing the formula. """ return self._rational_threshold def to_dict(self) -> Dict[str, Any]: return { "n_zones": self.n_zones, "horizon_days": self.horizon_days, "max_steps": self.max_steps, "prior_belief": self.prior_belief, "belief_floor": self.belief_floor, "belief_update_radius": self.belief_update_radius, "belief_increase_rate": self.belief_increase_rate, "belief_decrease_rate": self.belief_decrease_rate, "alert_value": self.alert_value, "false_alert_penalty": self.false_alert_penalty, "miss_penalty": self.miss_penalty, "inspection_cost": self.inspection_cost, "zone_visit_bonus": self.zone_visit_bonus, "unvisited_zone_penalty": self.unvisited_zone_penalty, "economic_randomization": self.economic_randomization, "clean_episode_ratio": self.clean_episode_ratio, "event_spatial_correlation": self.event_spatial_correlation, "shuffle_zone_order": self.shuffle_zone_order, "alert_value_range": list(self.alert_value_range), "miss_penalty_range": list(self.miss_penalty_range), "soft_reset": self.soft_reset, "seed": self.seed, "real_data_ratio": self.real_data_ratio, "era5_ratio": self.era5_ratio, "force_data_source": ( self.force_data_source.value if self.force_data_source is not None else None ), "inject_noise": self.inject_noise, "noise_scale": self.noise_scale, "use_satellite_precip": self.use_satellite_precip, "use_satellite_soil": self.use_satellite_soil, "include_basin_context": self.include_basin_context, "require_real_basin_context": self.require_real_basin_context, "forecast_backend": self.forecast_backend, "use_climatology_anomalies": self.use_climatology_anomalies, "climatology_years": self.climatology_years, "_schema_version": SCHEMA_VERSION, } @classmethod def from_dict(cls, d: Dict[str, Any]) -> "ForecastConfig": sv, d = _copy_and_pop_schema(d) _check_schema(sv, "ForecastConfig") if "alert_value_range" in d and isinstance(d["alert_value_range"], list): d["alert_value_range"] = tuple(d["alert_value_range"]) if "miss_penalty_range" in d and isinstance(d["miss_penalty_range"], list): d["miss_penalty_range"] = tuple(d["miss_penalty_range"]) if "force_data_source" in d and d["force_data_source"] is not None: d["force_data_source"] = DataSource(d["force_data_source"]) return cls(**{k: v for k, v in d.items() if not k.startswith("_")}) # --------------------------------------------------------------------------- # GeoPolygon # --------------------------------------------------------------------------- @dataclass class GeoPolygon: """Lightweight polygon -- no shapely dependency at this layer. Vertices are (lat, lon) pairs in decimal degrees, WGS-84. Vertex coordinates are coerced to float at construction so JSON-deserialized strings ('3.0') do not silently fail later. """ vertices: List[Tuple[float, float]] # [(lat degrees, lon degrees), ...] zone_id: str label: str = "" def __post_init__(self) -> None: self.vertices = [(float(v[0]), float(v[1])) for v in self.vertices] if len(self.vertices) < 3: raise ValueError( f"GeoPolygon '{self.zone_id}' needs >= 3 vertices, " f"got {len(self.vertices)}" ) for lat, lon in self.vertices: if not (-90.0 <= lat <= 90.0): raise ValueError( f"GeoPolygon '{self.zone_id}': latitude {lat} out of [-90, 90]" ) if not (-180.0 <= lon <= 180.0): raise ValueError( f"GeoPolygon '{self.zone_id}': longitude {lon} out of [-180, 180]" ) @property def centroid(self) -> Tuple[float, float]: """Arithmetic centroid (lat, lon).""" lats = [v[0] for v in self.vertices] lons = [v[1] for v in self.vertices] return (sum(lats) / len(lats), sum(lons) / len(lons)) @property def approx_area_km2(self) -> float: """Shoelace formula on a flat-earth approximation. Accurate to ~1% for zones < 200 km across. """ lat_c, _ = self.centroid km_per_deg_lat = 111.0 km_per_deg_lon = 111.0 * math.cos(math.radians(lat_c)) n = len(self.vertices) area = 0.0 for i in range(n): x0 = self.vertices[i][1] * km_per_deg_lon y0 = self.vertices[i][0] * km_per_deg_lat x1 = self.vertices[(i + 1) % n][1] * km_per_deg_lon y1 = self.vertices[(i + 1) % n][0] * km_per_deg_lat area += x0 * y1 - x1 * y0 return abs(area) / 2.0 def contains_point(self, lat: float, lon: float) -> bool: """Ray-casting point-in-polygon test (WGS-84 lat/lon).""" n = len(self.vertices) inside = False j = n - 1 for i in range(n): xi, yi = self.vertices[i][1], self.vertices[i][0] xj, yj = self.vertices[j][1], self.vertices[j][0] if ((yi > lat) != (yj > lat)) and ( lon < (xj - xi) * (lat - yi) / (yj - yi + 1e-12) + xi ): inside = not inside j = i return inside def to_dict(self) -> Dict[str, Any]: return { "vertices": [list(v) for v in self.vertices], "zone_id": self.zone_id, "label": self.label, } @classmethod def from_dict(cls, d: Dict[str, Any]) -> "GeoPolygon": return cls( vertices=[(float(v[0]), float(v[1])) for v in d["vertices"]], zone_id=d["zone_id"], label=d.get("label", ""), ) # --------------------------------------------------------------------------- # ZoneObs # --------------------------------------------------------------------------- @dataclass class ZoneObs: """ One observation snapshot for one sourcing zone at one point in time. The single object that crosses all layer boundaries. The pipeline produces it. The env consumes it. The model reads it. The scorer reads it. The transport serialises it. Field naming convention: *_mm = millimetres *_c = degrees Celsius *_pct = percentage 0-100 *_idx = dimensionless index (standardised anomaly or ratio) *_prob = probability 0.0-1.0 *_days = duration in days *_ms = metres per second *_anom = z-score anomaly vs climatological mean """ # --- Identity --- zone_id: str valid_time: datetime # UTC timestamp of observation window start source: DataSource = DataSource.UNKNOWN # --- Precipitation --- precip_24h_mm: float = 0.0 # total precip last 24 h (mm) precip_7d_mm: float = 0.0 # total precip last 7 days (mm) precip_14d_mm: float = 0.0 # total precip last 14 days (mm) precip_30d_mm: float = 0.0 # total precip last 30 days (mm) precip_anomaly_idx: float = 0.0 # z-score vs ERA5 climatological mean # negative = drought, positive = excess # --- Temperature --- temp_mean_c: float = 0.0 # daily mean (degrees C) temp_max_c: float = 0.0 # daily maximum (degrees C) temp_min_c: float = 0.0 # daily minimum (degrees C) temp_anomaly_idx: float = 0.0 # z-score vs climatological mean # --- Heat stress --- gdd_accumulated: float = 0.0 # growing degree-days since sowing # base temp is crop-specific -> extras['gdd_base_c'] heat_stress_days: int = 0 # days where temp_max_c > threshold (default 35 C) cold_stress_days: int = 0 # days where temp_min_c < threshold (default 15 C) # --- Soil / moisture --- soil_moisture_pct: float = 0.0 # volumetric water content top 10cm, 0-100 soil_moisture_anom: float = 0.0 # z-score vs climatological mean evapotranspiration_mm: float = 0.0 # reference ET0 (FAO-56 Penman-Monteith), mm/day # --- Direct satellite retrievals (schema v3+, per-zone, optional) --- # These are raw provenance/audit fields, NOT yet consumed by # crop_risk_scorer.py's drought_signal()/flood_signal(). Wiring them in # requires a per-zone climatological baseline (mean/std) to convert a raw # satellite reading into an anomaly comparable to precip_anomaly_idx / # soil_moisture_anom -- that baseline does not exist anywhere in this # pipeline yet. Do not fabricate one; compute it from a real climatology # (e.g. multi-year ERA5 or IMERG mean) before using these fields for risk # scoring. None means "no satellite pass available", NOT zero. precip_satellite_mm: Optional[float] = None # IMERG/CHIRPS daily total (mm) soil_moisture_satellite_pct: Optional[float] = None # SMAP L3/L4 retrieval (%, 0-100) # --- Wind --- wind_speed_max_ms: float = 0.0 # maximum gust in window (m/s) wind_speed_mean_ms: float = 0.0 # mean 10m wind speed (m/s) # --- Humidity --- rh_mean_pct: float = 0.0 # relative humidity daily mean, 0-100 rh_max_pct: float = 0.0 # daily maximum, 0-100; key fungi risk driver # FIX (fungi/regime-insensitivity, verified on a 3,623-point real-Indonesia # backtest): fungi_risk_signal() used rh_max_pct against a fixed absolute # threshold. In a tropical maritime climate rh_max_pct sits near # saturation (90-100%) on most days regardless of ENSO phase, so the # signal was flat across El Nino/La Nina and won the alert_level max() # in 73.6% of points -- masking the correctly regime-sensitive # drought/flood signals in the system's actual output (alert_level). # z-score vs climatological mean, built from rh_mean_pct (NOT rh_max_pct # -- Open-Meteo's archive API documents max-RH aggregation as not # reliably available across models, while mean-RH is; daily max and mean # RH anomalies are highly correlated, so mean's anomaly is used as the # proxy correction applied to the max-based absolute term). Default 0.0 # is a true no-op in fungi_risk_signal() (additive, clipped, zero when # unset) -- synthetic data and any caller that never populates this are # byte-identical to the pre-fix formula. Populated for real data by # climatology.apply_climatology_anomalies(). rh_anomaly_idx: float = 0.0 # --- Vegetation (optional -- populated when satellite pass available) --- ndvi: Optional[float] = None # NDVI -1.0 to 1.0; None if no recent pass ndvi_anomaly_idx: Optional[float] = None # z-score vs same-DOY climatology ndvi_trend_14d: Optional[float] = None # linear slope over 14 days (NDVI/day) # --- Flood / waterlogging --- flood_extent_pct: float = 0.0 # % of zone with standing water (SAR-derived), 0-100 drainage_risk_idx: float = 0.0 # composite: slope + soil type + recent precip, 0-1 # --- Crop context --- crop_stage: CropStage = CropStage.UNKNOWN days_to_harvest: Optional[int] = None # None = unknown; 0 = harvest now planting_date: Optional[datetime] = None # --- Observation quality --- quality_flag: int = 0 # 0=good, 1=interpolated, 2=gap-filled, 3=synthetic cloud_cover_pct: float = 0.0 # cloud fraction 0-100; high values degrade NDVI # --- Extensibility hatch --- extras: Dict[str, Any] = field(default_factory=dict) # See KNOWN_EXTRAS at module level for documented keys. def __post_init__(self) -> None: # --- Defensive copy: prevent caller mutation of extras dict --- self.extras = dict(self.extras) if self.valid_time.tzinfo is None: logger.warning( f"ZoneObs('{self.zone_id}'): valid_time has no timezone, assuming UTC." ) self.valid_time = self.valid_time.replace(tzinfo=timezone.utc) self.soil_moisture_pct = float(_clip(self.soil_moisture_pct, 0.0, 100.0)) self.rh_mean_pct = float(_clip(self.rh_mean_pct, 0.0, 100.0)) self.rh_max_pct = float(_clip(self.rh_max_pct, 0.0, 100.0)) self.flood_extent_pct = float(_clip(self.flood_extent_pct, 0.0, 100.0)) self.cloud_cover_pct = float(_clip(self.cloud_cover_pct, 0.0, 100.0)) self.drainage_risk_idx = float(_clip(self.drainage_risk_idx, 0.0, 1.0)) # --- Anomaly clipping: prevent z-score explosions from destabilising RL --- self.precip_anomaly_idx = float(_clip(self.precip_anomaly_idx, -5.0, 5.0)) self.temp_anomaly_idx = float(_clip(self.temp_anomaly_idx, -5.0, 5.0)) self.soil_moisture_anom = float(_clip(self.soil_moisture_anom, -5.0, 5.0)) self.rh_anomaly_idx = float(_clip(self.rh_anomaly_idx, -5.0, 5.0)) if self.ndvi is not None: self.ndvi = float(_clip(self.ndvi, -1.0, 1.0)) if self.soil_moisture_satellite_pct is not None: self.soil_moisture_satellite_pct = float( _clip(self.soil_moisture_satellite_pct, 0.0, 100.0) ) if self.precip_satellite_mm is not None and self.precip_satellite_mm < 0.0: logger.warning( f"ZoneObs('{self.zone_id}'): precip_satellite_mm=" f"{self.precip_satellite_mm:.4f} < 0, clipping to 0." ) self.precip_satellite_mm = 0.0 for attr in ( "precip_24h_mm", "precip_7d_mm", "precip_14d_mm", "precip_30d_mm", "evapotranspiration_mm", "wind_speed_max_ms", "wind_speed_mean_ms", "gdd_accumulated", ): val = getattr(self, attr) if val < 0.0: logger.warning( f"ZoneObs('{self.zone_id}'): {attr}={val:.4f} < 0, clipping to 0." ) setattr(self, attr, 0.0) if self.heat_stress_days < 0: self.heat_stress_days = 0 if self.cold_stress_days < 0: self.cold_stress_days = 0 if self.quality_flag not in (0, 1, 2, 3): logger.warning( f"ZoneObs('{self.zone_id}'): quality_flag={self.quality_flag} " f"not in {{0,1,2,3}}, setting to 3." ) self.quality_flag = 3 if self.planting_date is not None and self.planting_date.tzinfo is None: self.planting_date = self.planting_date.replace(tzinfo=timezone.utc) # --- Monotonic precipitation aggregates sanity check --- if not ( self.precip_30d_mm >= self.precip_14d_mm >= self.precip_7d_mm >= self.precip_24h_mm ): logger.warning( f"ZoneObs('{self.zone_id}'): non-monotonic precipitation aggregates " f"(24h={self.precip_24h_mm:.2f}, 7d={self.precip_7d_mm:.2f}, " f"14d={self.precip_14d_mm:.2f}, 30d={self.precip_30d_mm:.2f})" ) @classmethod def validate(cls, obs: "ZoneObs", strict: bool = False) -> List[str]: """Return a list of validation warnings (empty = clean). Called by era5_data_pipeline.py before inserting obs into the env. strict=True raises ValueError instead of returning issues. """ issues: List[str] = [] if not obs.zone_id: issues.append("zone_id is empty") if obs.temp_max_c < obs.temp_min_c: issues.append(f"temp_max_c={obs.temp_max_c} < temp_min_c={obs.temp_min_c}") if obs.precip_14d_mm < obs.precip_7d_mm: issues.append(f"precip_14d_mm < precip_7d_mm") if obs.precip_30d_mm < obs.precip_14d_mm: issues.append(f"precip_30d_mm < precip_14d_mm") if obs.wind_speed_max_ms < obs.wind_speed_mean_ms: issues.append(f"wind_speed_max_ms < wind_speed_mean_ms") if obs.rh_max_pct < obs.rh_mean_pct: issues.append(f"rh_max_pct < rh_mean_pct") if obs.days_to_harvest is not None and obs.days_to_harvest < 0: issues.append(f"days_to_harvest={obs.days_to_harvest} < 0") if obs.quality_flag >= 2 and obs.source.is_observational(): issues.append( f"quality_flag={obs.quality_flag} (gap-filled/synthetic) " f"but source={obs.source.value} is observational" ) if strict and issues: raise ValueError(f"ZoneObs('{obs.zone_id}') strict validation failed: {issues}") return issues def to_dict(self) -> Dict[str, Any]: """Lossless serialisation to JSON-safe dict. Never mutates self.""" d = asdict(self) d["source"] = self.source.value d["crop_stage"] = self.crop_stage.value d["valid_time"] = self.valid_time.isoformat() d["planting_date"] = self.planting_date.isoformat() if self.planting_date else None d["_schema_version"] = SCHEMA_VERSION return d @classmethod def from_dict(cls, d: Dict[str, Any]) -> "ZoneObs": """Reconstruct from to_dict(). Never mutates the caller's dict.""" sv, d = _copy_and_pop_schema(d) _check_schema(sv, "ZoneObs") d["valid_time"] = datetime.fromisoformat(d["valid_time"]) d["source"] = DataSource(d["source"]) d["crop_stage"] = CropStage(d["crop_stage"]) d["planting_date"] = ( datetime.fromisoformat(d["planting_date"]) if d.get("planting_date") else None ) return cls(**d) # --- Diagnostic helpers --- def is_harvest_window(self, lookahead_days: int = 21) -> bool: if self.days_to_harvest is None: return self.crop_stage in (CropStage.MATURATION, CropStage.HARVEST) return 0 <= self.days_to_harvest <= lookahead_days def has_reliable_ndvi(self) -> bool: return self.ndvi is not None and self.cloud_cover_pct < 30.0 def drought_signal(self) -> float: """Composite drought signal [0, 1]. Heuristic for diagnostics only.""" return float( 0.6 * _clip(-self.precip_anomaly_idx / 3.0, 0.0, 1.0) + 0.4 * _clip(-self.soil_moisture_anom / 3.0, 0.0, 1.0) ) def flood_signal(self) -> float: """Composite flood signal [0, 1].""" return float( 0.4 * _clip(self.precip_anomaly_idx / 3.0, 0.0, 1.0) + 0.4 * (self.flood_extent_pct / 100.0) + 0.2 * _clip(self.drainage_risk_idx, 0.0, 1.0) ) def fungi_risk_signal(self) -> float: """Composite fungi / post-harvest moisture risk [0, 1]. FIX (fungi/regime-insensitivity -- see rh_anomaly_idx's field comment for the verified root cause): the pure absolute-threshold term below (rh_s) is UNCHANGED -- it stays the sole driver whenever rh_anomaly_idx is at its 0.0 default, which is every synthetic ZoneObs and any real-data caller that hasn't run climatology anomaly application. This is a true no-op for existing behavior, not a reweighting: anomaly_adj is additive and clipped to +/-0.3, so it is exactly 0 when rh_anomaly_idx is exactly 0. When populated (real data via climatology.apply_climatology_anomalies), anomaly_adj pulls the absolute term down during anomalously DRY periods (e.g. El Nino, even though rh_max_pct often stays nominally high in absolute terms in a tropical climate) and up during anomalously WET periods (La Nina) -- restoring the regime sensitivity the pure absolute threshold structurally couldn't have. The absolute floor is kept deliberately: fungal risk still needs a real minimum humidity regardless of how anomalous conditions are. """ rh_s = _clip((self.rh_max_pct - 70.0) / 30.0, 0.0, 1.0) anomaly_adj = _clip(self.rh_anomaly_idx / 3.0, -0.3, 0.3) rh_s_adjusted = _clip(rh_s + anomaly_adj, 0.0, 1.0) mult = ( 1.0 if self.crop_stage in (CropStage.GRAIN_FILLING, CropStage.MATURATION) else 0.5 ) return float(rh_s_adjusted * mult) def composite_risk(self) -> float: """Single-number composite risk [0, 1] for belief map initialisation.""" return float(_clip( 0.35 * self.drought_signal() + 0.40 * self.flood_signal() + 0.25 * self.fungi_risk_signal(), 0.0, 1.0, )) # --------------------------------------------------------------------------- # ForecastResult # --------------------------------------------------------------------------- @dataclass(frozen=True) class ForecastResult: """ Probabilistic forecast for one zone over horizon_days. Frozen: produced by timesfm_wrapper.py, never mutated downstream. All sequences are horizon_days elements (default 30). """ zone_id: str forecast_time: datetime horizon_days: int = 30 # Point forecasts (median) precip_mm: Tuple[float, ...] = field(default_factory=tuple) temp_mean_c: Tuple[float, ...] = field(default_factory=tuple) rh_mean_pct: Tuple[float, ...] = field(default_factory=tuple) # Uncertainty (p10 to p90) precip_p10: Tuple[float, ...] = field(default_factory=tuple) precip_p90: Tuple[float, ...] = field(default_factory=tuple) temp_p10: Tuple[float, ...] = field(default_factory=tuple) temp_p90: Tuple[float, ...] = field(default_factory=tuple) # Exceedance probabilities, validated to [0, 1] prob_heavy_rain: Tuple[float, ...] = field(default_factory=tuple) prob_drought_day: Tuple[float, ...] = field(default_factory=tuple) prob_high_humidity: Tuple[float, ...] = field(default_factory=tuple) model_id: str = "timesfm-2.5-200m" crps_score: Optional[float] = None source: DataSource = DataSource.SYNTHETIC extras: Dict[str, Any] = field(default_factory=dict) _SEQUENCE_FIELDS: ClassVar[Tuple[str, ...]] = ( "precip_mm", "temp_mean_c", "rh_mean_pct", "precip_p10", "precip_p90", "temp_p10", "temp_p90", "prob_heavy_rain", "prob_drought_day", "prob_high_humidity", ) _PROB_FIELDS: ClassVar[Tuple[str, ...]] = ( "prob_heavy_rain", "prob_drought_day", "prob_high_humidity", ) def __post_init__(self) -> None: # --- Defensive copy: prevent caller mutation of extras dict --- object.__setattr__(self, "extras", dict(self.extras)) # All non-empty sequences must be the same length seqs = [ (name, getattr(self, name)) for name in self._SEQUENCE_FIELDS if getattr(self, name) ] if seqs: lengths = {len(s) for _, s in seqs} if len(lengths) > 1: raise ValueError( f"ForecastResult('{self.zone_id}'): sequence length mismatch: " f"{ {n: len(s) for n, s in seqs} }" ) # Enforce that all sequences match horizon_days exactly expected = self.horizon_days for name, seq in seqs: if len(seq) != expected: raise ValueError( f"ForecastResult('{self.zone_id}'): {name} length={len(seq)} " f"!= horizon_days={expected}. Truncate or pad before constructing." ) # Probability values must be in [0, 1] for fname in self._PROB_FIELDS: for i, v in enumerate(getattr(self, fname)): if not (0.0 <= v <= 1.0): raise ValueError( f"ForecastResult('{self.zone_id}'): " f"{fname}[{i}]={v:.4f} outside [0, 1]. " f"Clip before constructing ForecastResult." ) # p10 <= p90 for i, (lo, hi) in enumerate(zip(self.precip_p10, self.precip_p90)): if lo > hi: raise ValueError( f"ForecastResult('{self.zone_id}'): " f"precip_p10[{i}]={lo} > precip_p90[{i}]={hi}" ) def peak_precip_day(self) -> Optional[int]: if not self.precip_mm: return None return int(max(range(len(self.precip_mm)), key=lambda i: self.precip_mm[i])) def cumulative_precip_mm(self, window_days: int = 14) -> float: return float(sum(self.precip_mm[:window_days])) def max_consecutive_rain_days(self, threshold_mm: float = 10.0) -> int: max_run = run = 0 for p in self.precip_mm: run = run + 1 if p > threshold_mm else 0 max_run = max(max_run, run) return max_run def mean_exceedance_prob( self, field_name: str, window_days: Optional[int] = None ) -> float: """Mean probability of exceedance over window_days (default: full horizon).""" seq = getattr(self, field_name, ()) if not seq: return 0.0 window = seq[:window_days] if window_days else seq return float(sum(window) / len(window)) def to_dict(self) -> Dict[str, Any]: d = asdict(self) d["forecast_time"] = self.forecast_time.isoformat() d["source"] = self.source.value d["_schema_version"] = SCHEMA_VERSION return d @classmethod def from_dict(cls, d: Dict[str, Any]) -> "ForecastResult": sv, d = _copy_and_pop_schema(d) _check_schema(sv, "ForecastResult") d["forecast_time"] = datetime.fromisoformat(d["forecast_time"]) d["source"] = DataSource(d["source"]) for k in cls._SEQUENCE_FIELDS: if k in d and isinstance(d[k], list): d[k] = tuple(float(v) for v in d[k]) return cls(**{k: v for k, v in d.items() if not k.startswith("_")}) # --------------------------------------------------------------------------- # RiskScore # --------------------------------------------------------------------------- @dataclass(frozen=True) class RiskScore: """ Composite risk assessment for one sourcing zone. Frozen: terminal output of crop_risk_scorer.py. Never mutated. All scalar scores in [0.0, 1.0] unless noted. """ zone_id: str scored_at: datetime # Supply risk supply_shortfall_prob: float = 0.0 # P(zone delivers < 80% of contracted volume) drought_risk: float = 0.0 # [0, 1] flood_risk: float = 0.0 # [0, 1] supply_risk_composite: float = 0.0 # weighted dashboard score [0, 1] # Quality risk fungi_contamination_prob: float = 0.0 # P(moisture-related quality downgrade) harvest_delay_days: float = 0.0 # expected delay in days; >= 0 quality_risk_composite: float = 0.0 # [0, 1] # Timing optimal_harvest_window_start: Optional[datetime] = None optimal_harvest_window_end: Optional[datetime] = None # Decision output alert_level: AlertLevel = AlertLevel.NONE action_notes: str = "" # Confidence confidence: float = 0.5 # [0, 1] extras: Dict[str, Any] = field(default_factory=dict) _PROB_FIELDS: ClassVar[Tuple[str, ...]] = ( "supply_shortfall_prob", "drought_risk", "flood_risk", "supply_risk_composite", "fungi_contamination_prob", "quality_risk_composite", "confidence", ) def __post_init__(self) -> None: # --- Defensive copy: prevent caller mutation of extras dict --- object.__setattr__(self, "extras", dict(self.extras)) for attr in self._PROB_FIELDS: val = getattr(self, attr) clipped = _clip(val, 0.0, 1.0) if abs(clipped - val) > 1e-9: logger.warning( f"RiskScore('{self.zone_id}'): {attr}={val:.4f} " f"outside [0,1], clipped to {clipped:.4f}." ) object.__setattr__(self, attr, float(clipped)) if self.harvest_delay_days < 0.0: object.__setattr__(self, "harvest_delay_days", 0.0) # RiskScore is fully internal (produced by crop_risk_scorer.py only), # so we enforce timezone-awareness strictly here rather than silently fixing. if self.scored_at.tzinfo is None: raise ValueError( f"RiskScore('{self.zone_id}'): scored_at must be timezone-aware (UTC). " f"Use datetime.now(tz=timezone.utc) or .replace(tzinfo=timezone.utc)." ) for dt_attr in ("optimal_harvest_window_start", "optimal_harvest_window_end"): dt = getattr(self, dt_attr) if dt is not None and dt.tzinfo is None: raise ValueError( f"RiskScore('{self.zone_id}'): {dt_attr} must be timezone-aware (UTC)." ) def is_actionable(self) -> bool: """ Elevated attention: ADVISORY and above (alert_level > WATCH). Historical name — prefer is_elevated() / is_product_actionable() for new code. Product bus emission uses product_alert_service gate (WARNING+ or hazard thresholds), not this method alone. """ return self.alert_level > AlertLevel.WATCH def is_elevated(self) -> bool: """Internal / elevated attention: ADVISORY and above.""" return self.alert_level.severity() >= AlertLevel.ADVISORY.severity() def is_product_actionable(self) -> bool: """ Client-facing product floor on AlertLevel alone: WARNING and above. Hazard-threshold override (drought/flood freeze bars) is applied in product_alert_service.is_product_actionable(score, gate), not here. """ return self.alert_level.severity() >= AlertLevel.WARNING.severity() def harvest_window_days(self) -> Optional[int]: """Duration of optimal harvest window in days. None if not set.""" if self.optimal_harvest_window_start and self.optimal_harvest_window_end: return max( 0, (self.optimal_harvest_window_end - self.optimal_harvest_window_start).days, ) return None def to_dict(self) -> Dict[str, Any]: d = asdict(self) d["scored_at"] = self.scored_at.isoformat() d["alert_level"] = self.alert_level.value d["optimal_harvest_window_start"] = ( self.optimal_harvest_window_start.isoformat() if self.optimal_harvest_window_start else None ) d["optimal_harvest_window_end"] = ( self.optimal_harvest_window_end.isoformat() if self.optimal_harvest_window_end else None ) d["_schema_version"] = SCHEMA_VERSION return d @classmethod def from_dict(cls, d: Dict[str, Any]) -> "RiskScore": sv, d = _copy_and_pop_schema(d) _check_schema(sv, "RiskScore") d["scored_at"] = datetime.fromisoformat(d["scored_at"]) d["alert_level"] = AlertLevel(d["alert_level"]) d["optimal_harvest_window_start"] = ( datetime.fromisoformat(d["optimal_harvest_window_start"]) if d.get("optimal_harvest_window_start") else None ) d["optimal_harvest_window_end"] = ( datetime.fromisoformat(d["optimal_harvest_window_end"]) if d.get("optimal_harvest_window_end") else None ) return cls(**{k: v for k, v in d.items() if not k.startswith("_")}) # --------------------------------------------------------------------------- # BasinContext (schema v3+) # --------------------------------------------------------------------------- _HELIO_REGIMES = frozenset({"quiet", "active", "storm"}) def derive_helio_regime(kp_index: float, goes_xray_flux: float) -> str: """Map Kp + GOES X-ray flux to a coarse helio regime label. Thresholds chosen so the *default / typical* state is ``quiet`` and storms are rare — the same anti-saturation principle used for the humidity anomaly fix. These labels are context features only; they do not directly drive crop-risk scores until a later policy/scorer change explicitly consumes them. Rules (NOAA-style G-scale / GOES class approximation): storm — Kp >= 5 (G1+) OR X-ray >= 1e-5 (M-class+) active — Kp >= 3 (unsettled) OR X-ray >= 1e-6 (C-class+) quiet — otherwise """ if kp_index >= 5.0 or goes_xray_flux >= 1e-5: return "storm" if kp_index >= 3.0 or goes_xray_flux >= 1e-6: return "active" return "quiet" @dataclass(frozen=True) class BasinContext: """ Basin-scale / teleconnection context for one episode. Deliberately NOT a ZoneObs field. ENSO, IOD, ITCZ position, monsoon pressure systems, and heliophysics indices are not zone-specific -- they are the same value for every zone in an episode. Storing them per-zone would duplicate an identical scalar across every zone and risk the copies drifting apart across a training run. One BasinContext per EpisodeContext, consumed as a distinct input channel (not folded into per-zone belief/forecast arrays) by weather_forecast_env.py and gru_weather_policy.py. Field bounds are generous (not tight physical limits) since these are RL policy inputs -- clipping only guards against clearly corrupt values, not against genuine extremes (e.g. an ONI of 2.8 during a strong El Nino is real, not an error). Helio fields (solar_wind_speed_kms, kp_index, goes_xray_flux, helio_regime) default to quiet-Sun values so missing / synthetic / offline paths never inject artificial "storm" context the way the old absolute humidity term saturated fungi risk. They ride along as optional context; the crop-risk scorer does not consume them until an explicit later integration step. """ valid_date: datetime enso_oni: float = 0.0 # Oceanic (or Relative Oceanic) Nino Index, degrees C anomaly iod_dmi: float = 0.0 # Indian Ocean Dipole Mode Index, degrees C itcz_latitude_deg: float = 0.0 # approximate ITCZ position, degrees N (negative = south) mslp_regional_hpa: float = 1013.25 # area-averaged regional MSLP, hPa (monsoon high/low proxy) # Helio / space-weather (quiet-Sun defaults — anti-saturation design) solar_wind_speed_kms: float = 400.0 # typical quiet-Sun ~300–450 km/s kp_index: float = 2.0 # planetary K-index [0, 9]; ~2 is quiet goes_xray_flux: float = 1e-7 # W/m²; background / low-C floor helio_regime: str = "quiet" # "quiet" | "active" | "storm" source: DataSource = DataSource.SYNTHETIC extras: Dict[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: object.__setattr__(self, "extras", dict(self.extras)) if self.valid_date.tzinfo is None: object.__setattr__( self, "valid_date", self.valid_date.replace(tzinfo=timezone.utc) ) object.__setattr__(self, "enso_oni", float(_clip(self.enso_oni, -5.0, 5.0))) object.__setattr__(self, "iod_dmi", float(_clip(self.iod_dmi, -5.0, 5.0))) object.__setattr__(self, "itcz_latitude_deg", float(_clip(self.itcz_latitude_deg, -30.0, 30.0))) object.__setattr__(self, "mslp_regional_hpa", float(_clip(self.mslp_regional_hpa, 900.0, 1100.0))) # Helio clipping — physical ranges, not risk-amplifying floors. object.__setattr__( self, "solar_wind_speed_kms", float(_clip(self.solar_wind_speed_kms, 200.0, 1200.0)), ) object.__setattr__(self, "kp_index", float(_clip(self.kp_index, 0.0, 9.0))) object.__setattr__( self, "goes_xray_flux", float(_clip(self.goes_xray_flux, 1e-9, 1e-3)), ) regime = self.helio_regime if self.helio_regime in _HELIO_REGIMES else "quiet" object.__setattr__(self, "helio_regime", regime) def to_dict(self) -> Dict[str, Any]: return { "valid_date": self.valid_date.isoformat(), "enso_oni": self.enso_oni, "iod_dmi": self.iod_dmi, "itcz_latitude_deg": self.itcz_latitude_deg, "mslp_regional_hpa": self.mslp_regional_hpa, "solar_wind_speed_kms": self.solar_wind_speed_kms, "kp_index": self.kp_index, "goes_xray_flux": self.goes_xray_flux, "helio_regime": self.helio_regime, "source": self.source.value, "extras": self.extras, "_schema_version": SCHEMA_VERSION, } @classmethod def from_dict(cls, d: Dict[str, Any]) -> "BasinContext": sv, d = _copy_and_pop_schema(d) _check_schema(sv, "BasinContext") d["valid_date"] = datetime.fromisoformat(d["valid_date"]) d["source"] = DataSource(d.get("source", "synthetic")) # Helio fields optional for backward compatibility with records # serialised before they existed — dataclass defaults apply when # keys are absent. return cls(**{k: v for k, v in d.items() if not k.startswith("_")}) def make_synthetic_basin_context( valid_date: Optional[datetime] = None, seed: Optional[int] = None, ) -> BasinContext: """Deterministic synthetic BasinContext for tests and default episodes. Helio draws are deliberately *quiet-biased* (Kp mostly 0.5–3.5, X-ray background-to-low-C) so synthetic training trajectories do not inject a constant "active/storm" prior the way the old absolute humidity term saturated fungi risk. Occasional higher draws still occur so the policy can see regime transitions. """ rng = random.Random(seed if seed is not None else 0) if valid_date is None: valid_date = datetime(2020, 1, 1, tzinfo=timezone.utc) kp = float(rng.uniform(0.5, 3.5)) if rng.random() < 0.10: kp = float(rng.uniform(4.0, 7.0)) sw = float(rng.uniform(320.0, 480.0)) if kp >= 5.0: sw = float(rng.uniform(500.0, 800.0)) log_xray = rng.uniform(-8.0, -6.5) if kp >= 5.0: log_xray = rng.uniform(-5.5, -4.5) xray = float(10.0 ** log_xray) regime = derive_helio_regime(kp, xray) return BasinContext( valid_date=valid_date, enso_oni=rng.uniform(-1.5, 1.5), iod_dmi=rng.uniform(-1.0, 1.0), itcz_latitude_deg=rng.uniform(-10.0, 10.0), mslp_regional_hpa=rng.uniform(1005.0, 1020.0), solar_wind_speed_kms=sw, kp_index=kp, goes_xray_flux=xray, helio_regime=regime, source=DataSource.SYNTHETIC, ) # --------------------------------------------------------------------------- # EpisodeContext # --------------------------------------------------------------------------- @dataclass class EpisodeContext: """ Typed episode initialiser for weather_forecast_env.py. Analogous to the wafer_data dict in the MEMS env, but fully typed and validated. The pipeline produces this; the env consumes it at _soft_reset(). Economic parameters live in config (ForecastConfig) rather than as loose floats, so the rational termination threshold is always validated. Multi-zone real eval (Blocker B) -------------------------------- Prefer ``zone_obs`` / ``zone_forecasts`` parallel lists (same length as ``zone_ids``). Legacy single-zone callers may leave those lists empty and only set ``obs`` / ``forecast``; ``resolved_zone_obs()`` then returns ``[obs]``. Injecting multi-zone real data without the lists is refused by WeatherForecastEnv (padding one zone into N slots is not multi-zone). """ obs: ZoneObs forecast: ForecastResult config: ForecastConfig = field(default_factory=ForecastConfig) ground_truth: Optional[RiskScore] = None # None during live deployment; set during backtesting / eval # Spatial context -- mirrors cluster_model in the MEMS env zone_ids: List[str] = field(default_factory=list) adjacency: Dict[str, List[str]] = field(default_factory=dict) # adjacency[zone_id] = [neighbouring zone_ids] data_source: DataSource = DataSource.SYNTHETIC # Basin-scale / teleconnection context (schema v3+). One per episode, # shared across all zones -- see BasinContext docstring. None is a valid # and expected value (e.g. legacy v2-shaped callers, or callers that # haven't wired up basin data yet); consumers must supply their own # neutral/climatological default when it is None, not treat it as an error. basin_context: Optional[BasinContext] = None # Per-zone real/synthetic payloads for multi-zone injection (optional). # When empty, resolved_* fall back to [obs] / [forecast] (single-zone). zone_obs: List[ZoneObs] = field(default_factory=list) zone_forecasts: List[ForecastResult] = field(default_factory=list) def __post_init__(self) -> None: if not self.obs.zone_id: raise ValueError("EpisodeContext: obs.zone_id is empty") if self.obs.zone_id != self.forecast.zone_id: raise ValueError( f"EpisodeContext: obs.zone_id='{self.obs.zone_id}' != " f"forecast.zone_id='{self.forecast.zone_id}'" ) if ( self.ground_truth is not None and self.ground_truth.zone_id != self.obs.zone_id ): raise ValueError( f"EpisodeContext: ground_truth.zone_id='{self.ground_truth.zone_id}'" f" != obs.zone_id='{self.obs.zone_id}'" ) if self.obs.zone_id not in self.zone_ids: self.zone_ids = [self.obs.zone_id] + list(self.zone_ids) # --- Multi-zone list integrity (optional fields) --- if self.zone_obs or self.zone_forecasts: if len(self.zone_obs) != len(self.zone_forecasts): raise ValueError( f"EpisodeContext: len(zone_obs)={len(self.zone_obs)} != " f"len(zone_forecasts)={len(self.zone_forecasts)}" ) if len(self.zone_obs) != len(self.zone_ids): raise ValueError( f"EpisodeContext: len(zone_obs)={len(self.zone_obs)} != " f"len(zone_ids)={len(self.zone_ids)}" ) for i, (zo, zf, zid) in enumerate( zip(self.zone_obs, self.zone_forecasts, self.zone_ids) ): if zo.zone_id != zid: raise ValueError( f"EpisodeContext: zone_obs[{i}].zone_id={zo.zone_id!r} " f"!= zone_ids[{i}]={zid!r}" ) if zf.zone_id != zid: raise ValueError( f"EpisodeContext: zone_forecasts[{i}].zone_id={zf.zone_id!r} " f"!= zone_ids[{i}]={zid!r}" ) if zo.zone_id != zf.zone_id: raise ValueError( f"EpisodeContext: zone_obs[{i}] / zone_forecasts[{i}] " f"zone_id mismatch" ) # Keep primary obs/forecast aligned with slot 0 for legacy readers. if self.zone_obs[0].zone_id != self.obs.zone_id: # Prefer list as authority when multi-zone lists are present. self.obs = self.zone_obs[0] self.forecast = self.zone_forecasts[0] # --- Adjacency integrity: all keys and neighbours must be in zone_ids --- for z, neighbours in self.adjacency.items(): if z not in self.zone_ids: raise ValueError( f"EpisodeContext: adjacency key '{z}' not in zone_ids={self.zone_ids}" ) for n in neighbours: if n not in self.zone_ids: raise ValueError( f"EpisodeContext: adjacency neighbour '{n}' (of '{z}') " f"not in zone_ids={self.zone_ids}" ) @property def n_zones(self) -> int: return len(self.zone_ids) def resolved_zone_obs(self) -> List[ZoneObs]: """Per-zone obs for env injection; falls back to [obs] if lists empty.""" if self.zone_obs: return list(self.zone_obs) return [self.obs] def resolved_zone_forecasts(self) -> List[ForecastResult]: """Per-zone forecasts for env injection; falls back to [forecast].""" if self.zone_forecasts: return list(self.zone_forecasts) return [self.forecast] def to_dict(self) -> Dict[str, Any]: return { "obs": self.obs.to_dict(), "forecast": self.forecast.to_dict(), "config": self.config.to_dict(), "ground_truth": self.ground_truth.to_dict() if self.ground_truth else None, "zone_ids": self.zone_ids, "adjacency": self.adjacency, "data_source": self.data_source.value, "basin_context": self.basin_context.to_dict() if self.basin_context else None, "zone_obs": [z.to_dict() for z in self.zone_obs], "zone_forecasts": [f.to_dict() for f in self.zone_forecasts], "_schema_version": SCHEMA_VERSION, } @classmethod def from_dict(cls, d: Dict[str, Any]) -> "EpisodeContext": sv, d = _copy_and_pop_schema(d) _check_schema(sv, "EpisodeContext") zone_obs_raw = d.get("zone_obs") or [] zone_fc_raw = d.get("zone_forecasts") or [] return cls( obs=ZoneObs.from_dict(d["obs"]), forecast=ForecastResult.from_dict(d["forecast"]), config=ForecastConfig.from_dict(d["config"]), ground_truth=( RiskScore.from_dict(d["ground_truth"]) if d.get("ground_truth") else None ), zone_ids=d.get("zone_ids", []), adjacency=d.get("adjacency", {}), data_source=DataSource(d.get("data_source", "synthetic")), basin_context=( BasinContext.from_dict(d["basin_context"]) if d.get("basin_context") else None ), zone_obs=[ZoneObs.from_dict(x) for x in zone_obs_raw], zone_forecasts=[ForecastResult.from_dict(x) for x in zone_fc_raw], ) # --------------------------------------------------------------------------- # Synthetic generators # --------------------------------------------------------------------------- def make_synthetic_zone_obs( zone_id: str = "synthetic_zone_0", crop_stage: CropStage = CropStage.GRAIN_FILLING, drought: bool = False, flood: bool = False, fungi: bool = False, seed: Optional[int] = None, ) -> ZoneObs: """Deterministic synthetic ZoneObs for unit tests and env smoke tests. Seeded via zlib.crc32 (not hash()) for process-stable reproducibility. In production, era5_data_pipeline.py replaces this entirely. """ rng = random.Random(seed if seed is not None else _stable_seed(zone_id)) # Deterministic timestamp: seeded offset from a fixed base, not datetime.now(). # datetime.now() breaks training reproducibility -- two runs with the same seed # produce ZoneObs objects that compare unequal on valid_time. _BASE_TIME = datetime(2020, 1, 1, tzinfo=timezone.utc) _synthetic_valid_time = _BASE_TIME + timedelta(days=rng.randint(0, 3650)) base_precip = ( rng.uniform(60.0, 120.0) if flood else rng.uniform(0.0, 2.0) if drought else 5.0 ) rh = rng.uniform(85.0, 98.0) if fungi else rng.uniform(55.0, 75.0) return ZoneObs( zone_id=zone_id, valid_time=_synthetic_valid_time, source=DataSource.SYNTHETIC, precip_24h_mm=base_precip, precip_7d_mm=base_precip * 6.5, precip_14d_mm=base_precip * 12.0, precip_30d_mm=base_precip * 24.0, precip_anomaly_idx=3.0 if flood else (-2.5 if drought else rng.uniform(-0.5, 0.5)), temp_mean_c=rng.uniform(26.0, 32.0), temp_max_c=rng.uniform(31.0, 36.0), temp_min_c=rng.uniform(22.0, 26.0), temp_anomaly_idx=rng.uniform(-0.5, 0.5), gdd_accumulated=rng.uniform(400.0, 900.0), heat_stress_days=rng.randint(0, 5), cold_stress_days=0, soil_moisture_pct=rng.uniform(10.0, 25.0) if drought else rng.uniform(40.0, 70.0), soil_moisture_anom=-2.0 if drought else rng.uniform(-0.5, 0.5), evapotranspiration_mm=rng.uniform(4.0, 7.0), wind_speed_max_ms=rng.uniform(3.0, 8.0), wind_speed_mean_ms=rng.uniform(1.0, 3.5), rh_mean_pct=rh * 0.9, rh_max_pct=rh, ndvi=rng.uniform(0.35, 0.80), ndvi_anomaly_idx=rng.uniform(-0.3, 0.3), flood_extent_pct=rng.uniform(20.0, 60.0) if flood else 0.0, drainage_risk_idx=rng.uniform(0.5, 0.9) if flood else rng.uniform(0.0, 0.3), crop_stage=crop_stage, days_to_harvest=rng.randint(7, 45), quality_flag=3, cloud_cover_pct=rng.uniform(0.0, 20.0), ) def make_synthetic_forecast_result( zone_id: str = "synthetic_zone_0", valid_time: Optional[datetime] = None, horizon_days: int = 30, drought: bool = False, flood: bool = False, seed: Optional[int] = None, ) -> ForecastResult: """Deterministic synthetic ForecastResult for testing.""" rng = random.Random( seed if seed is not None else _stable_seed(zone_id + "_forecast") ) if valid_time is None: _BASE_TIME = datetime(2020, 1, 1, tzinfo=timezone.utc) t = _BASE_TIME + timedelta(days=rng.randint(0, 3650)) else: t = valid_time precip = tuple( max(0.0, rng.uniform(30.0, 80.0) if flood else rng.uniform(0.0, 3.0) if drought else max(0.0, rng.gauss(8.0, 5.0))) for _ in range(horizon_days) ) temp = tuple(rng.uniform(26.0, 32.0) for _ in range(horizon_days)) rh = tuple(rng.uniform(60.0, 90.0) for _ in range(horizon_days)) p10 = tuple(max(0.0, p * rng.uniform(0.3, 0.7)) for p in precip) p90 = tuple(p * rng.uniform(1.3, 2.0) for p in precip) prob_rain = tuple( float(_clip(p / 60.0 + rng.uniform(-0.05, 0.05), 0.0, 1.0)) for p in precip ) prob_drought = tuple( float(_clip(0.8 if drought else rng.uniform(0.0, 0.15), 0.0, 1.0)) for _ in range(horizon_days) ) prob_humid = tuple( float(_clip((r - 70.0) / 30.0 + rng.uniform(-0.05, 0.05), 0.0, 1.0)) for r in rh ) return ForecastResult( zone_id=zone_id, forecast_time=t, horizon_days=horizon_days, precip_mm=precip, temp_mean_c=temp, rh_mean_pct=rh, precip_p10=p10, precip_p90=p90, temp_p10=tuple(v - rng.uniform(1.0, 3.0) for v in temp), temp_p90=tuple(v + rng.uniform(1.0, 3.0) for v in temp), prob_heavy_rain=prob_rain, prob_drought_day=prob_drought, prob_high_humidity=prob_humid, source=DataSource.SYNTHETIC, ) def make_synthetic_episode_context( zone_id: str = "synthetic_zone_0", config: Optional[ForecastConfig] = None, drought: bool = False, flood: bool = False, fungi: bool = False, seed: Optional[int] = None, ) -> EpisodeContext: """Construct a complete synthetic EpisodeContext for env smoke tests.""" cfg = config or ForecastConfig() obs = make_synthetic_zone_obs(zone_id, drought=drought, flood=flood, fungi=fungi, seed=seed) fcast = make_synthetic_forecast_result(zone_id, valid_time=obs.valid_time, drought=drought, flood=flood, seed=seed) basin = ( make_synthetic_basin_context(valid_date=obs.valid_time, seed=seed) if cfg.include_basin_context else None ) return EpisodeContext( obs=obs, forecast=fcast, config=cfg, zone_ids=[zone_id], data_source=DataSource.SYNTHETIC, basin_context=basin, ) # --------------------------------------------------------------------------- # Self-test (python zone_observation.py) # --------------------------------------------------------------------------- if __name__ == "__main__": import sys logging.basicConfig(level=logging.WARNING) print(f"zone_observation.py schema_version={SCHEMA_VERSION}\n") failures: List[str] = [] def _assert(condition: bool, msg: str) -> None: if not condition: failures.append(msg) print(f" FAIL: {msg}") # 1. ZoneObs round-trip + non-mutation obs = make_synthetic_zone_obs("test_flood", flood=True, seed=42) d = obs.to_dict() had_sv = "_schema_version" in d obs2 = ZoneObs.from_dict(d) still_has_sv = "_schema_version" in d _assert(had_sv and still_has_sv, "ZoneObs.from_dict mutated caller dict") _assert(obs.zone_id == obs2.zone_id, "ZoneObs zone_id round-trip") _assert(abs(obs.precip_24h_mm - obs2.precip_24h_mm) < 1e-9, "ZoneObs precip precision") _assert(obs.crop_stage == obs2.crop_stage, "ZoneObs crop_stage round-trip") _assert(obs.source == obs2.source, "ZoneObs source round-trip") print(f" ZoneObs flood={obs.flood_signal():.3f} drought={obs.drought_signal():.3f}" f" fungi={obs.fungi_risk_signal():.3f} composite={obs.composite_risk():.3f}") # 2. Deterministic seeding a = make_synthetic_zone_obs("stable", seed=99) b = make_synthetic_zone_obs("stable", seed=99) _assert(a.precip_24h_mm == b.precip_24h_mm, "Explicit seed not deterministic") c = make_synthetic_zone_obs("crc_zone") d2 = make_synthetic_zone_obs("crc_zone") _assert(c.precip_24h_mm == d2.precip_24h_mm, "zlib.crc32 seed not stable") print(" Deterministic seeding OK") # 3. ZoneObs.validate() obs_v = make_synthetic_zone_obs("val_zone", seed=1) obs_v.precip_7d_mm = obs_v.precip_14d_mm + 50.0 issues = ZoneObs.validate(obs_v) _assert(len(issues) > 0, "validate() missed precip_14d < precip_7d") print(f" ZoneObs.validate() caught {len(issues)} issue(s)") # 4. GeoPolygon string-vertex coercion + contains_point poly = GeoPolygon( vertices=[("3.0", "101.0"), (3.1, 101.0), (3.1, 101.1), (3.0, 101.1)], zone_id="sel_A1", ) _assert(isinstance(poly.centroid[0], float), "GeoPolygon centroid not float") _assert(poly.contains_point(3.05, 101.05), "GeoPolygon inside point") _assert(not poly.contains_point(4.0, 102.0), "GeoPolygon outside point") poly2 = GeoPolygon.from_dict(poly.to_dict()) _assert(poly.zone_id == poly2.zone_id, "GeoPolygon round-trip") print(f" GeoPolygon area={poly.approx_area_km2:.1f} km2 contains_point OK") # 5. ForecastResult round-trip + prob validation + non-mutation fr = make_synthetic_forecast_result("test_flood", flood=True, seed=42) d_fr = fr.to_dict() had_sv_fr = "_schema_version" in d_fr fr2 = ForecastResult.from_dict(d_fr) _assert(had_sv_fr and "_schema_version" in d_fr, "ForecastResult.from_dict mutated dict") _assert(fr.precip_mm == fr2.precip_mm, "ForecastResult precip round-trip") _assert(fr.source == fr2.source, "ForecastResult source round-trip") try: ForecastResult( zone_id="x", forecast_time=datetime.now(tz=timezone.utc), precip_mm=tuple([0.0]*30), temp_mean_c=tuple([29.0]*30), rh_mean_pct=tuple([70.0]*30), precip_p10=tuple([0.0]*30), precip_p90=tuple([1.0]*30), prob_heavy_rain=tuple([5.0]*30), prob_drought_day=tuple([0.0]*30), prob_high_humidity=tuple([0.0]*30), ) _assert(False, "ForecastResult accepted prob > 1.0") except ValueError: pass print(f" ForecastResult peak_day={fr.peak_precip_day()}" f" cumul14d={fr.cumulative_precip_mm(14):.1f}mm prob_validation OK") # 6. RiskScore round-trip + ordering + harvest_window_days now = datetime.now(tz=timezone.utc) rs = RiskScore( zone_id="test_flood", scored_at=now, supply_shortfall_prob=0.35, drought_risk=0.05, flood_risk=0.78, supply_risk_composite=0.55, fungi_contamination_prob=0.42, harvest_delay_days=6.0, quality_risk_composite=0.42, optimal_harvest_window_start=now + timedelta(days=14), optimal_harvest_window_end=now + timedelta(days=21), alert_level=AlertLevel.WARNING, confidence=0.80, ) d_rs = rs.to_dict() rs2 = RiskScore.from_dict(d_rs) _assert("_schema_version" in d_rs, "RiskScore.from_dict mutated dict") _assert(rs.alert_level == rs2.alert_level, "RiskScore alert_level round-trip") _assert(rs.harvest_window_days() == 7, "RiskScore harvest_window_days") _assert(rs.is_actionable(), "RiskScore.is_actionable() for WARNING") _assert(AlertLevel.WARNING > AlertLevel.WATCH, "AlertLevel ordering >") _assert(AlertLevel.NONE < AlertLevel.CRITICAL, "AlertLevel ordering <") print(f" RiskScore alert={rs.alert_level.value} window={rs.harvest_window_days()}d" f" actionable={rs.is_actionable()}") # 7. ForecastConfig rational threshold + new pipeline fields round-trip cfg = ForecastConfig() _assert(cfg.belief_floor < cfg.rational_termination_threshold, "Default ForecastConfig: belief_floor >= rational_threshold") cfg2 = ForecastConfig.from_dict(cfg.to_dict()) _assert(cfg.alert_value == cfg2.alert_value, "ForecastConfig round-trip") _assert(cfg2.real_data_ratio == 0.7, "ForecastConfig real_data_ratio round-trip") _assert(cfg2.era5_ratio == 0.5, "ForecastConfig era5_ratio round-trip") _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") # Verify force_data_source serialises/deserialises correctly when set cfg_era5 = ForecastConfig(force_data_source=DataSource.ERA5_REANALYSIS) cfg_era5_back = ForecastConfig.from_dict(cfg_era5.to_dict()) _assert( cfg_era5_back.force_data_source == DataSource.ERA5_REANALYSIS, "ForecastConfig force_data_source=ERA5 round-trip" ) # New pipeline-integration fields round-trip + validation cfg_new = ForecastConfig( forecast_backend="openmeteo", use_climatology_anomalies=True, climatology_years=15, ) cfg_new_back = ForecastConfig.from_dict(cfg_new.to_dict()) _assert(cfg_new_back.forecast_backend == "openmeteo", "forecast_backend round-trip") _assert(cfg_new_back.use_climatology_anomalies is True, "use_climatology_anomalies round-trip") _assert(cfg_new_back.climatology_years == 15, "climatology_years round-trip") _assert(cfg2.forecast_backend == "synthetic", "forecast_backend default should be 'synthetic' (back-compat)") try: ForecastConfig(forecast_backend="not_a_backend") _assert(False, "ForecastConfig accepted invalid forecast_backend") except ValueError: pass print(f" ForecastConfig rational_threshold={cfg.rational_termination_threshold:.4f}" f" belief_floor={cfg.belief_floor:.4f} pipeline fields OK") # 8. EpisodeContext round-trip + validation ec = make_synthetic_episode_context("zone_A", seed=7) d_ec = ec.to_dict() ec2 = EpisodeContext.from_dict(d_ec) _assert(ec.obs.zone_id == ec2.obs.zone_id, "EpisodeContext zone_id round-trip") _assert(ec.config.alert_value == ec2.config.alert_value, "EpisodeContext config round-trip") _assert(ec.n_zones == 1, "EpisodeContext n_zones") try: EpisodeContext( obs=make_synthetic_zone_obs("zone_A"), forecast=make_synthetic_forecast_result("zone_B"), config=ForecastConfig(), ) _assert(False, "EpisodeContext accepted zone_id mismatch") except ValueError: pass print(f" EpisodeContext n_zones={ec.n_zones} zone_mismatch_check OK") # 9. Full JSON round-trip ec_json = json.dumps(ec.to_dict()) ec_back = EpisodeContext.from_dict(json.loads(ec_json)) _assert(ec.obs.zone_id == ec_back.obs.zone_id, "EpisodeContext JSON zone_id round-trip") _assert(ec.forecast.precip_mm == ec_back.forecast.precip_mm, "ForecastResult precip JSON round-trip") print(" Full JSON serialisation round-trip OK") # 10. BasinContext round-trip + clipping + helio + EpisodeContext integration bc = make_synthetic_basin_context(seed=3) d_bc = bc.to_dict() bc2 = BasinContext.from_dict(d_bc) _assert("_schema_version" in d_bc, "BasinContext.from_dict mutated dict") _assert(abs(bc.enso_oni - bc2.enso_oni) < 1e-9, "BasinContext enso_oni round-trip") _assert(abs(bc.iod_dmi - bc2.iod_dmi) < 1e-9, "BasinContext iod_dmi round-trip") _assert(bc.source == bc2.source, "BasinContext source round-trip") _assert(abs(bc.kp_index - bc2.kp_index) < 1e-9, "BasinContext kp_index round-trip") _assert(bc.helio_regime == bc2.helio_regime, "BasinContext helio_regime round-trip") _assert(bc.helio_regime in ("quiet", "active", "storm"), f"invalid helio_regime {bc.helio_regime!r}") bc_extreme = BasinContext(valid_date=now, enso_oni=99.0, mslp_regional_hpa=1.0) _assert(bc_extreme.enso_oni <= 5.0, "BasinContext enso_oni not clipped") _assert(bc_extreme.mslp_regional_hpa >= 900.0, "BasinContext mslp_regional_hpa not clipped") _assert(bc_extreme.kp_index == 2.0, "BasinContext kp default should be quiet-Sun 2.0") _assert(bc_extreme.helio_regime == "quiet", "BasinContext helio default should be quiet") _assert(derive_helio_regime(6.0, 1e-7) == "storm", "derive_helio_regime storm by Kp") _assert(derive_helio_regime(1.0, 2e-5) == "storm", "derive_helio_regime storm by X-ray") _assert(derive_helio_regime(3.5, 1e-7) == "active", "derive_helio_regime active") _assert(derive_helio_regime(1.0, 1e-8) == "quiet", "derive_helio_regime quiet") cfg_basin = ForecastConfig(include_basin_context=True) _assert(cfg_basin.require_real_basin_context is False, "require_real_basin_context should default False") ec_basin = make_synthetic_episode_context("zone_basin", config=cfg_basin, seed=11) _assert(ec_basin.basin_context is not None, "make_synthetic_episode_context did not attach basin_context when opted in") d_ec_basin = ec_basin.to_dict() ec_basin2 = EpisodeContext.from_dict(d_ec_basin) _assert(ec_basin2.basin_context is not None, "EpisodeContext.basin_context lost in round-trip") _assert( abs(ec_basin.basin_context.enso_oni - ec_basin2.basin_context.enso_oni) < 1e-9, "EpisodeContext.basin_context.enso_oni round-trip" ) _assert( ec_basin.basin_context.helio_regime == ec_basin2.basin_context.helio_regime, "EpisodeContext.basin_context.helio_regime round-trip" ) ec_no_basin = make_synthetic_episode_context("zone_no_basin", seed=11) _assert(ec_no_basin.basin_context is None, "basin_context should default to None when include_basin_context=False") print(f" BasinContext oni={bc.enso_oni:.2f} dmi={bc.iod_dmi:.2f} " f"kp={bc.kp_index:.1f} regime={bc.helio_regime} " f"round-trip OK, EpisodeContext integration OK") # 11. New optional ZoneObs satellite fields: None-by-default, clipping, round-trip obs_sat = ZoneObs( zone_id="sat_zone", valid_time=now, soil_moisture_satellite_pct=150.0, # out of range -> should clip to 100 precip_satellite_mm=12.5, ) _assert(obs_sat.soil_moisture_satellite_pct == 100.0, "soil_moisture_satellite_pct not clipped to 100") _assert(obs_sat.precip_satellite_mm == 12.5, "precip_satellite_mm unexpectedly altered") obs_plain = make_synthetic_zone_obs("plain_zone", seed=5) _assert(obs_plain.precip_satellite_mm is None, "precip_satellite_mm should default to None, not 0.0") _assert(obs_plain.soil_moisture_satellite_pct is None, "soil_moisture_satellite_pct should default to None, not 0.0") d_sat = obs_sat.to_dict() obs_sat2 = ZoneObs.from_dict(d_sat) _assert(obs_sat2.precip_satellite_mm == obs_sat.precip_satellite_mm, "precip_satellite_mm round-trip") print(" ZoneObs satellite fields: None-default, clipping, round-trip OK") print() if failures: print(f"FAILED {len(failures)} test(s):") for f in failures: print(f" - {f}") sys.exit(1) else: print(f"All {11} test groups passed.")