""" 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: 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) 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: 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 # --- 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 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 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) 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 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) ) 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, "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: 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) _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: 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: 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")) 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: 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") 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" ) 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.")