""" zone_observation.py =================== """ from __future__ import annotations import json import logging import math import random import zlib 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 logger = logging.getLogger(__name__) 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): 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): ERA5_REANALYSIS = "era5_reanalysis" # ECMWF ERA5 via CDS or Open-Meteo OPENMETEO_LIVE = "openmeteo_live" # Open-Meteo forecast API (free tier) OPENMETEO_ARCHIVE = "openmeteo_archive" # Open-Meteo historical archive API 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 TIGGE_ECMF = "tigge_ecmf" # TIGGE ECMWF control (CC BY 4.0) SYNTHETIC = "synthetic" # Generated by make_synthetic_* for training UNKNOWN = "unknown" def is_observational(self) -> bool: return self not in (DataSource.SYNTHETIC, DataSource.UNKNOWN) # --------------------------------------------------------------------------- # Utilities # --------------------------------------------------------------------------- def _clip(value: float, lo: float, hi: float) -> float: return max(lo, min(hi, value)) def _stable_seed(key: str) -> int: return zlib.crc32(key.encode("utf-8")) & 0x7FFFFFFF def _copy_and_pop_schema(d: Dict[str, Any]) -> Tuple[Optional[int], Dict[str, Any]]: 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: 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: # --- 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 belief_prior_weight: float = 0.70 # weight on episode prior vs. zone signal when # seeding zone_belief at reset: # prior_w*prior + (1-prior_w)*signal # --- 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) 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 uncertainty_decay: float = 0.70 # forecast_uncertainty[zone] *= this on inspection info_gain_scale: float = 5.0 # multiplier on belief info-gain in step reward uncertainty_penalty_scale: float = 5.0 # multiplier on mean uncertainty at termination economic_randomization: bool = False clean_episode_ratio: float = 0.7 event_spatial_correlation: float = 0.85 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: bool = True seed: Optional[int] = None 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, 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 include_basin_context: bool = False # attach ENSO/IOD/monsoon/helio context to EpisodeContext require_real_basin_context: bool = False forecast_backend: str = "synthetic" 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 not (0.0 <= self.belief_prior_weight <= 1.0): raise ValueError( f"ForecastConfig: belief_prior_weight={self.belief_prior_weight} " f"must be in [0, 1]" ) if not (0.0 <= self.uncertainty_decay <= 1.0): raise ValueError( f"ForecastConfig: uncertainty_decay={self.uncertainty_decay} " f"must be in [0, 1]" ) if self.info_gain_scale < 0: raise ValueError(f"ForecastConfig: info_gain_scale must be >= 0") if self.uncertainty_penalty_scale < 0: raise ValueError(f"ForecastConfig: uncertainty_penalty_scale 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.belief_prior_weight = float(_clip(self.belief_prior_weight, 0.0, 1.0)) self.uncertainty_decay = float(_clip(self.uncertainty_decay, 0.0, 1.0)) self.info_gain_scale = float(_clip(self.info_gain_scale, 0.0, 1_000.0)) self.uncertainty_penalty_scale = float(_clip(self.uncertainty_penalty_scale, 0.0, 1_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) ) 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: 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, "belief_prior_weight": self.belief_prior_weight, "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, "uncertainty_decay": self.uncertainty_decay, "info_gain_scale": self.info_gain_scale, "uncertainty_penalty_scale": self.uncertainty_penalty_scale, "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, "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, "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: 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]: 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: 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: 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: zone_id: str valid_time: datetime # UTC timestamp of observation window start source: DataSource = DataSource.UNKNOWN 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 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 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_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 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_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) 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 rh_anomaly_idx: float = 0.0 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_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_stage: CropStage = CropStage.UNKNOWN days_to_harvest: Optional[int] = None # None = unknown; 0 = harvest now planting_date: Optional[datetime] = None 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 extras: Dict[str, Any] = field(default_factory=dict) def __post_init__(self) -> None: 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)) 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) 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]: 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]: 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": 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) 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: 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: 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: 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: 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: zone_id: str forecast_time: datetime horizon_days: int = 30 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) 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) 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: object.__setattr__(self, "extras", dict(self.extras)) 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} }" ) 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." ) 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." ) 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: 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) if sv is not None: _check_schema(sv, "ForecastResult") raw_time = d.get("forecast_time") if isinstance(raw_time, str): d["forecast_time"] = datetime.fromisoformat(raw_time) src = d.get("source", "unknown") try: d["source"] = DataSource(src) except Exception: d["source"] = DataSource.UNKNOWN for k in cls._SEQUENCE_FIELDS: if k in d and isinstance(d[k], (list, tuple)): d[k] = tuple(float(v) if v is not None else 0.0 for v in d[k]) known = {f.name for f in _dc_fields(cls)} return cls(**{k: v for k, v in d.items() if k in known}) # --------------------------------------------------------------------------- # RiskScore # --------------------------------------------------------------------------- @dataclass(frozen=True) class RiskScore: zone_id: str scored_at: datetime 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] 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] optimal_harvest_window_start: Optional[datetime] = None optimal_harvest_window_end: Optional[datetime] = None alert_level: AlertLevel = AlertLevel.NONE action_notes: str = "" 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: 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) 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: return self.alert_level > AlertLevel.WATCH def is_elevated(self) -> bool: return self.alert_level.severity() >= AlertLevel.ADVISORY.severity() def is_product_actionable(self) -> bool: return self.alert_level.severity() >= AlertLevel.WARNING.severity() def harvest_window_days(self) -> Optional[int]: 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: """Classify heliophysical regime from Kp and GOES X-ray flux.""" kp = float(kp_index) xray = float(goes_xray_flux) if goes_xray_flux is not None else 1e-7 if kp >= 5.0 or xray >= 1e-5: return "storm" if kp >= 3.0 or xray >= 5e-7: return "active" return "quiet" @dataclass class BasinContext: """Basin-scale teleconnections + heliophysical context (schema v3+).""" 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")) 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, seed: Optional[int] = None, ) -> BasinContext: 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: obs: ZoneObs forecast: ForecastResult config: ForecastConfig = field(default_factory=ForecastConfig) ground_truth: Optional[RiskScore] = None zone_ids: List[str] = field(default_factory=list) adjacency: Dict[str, List[str]] = field(default_factory=dict) data_source: DataSource = DataSource.SYNTHETIC basin_context: Optional[BasinContext] = None 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" ) if self.zone_obs[0].zone_id != self.obs.zone_id: self.obs = self.zone_obs[0] self.forecast = self.zone_forecasts[0] 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]: if self.zone_obs: return list(self.zone_obs) return [self.obs] def resolved_zone_forecasts(self) -> List[ForecastResult]: 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: rng = random.Random(seed if seed is not None else _stable_seed(zone_id)) _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: 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: 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") _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( cfg_era5_back.force_data_source == DataSource.ERA5_REANALYSIS, "ForecastConfig force_data_source=ERA5 round-trip" ) _assert(DataSource.OPENMETEO_ARCHIVE.value == "openmeteo_archive", "OPENMETEO_ARCHIVE enum member missing") _assert(DataSource.OPENMETEO_ARCHIVE.is_observational(), "OPENMETEO_ARCHIVE must be observational") _assert(DataSource.OPENMETEO_ARCHIVE != DataSource.OPENMETEO_LIVE, "ARCHIVE and LIVE must be distinct sources") _assert(DataSource.TIGGE_ECMF.value == "tigge_ecmf", "TIGGE_ECMF enum member missing") _assert(DataSource.TIGGE_ECMF.is_observational(), "TIGGE_ECMF must be observational") fr_tigge = make_synthetic_forecast_result("tigge_zone", horizon_days=15, seed=1) d_tigge = fr_tigge.to_dict() d_tigge["source"] = "tigge_ecmf" d_tigge["horizon_days"] = 15 for k in ForecastResult._SEQUENCE_FIELDS: if k in d_tigge and isinstance(d_tigge[k], (list, tuple)): d_tigge[k] = list(d_tigge[k])[:15] fr_tigge2 = ForecastResult.from_dict(d_tigge) _assert(fr_tigge2.source == DataSource.TIGGE_ECMF, "ForecastResult.from_dict did not accept source=tigge_ecmf") d_unknown = fr_tigge.to_dict() d_unknown["source"] = "not_a_real_source" d_unknown["horizon_days"] = 15 for k in ForecastResult._SEQUENCE_FIELDS: if k in d_unknown and isinstance(d_unknown[k], (list, tuple)): d_unknown[k] = list(d_unknown[k])[:15] fr_unk = ForecastResult.from_dict(d_unknown) _assert(fr_unk.source == DataSource.UNKNOWN, "ForecastResult.from_dict should coerce unknown source to UNKNOWN") cfg_arch = ForecastConfig(force_data_source=DataSource.OPENMETEO_ARCHIVE) cfg_arch_back = ForecastConfig.from_dict(cfg_arch.to_dict()) _assert( cfg_arch_back.force_data_source == DataSource.OPENMETEO_ARCHIVE, "ForecastConfig force_data_source=OPENMETEO_ARCHIVE round-trip" ) 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 _assert(cfg2.belief_prior_weight == 0.70, "belief_prior_weight default round-trip") _assert(cfg2.uncertainty_decay == 0.70, "uncertainty_decay default round-trip") _assert(cfg2.info_gain_scale == 5.0, "info_gain_scale default round-trip") _assert(cfg2.uncertainty_penalty_scale == 5.0, "uncertainty_penalty_scale default round-trip") cfg_belief = ForecastConfig( belief_prior_weight=0.35, uncertainty_decay=0.5, info_gain_scale=2.0, uncertainty_penalty_scale=8.0, ) cfg_belief_back = ForecastConfig.from_dict(cfg_belief.to_dict()) _assert(cfg_belief_back.belief_prior_weight == 0.35, "non-default belief_prior_weight round-trip") _assert(cfg_belief_back.uncertainty_decay == 0.5, "non-default uncertainty_decay round-trip") _assert(cfg_belief_back.info_gain_scale == 2.0, "non-default info_gain_scale round-trip") _assert(cfg_belief_back.uncertainty_penalty_scale == 8.0, "non-default uncertainty_penalty_scale round-trip") try: ForecastConfig(belief_prior_weight=1.5) _assert(False, "ForecastConfig accepted belief_prior_weight out of [0,1]") except ValueError: pass try: ForecastConfig(info_gain_scale=-1.0) _assert(False, "ForecastConfig accepted negative info_gain_scale") 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") # 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, 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.")