""" era5_data_pipeline.py ===================== """ from __future__ import annotations import hashlib import json import logging import math import os import random from datetime import datetime, timedelta, timezone from pathlib import Path from typing import Any, Dict, List, Optional, Tuple import zone_observation as _zo assert _zo.SCHEMA_VERSION == 3, ( f"Schema mismatch: expected 3, got {_zo.SCHEMA_VERSION}" ) from zone_observation import ( BasinContext, DataSource, EpisodeContext, ForecastConfig, GeoPolygon, ZoneObs, derive_helio_regime, make_synthetic_basin_context, make_synthetic_episode_context, make_synthetic_zone_obs, _stable_seed, ) logger = logging.getLogger(__name__) # --------------------------------------------------------------------------- # Optional dependencies # --------------------------------------------------------------------------- try: import requests REQUESTS_AVAILABLE = True except ImportError: REQUESTS_AVAILABLE = False logger.warning("requests not installed — Open-Meteo unavailable") try: import cdsapi CDSAPI_AVAILABLE = True except ImportError: CDSAPI_AVAILABLE = False logger.info("cdsapi not installed — ERA5 will fall back to synthetic") try: import numpy as np NUMPY_AVAILABLE = True except ImportError: NUMPY_AVAILABLE = False try: import ee # Google Earth Engine — IMERG / CHIRPS / SMAP access EE_AVAILABLE = True except ImportError: EE_AVAILABLE = False logger.info("earthengine-api not installed — satellite sources unavailable") _EE_INITIALIZED = False def _ensure_ee_initialized() -> None: """Lazily call ee.Initialize() once per process.""" global _EE_INITIALIZED if _EE_INITIALIZED: return if not EE_AVAILABLE: raise RuntimeError("earthengine-api not installed") project = os.environ.get("EARTHENGINE_PROJECT") if project: ee.Initialize(project=project) else: ee.Initialize() _EE_INITIALIZED = True # --------------------------------------------------------------------------- # Zone registry # --------------------------------------------------------------------------- _ZONE_REGISTRY: Dict[str, GeoPolygon] = {} def register_zone(polygon: GeoPolygon) -> None: """Register a sourcing zone polygon for lat/lon resolution.""" _ZONE_REGISTRY[polygon.zone_id] = polygon logger.info( f"Registered zone {polygon.zone_id} centroid={polygon.centroid}" ) def _resolve_latlon(zone_id: str) -> Tuple[float, float]: if zone_id not in _ZONE_REGISTRY: raise KeyError( f"Zone '{zone_id}' not registered. " f"Call register_zone() before fetching data." ) return _ZONE_REGISTRY[zone_id].centroid # --------------------------------------------------------------------------- # Config / cache # --------------------------------------------------------------------------- _CACHE_DIR = Path(os.environ.get("WEATHER_CACHE_DIR", ".cache/era5")) _CACHE_DIR.mkdir(parents=True, exist_ok=True) _ERA5_CACHE_DIR = _CACHE_DIR / "era5_nc" _ERA5_CACHE_DIR.mkdir(parents=True, exist_ok=True) _OPENMETEO_URL = "https://api.open-meteo.com/v1/forecast" _OPENMETEO_ARCHIVE_URL = "https://archive-api.open-meteo.com/v1/archive" _TIMEOUT_S = int(os.environ.get("WEATHER_HTTP_TIMEOUT", "30")) _CACHE_TTL_DAYS = int(os.environ.get("WEATHER_CACHE_TTL_DAYS", "7")) _ERA5_TTL_DAYS = int(os.environ.get("WEATHER_ERA5_TTL_DAYS", "30")) # ERA5 bounding box padding in degrees around zone centroid _ERA5_BOX_PAD = float(os.environ.get("WEATHER_ERA5_BOX_PAD", "0.5")) # --------------------------------------------------------------------------- # Cache utilities # --------------------------------------------------------------------------- def _cache_key(url: str, params: Dict[str, Any]) -> str: payload = f"{url}{json.dumps(params, sort_keys=True)}" return hashlib.sha256(payload.encode()).hexdigest() def _cached_get(url: str, params: Dict[str, Any]) -> Dict[str, Any]: """HTTP GET with file-based JSON cache.""" if not REQUESTS_AVAILABLE: raise RuntimeError("requests not installed — cannot fetch HTTP data") path = _CACHE_DIR / f"{_cache_key(url, params)}.json" if path.exists(): age = ( datetime.now(timezone.utc) - datetime.fromtimestamp(path.stat().st_mtime, tz=timezone.utc) ).days if age < _CACHE_TTL_DAYS: with open(path) as f: return json.load(f) resp = requests.get(url, params=params, timeout=_TIMEOUT_S) resp.raise_for_status() data = resp.json() with open(path, "w") as f: json.dump(data, f) return data def _era5_cache_path(zone_id: str, date_range: Tuple[datetime, datetime]) -> Path: """Deterministic cache file path for an ERA5 download.""" key = _stable_seed( zone_id + date_range[0].date().isoformat() + date_range[1].date().isoformat() ) return _ERA5_CACHE_DIR / f"{zone_id}_{key}.nc" # --------------------------------------------------------------------------- # Timezone helper # --------------------------------------------------------------------------- def _ensure_utc(dt: datetime) -> datetime: """Normalise a datetime to UTC at the pipeline boundary.""" if dt.tzinfo is None: return dt.replace(tzinfo=timezone.utc) return dt.astimezone(timezone.utc) _ERA5_VARIABLES: List[str] = [ "total_precipitation", "2m_temperature", "2m_dewpoint_temperature", "10m_u_component_of_wind", "10m_v_component_of_wind", "surface_pressure", "volumetric_soil_water_layer_1", "potential_evaporation", ] def _dewpoint_to_rh(temp_c: float, dewpoint_c: float) -> float: """Magnus formula: relative humidity from temperature and dewpoint (%).""" a, b = 17.625, 243.04 # Magnus coefficients rh = 100.0 * math.exp( (a * dewpoint_c / (b + dewpoint_c)) - (a * temp_c / (b + temp_c)) ) return max(0.0, min(100.0, rh)) def _build_era5_obs( zone_id: str, start: datetime, nc_path: Path, ) -> ZoneObs: """Map a downloaded ERA5 NetCDF file to a ZoneObs.""" try: import netCDF4 as nc # type: ignore except ImportError: try: import xarray as xr # type: ignore return _build_era5_obs_xarray(zone_id, start, nc_path) except ImportError: raise ImportError( "netCDF4 or xarray required for ERA5 ingestion. " "Install with: pip install netCDF4 or pip install xarray" ) import numpy as np ds = nc.Dataset(str(nc_path), "r") try: def _daily_mean(var: str) -> float: if var not in ds.variables: return 0.0 data = ds.variables[var][:24].flatten() data = np.ma.filled(data, np.nan) valid = data[np.isfinite(data)] return float(np.mean(valid)) if len(valid) > 0 else 0.0 def _daily_max(var: str) -> float: if var not in ds.variables: return 0.0 data = ds.variables[var][:24].flatten() data = np.ma.filled(data, np.nan) valid = data[np.isfinite(data)] return float(np.max(valid)) if len(valid) > 0 else 0.0 def _daily_sum(var: str, scale: float = 1.0) -> float: if var not in ds.variables: return 0.0 data = ds.variables[var][:24].flatten() data = np.ma.filled(data, np.nan) valid = data[np.isfinite(data)] return float(np.sum(valid) * scale) if len(valid) > 0 else 0.0 temp_mean_c = _daily_mean("t2m") - 273.15 temp_max_c = _daily_max("t2m") - 273.15 if "t2m" in ds.variables: t_data = np.ma.filled(ds.variables["t2m"][:24].flatten(), np.nan) valid = t_data[np.isfinite(t_data)] temp_min_c = float(np.min(valid)) - 273.15 if len(valid) > 0 else temp_mean_c else: temp_min_c = temp_mean_c dewpoint_mean_c = _daily_mean("d2m") - 273.15 rh_mean = _dewpoint_to_rh(temp_mean_c, dewpoint_mean_c) dewpoint_max_c = _daily_max("d2m") - 273.15 rh_max = _dewpoint_to_rh(temp_min_c, dewpoint_max_c) precip_24h = _daily_sum("tp", scale=1000.0) total_hours = len(ds.variables.get("tp", [])) if "tp" in ds.variables else 24 if total_hours >= 24 * 7: def _window_sum(var: str, hours: int, scale: float = 1.0) -> float: if var not in ds.variables: return 0.0 data = np.ma.filled( ds.variables[var][:hours].flatten(), np.nan ) valid = data[np.isfinite(data)] return float(np.sum(valid) * scale) if len(valid) > 0 else 0.0 precip_7d = _window_sum("tp", 24*7, scale=1000.0) precip_14d = _window_sum("tp", 24*14, scale=1000.0) precip_30d = _window_sum("tp", 24*30, scale=1000.0) else: precip_7d = precip_24h * 7.0 precip_14d = precip_24h * 14.0 precip_30d = precip_24h * 30.0 u = _daily_mean("u10") v = _daily_mean("v10") wind_mean = math.sqrt(u**2 + v**2) u_max = _daily_max("u10") v_max = _daily_max("v10") wind_max = math.sqrt(u_max**2 + v_max**2) soil_pct = _daily_mean("swvl1") * 100.0 pe_m = _daily_sum("pev", scale=1.0) et0_mm = abs(pe_m) * 1000.0 return ZoneObs( zone_id=zone_id, valid_time=start, source=DataSource.ERA5_REANALYSIS, precip_24h_mm=max(0.0, precip_24h), precip_7d_mm=max(0.0, precip_7d), precip_14d_mm=max(0.0, precip_14d), precip_30d_mm=max(0.0, precip_30d), temp_mean_c=temp_mean_c, temp_max_c=max(temp_mean_c, temp_max_c), temp_min_c=min(temp_mean_c, temp_min_c), temp_anomaly_idx=0.0, precip_anomaly_idx=0.0, evapotranspiration_mm=max(0.0, et0_mm), wind_speed_mean_ms=max(0.0, wind_mean), wind_speed_max_ms=max(0.0, wind_max), rh_mean_pct=rh_mean, rh_max_pct=max(rh_mean, rh_max), soil_moisture_pct=max(0.0, min(100.0, soil_pct)), soil_moisture_anom=0.0, quality_flag=0, ) finally: ds.close() def _build_era5_obs_xarray( zone_id: str, start: datetime, nc_path: Path, ) -> ZoneObs: import xarray as xr import numpy as np ds = xr.open_dataset(str(nc_path)) try: def _mean(var: str) -> float: if var not in ds: return 0.0 return float(ds[var].values.flatten()[np.isfinite( ds[var].values.flatten() )].mean()) if len(ds[var].values.flatten()) > 0 else 0.0 def _max(var: str) -> float: if var not in ds: return 0.0 vals = ds[var].values.flatten() valid = vals[np.isfinite(vals)] return float(valid.max()) if len(valid) > 0 else 0.0 def _sum(var: str, scale: float = 1.0) -> float: if var not in ds: return 0.0 vals = ds[var].values.flatten() valid = vals[np.isfinite(vals)] return float(valid.sum() * scale) if len(valid) > 0 else 0.0 temp_mean_c = _mean("t2m") - 273.15 temp_max_c = _max("t2m") - 273.15 if "t2m" in ds: _t2m_flat = ds["t2m"].values.flatten() _t2m_valid = _t2m_flat[np.isfinite(_t2m_flat)] temp_min_c = float(np.min(_t2m_valid)) - 273.15 if len(_t2m_valid) > 0 else temp_mean_c else: temp_min_c = temp_mean_c dewpoint_mean_c = _mean("d2m") - 273.15 dewpoint_max_c = _max("d2m") - 273.15 rh_mean = _dewpoint_to_rh(temp_mean_c, dewpoint_mean_c) rh_max = _dewpoint_to_rh(temp_min_c, dewpoint_max_c) precip_24h = max(0.0, _sum("tp", scale=1000.0)) if "tp" in ds: total_hours = len(ds["tp"].values.flatten()) else: total_hours = 0 if total_hours >= 24 * 7: def _window_sum_xr(var: str, hours: int, scale: float = 1.0) -> float: if var not in ds: return 0.0 vals = ds[var].values.flatten()[:hours] valid = vals[np.isfinite(vals)] return float(valid.sum() * scale) if len(valid) > 0 else 0.0 precip_7d = max(0.0, _window_sum_xr("tp", 24 * 7, scale=1000.0)) precip_14d = max(0.0, _window_sum_xr("tp", 24 * 14, scale=1000.0)) precip_30d = max(0.0, _window_sum_xr("tp", 24 * 30, scale=1000.0)) else: precip_7d = precip_24h * 7.0 precip_14d = precip_24h * 14.0 precip_30d = precip_24h * 30.0 u = _mean("u10") v = _mean("v10") wind_mean = math.sqrt(u**2 + v**2) wind_max = math.sqrt(_max("u10")**2 + _max("v10")**2) soil_pct = _mean("swvl1") * 100.0 et0_mm = abs(_sum("pev")) * 1000.0 return ZoneObs( zone_id=zone_id, valid_time=start, source=DataSource.ERA5_REANALYSIS, precip_24h_mm=precip_24h, precip_7d_mm=precip_7d, precip_14d_mm=precip_14d, precip_30d_mm=precip_30d, temp_mean_c=temp_mean_c, temp_max_c=max(temp_mean_c, temp_max_c), temp_min_c=min(temp_mean_c, temp_min_c), evapotranspiration_mm=max(0.0, et0_mm), wind_speed_mean_ms=max(0.0, wind_mean), wind_speed_max_ms=max(0.0, wind_max), rh_mean_pct=rh_mean, rh_max_pct=max(rh_mean, rh_max), soil_moisture_pct=max(0.0, min(100.0, soil_pct)), quality_flag=0, ) finally: ds.close() def _fetch_openmeteo(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs: """Fetch from Open-Meteo API with full variable coverage.""" lat, lon = _resolve_latlon(zone_id) start = _ensure_utc(date_range[0]) end = _ensure_utc(date_range[1]) today = datetime.now(timezone.utc).date() use_archive = start.date() < today - timedelta(days=5) url = _OPENMETEO_ARCHIVE_URL if use_archive else _OPENMETEO_URL daily_vars = ",".join([ "temperature_2m_mean", "temperature_2m_max", "temperature_2m_min", "precipitation_sum", "precipitation_hours", "wind_speed_10m_mean", "wind_speed_10m_max", "relative_humidity_2m_mean", "relative_humidity_2m_max", "et0_fao_evapotranspiration", "shortwave_radiation_sum", "soil_moisture_0_to_7cm_mean", ]) params = { "latitude": lat, "longitude": lon, "daily": daily_vars, "start_date": start.date().isoformat(), "end_date": end.date().isoformat(), "timezone": "UTC", } data = _cached_get(url, params) daily = data.get("daily", {}) def _safe(key: str, i: int = 0) -> float: vals = daily.get(key, []) return float(vals[i]) if i < len(vals) and vals[i] is not None else 0.0 def _safe_sum(key: str, n: int) -> float: vals = daily.get(key, []) return float(sum( v for v in vals[:n] if v is not None )) temp_mean = _safe("temperature_2m_mean") temp_max = _safe("temperature_2m_max") temp_min = _safe("temperature_2m_min") temp_max = max(temp_mean, temp_max) temp_min = min(temp_mean, temp_min) rh_mean = _safe("relative_humidity_2m_mean") rh_max = _safe("relative_humidity_2m_max") rh_max = max(rh_mean, rh_max) wind_mean = _safe("wind_speed_10m_mean") wind_max = _safe("wind_speed_10m_max") wind_max = max(wind_mean, wind_max) return ZoneObs( zone_id=zone_id, valid_time=start, source=DataSource.OPENMETEO_LIVE, precip_24h_mm=_safe("precipitation_sum"), precip_7d_mm=_safe_sum("precipitation_sum", 7), precip_14d_mm=_safe_sum("precipitation_sum", 14), precip_30d_mm=_safe_sum("precipitation_sum", 30), temp_mean_c=temp_mean, temp_max_c=temp_max, temp_min_c=temp_min, evapotranspiration_mm=_safe("et0_fao_evapotranspiration"), wind_speed_mean_ms=wind_mean / 3.6, wind_speed_max_ms=wind_max / 3.6, rh_mean_pct=rh_mean, rh_max_pct=rh_max, soil_moisture_pct=max(0.0, min(100.0, _safe("soil_moisture_0_to_7cm_mean") * 100.0)), quality_flag=0, ) def _fetch_era5(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs: if not CDSAPI_AVAILABLE: logger.debug("cdsapi unavailable — falling back to synthetic for %s", zone_id) return _fetch_synthetic(zone_id, date_range) lat, lon = _resolve_latlon(zone_id) start = _ensure_utc(date_range[0]) end = _ensure_utc(date_range[1]) bbox = [ round(lat + _ERA5_BOX_PAD, 2), round(lon - _ERA5_BOX_PAD, 2), round(lat - _ERA5_BOX_PAD, 2), round(lon + _ERA5_BOX_PAD, 2), ] nc_path = _era5_cache_path(zone_id, date_range) if nc_path.exists(): age_days = ( datetime.now(timezone.utc) - datetime.fromtimestamp(nc_path.stat().st_mtime, tz=timezone.utc) ).days if age_days < _ERA5_TTL_DAYS: logger.debug("ERA5 cache hit for %s", zone_id) else: nc_path.unlink() if not nc_path.exists(): years: set = set() months: set = set() days: set = set() cursor = start.date() end_date = end.date() while cursor <= end_date: years.add(cursor.year) months.add(cursor.month) days.add(cursor.day) cursor += timedelta(days=1) logger.info( "Requesting ERA5 data for %s bbox=%s dates=%s to %s", zone_id, bbox, start.date().isoformat(), end.date().isoformat(), ) try: client = cdsapi.Client(quiet=True) client.retrieve( "reanalysis-era5-single-levels", { "product_type": "reanalysis", "variable": _ERA5_VARIABLES, "year": sorted(str(y) for y in years), "month": sorted(f"{m:02d}" for m in months), "day": sorted(f"{d:02d}" for d in days), "time": [f"{h:02d}:00" for h in range(24)], "area": bbox, "format": "netcdf", }, str(nc_path), ) logger.info("ERA5 download complete: %s", nc_path) except Exception as e: logger.warning( "ERA5 CDS request failed for %s: %s — falling back to synthetic", zone_id, e ) if nc_path.exists(): nc_path.unlink() return _fetch_synthetic(zone_id, date_range) try: return _build_era5_obs(zone_id, start, nc_path) except Exception as e: logger.warning( "ERA5 NetCDF parse failed for %s: %s — falling back to synthetic", zone_id, e ) return _fetch_synthetic(zone_id, date_range) def _fetch_imerg(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs: _ensure_ee_initialized() lat, lon = _resolve_latlon(zone_id) start = _ensure_utc(date_range[0]) region = ee.Geometry.Point([lon, lat]).buffer(_ERA5_BOX_PAD * 111_000) def _window_sum_mm(n_days: int) -> float: window_end = start + timedelta(days=n_days) coll = ( ee.ImageCollection("NASA/GPM_L3/IMERG_V07") .filterDate(start.isoformat(), window_end.isoformat()) .filterBounds(region) .select("precipitation") ) total_mm_image = coll.sum().multiply(0.5) stats = total_mm_image.reduceRegion( reducer=ee.Reducer.mean(), geometry=region, scale=11_000, bestEffort=True, ).getInfo() return float(stats.get("precipitation", 0.0) or 0.0) precip_24h = max(0.0, _window_sum_mm(1)) precip_7d = max(precip_24h, _window_sum_mm(7)) precip_14d = max(precip_7d, _window_sum_mm(14)) precip_30d = max(precip_14d, _window_sum_mm(30)) return ZoneObs( zone_id=zone_id, valid_time=start, source=DataSource.SATELLITE_PRECIP, precip_24h_mm=precip_24h, precip_7d_mm=precip_7d, precip_14d_mm=precip_14d, precip_30d_mm=precip_30d, precip_satellite_mm=precip_24h, quality_flag=0, ) def _fetch_smap(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs: _ensure_ee_initialized() lat, lon = _resolve_latlon(zone_id) start = _ensure_utc(date_range[0]) end = start + timedelta(days=1) region = ee.Geometry.Point([lon, lat]).buffer(_ERA5_BOX_PAD * 111_000) coll = ( ee.ImageCollection("NASA/SMAP/SPL4SMGP/007") .filterDate(start.isoformat(), end.isoformat()) .filterBounds(region) .select("sm_surface") ) stats = coll.mean().reduceRegion( reducer=ee.Reducer.mean(), geometry=region, scale=9_000, bestEffort=True, ).getInfo() vwc = stats.get("sm_surface") if vwc is None: raise RuntimeError(f"SMAP: no data returned for zone={zone_id} date={start.date()}") soil_pct = max(0.0, min(100.0, float(vwc) * 100.0)) return ZoneObs( zone_id=zone_id, valid_time=start, source=DataSource.SATELLITE_SOIL, soil_moisture_pct=soil_pct, soil_moisture_satellite_pct=soil_pct, quality_flag=0, ) def _fetch_synthetic(zone_id: str, date_range: Tuple[datetime, datetime]) -> ZoneObs: seed = _stable_seed(zone_id + date_range[0].isoformat()) return make_synthetic_zone_obs(zone_id=zone_id, seed=seed) def _inject_noise(obs: ZoneObs, rng: random.Random, scale: float) -> ZoneObs: def perturb(x: float) -> float: return x * (1.0 + rng.uniform(-scale, scale)) base = obs.to_dict() base.pop("_schema_version", None) base["precip_24h_mm"] = max(0.0, perturb(obs.precip_24h_mm)) base["precip_7d_mm"] = max(0.0, perturb(obs.precip_7d_mm)) base["precip_14d_mm"] = max(0.0, perturb(obs.precip_14d_mm)) base["precip_30d_mm"] = max(0.0, perturb(obs.precip_30d_mm)) base["temp_mean_c"] = perturb(obs.temp_mean_c) base["temp_max_c"] = perturb(obs.temp_max_c) base["temp_min_c"] = perturb(obs.temp_min_c) base["rh_mean_pct"] = min(100.0, max(0.0, perturb(obs.rh_mean_pct))) base["rh_max_pct"] = min(100.0, max(0.0, perturb(obs.rh_max_pct))) base["wind_speed_mean_ms"] = max(0.0, perturb(obs.wind_speed_mean_ms)) base["wind_speed_max_ms"] = max(0.0, perturb(obs.wind_speed_max_ms)) p24 = base["precip_24h_mm"] p7 = max(base["precip_7d_mm"], p24) p14 = max(base["precip_14d_mm"], p7) p30 = max(base["precip_30d_mm"], p14) base["precip_7d_mm"] = p7 base["precip_14d_mm"] = p14 base["precip_30d_mm"] = p30 t_mean = base["temp_mean_c"] base["temp_max_c"] = max(t_mean, base["temp_max_c"]) base["temp_min_c"] = min(t_mean, base["temp_min_c"]) base["rh_max_pct"] = max(base["rh_mean_pct"], base["rh_max_pct"]) base["wind_speed_max_ms"] = max( base["wind_speed_mean_ms"], base["wind_speed_max_ms"] ) return ZoneObs.from_dict(base) _NOAA_RONI_URL = "https://www.cpc.ncep.noaa.gov/data/indices/RONI.ascii.txt" _PSL_DMI_URLS: Tuple[str, ...] = ( "https://psl.noaa.gov/gcos_wgsp/Timeseries/Data/dmi.had.long.data", "https://psl.noaa.gov/data/timeseries/month/data/dmi.had.long.data", "https://psl.noaa.gov/data/correlation/dmi.data", ) _SWPC_PLASMA_URLS: Tuple[str, ...] = ( "https://services.swpc.noaa.gov/json/rtsw/rtsw_wind_1m.json", "https://services.swpc.noaa.gov/products/summary/solar-wind-speed.json", "https://services.swpc.noaa.gov/products/solar-wind/plasma-1-day.json", ) _SWPC_KP_URL = "https://services.swpc.noaa.gov/json/planetary_k_index_1m.json" _SWPC_XRAY_URL = "https://services.swpc.noaa.gov/json/goes/primary/xrays-1-day.json" _SEASON_FOR_MONTH: Dict[int, str] = { 1: "DJF", 2: "JFM", 3: "FMA", 4: "MAM", 5: "AMJ", 6: "MJJ", 7: "JJA", 8: "JAS", 9: "ASO", 10: "SON", 11: "OND", 12: "NDJ", } _SEASON_ORDER: Tuple[str, ...] = ( "DJF", "JFM", "FMA", "MAM", "AMJ", "MJJ", "JJA", "JAS", "ASO", "SON", "OND", "NDJ", ) def _parse_cpc_seasonal_ascii(text: str) -> Dict[Tuple[int, str], float]: """Parse NOAA CPC's seasonal index ASCII format.""" out: Dict[Tuple[int, str], float] = {} lines = text.strip().splitlines() for line in lines[1:]: parts = line.split() if len(parts) < 3: continue season = parts[0] try: year = int(parts[1]) value = float(parts[-1]) except (ValueError, IndexError): continue out[(year, season)] = value return out def _parse_psl_monthly_ascii( text: str, year: int, month: int, missing_below: float = -90.0 ) -> Optional[float]: for line in text.strip().splitlines(): parts = line.split() if len(parts) != 13: continue try: row_year = int(parts[0]) values = [float(v) for v in parts[1:]] except ValueError: continue if row_year == year: v = values[month - 1] return None if v <= missing_below else v return None def _lookup_roni_with_lag( table: Dict[Tuple[int, str], float], year: int, season: str ) -> Tuple[float, str]: if (year, season) in table: return table[(year, season)], f"{season} {year}" try: idx = _SEASON_ORDER.index(season) except ValueError: idx = 0 y, i = year, idx for _ in range(24): i -= 1 if i < 0: i = len(_SEASON_ORDER) - 1 y -= 1 key = (y, _SEASON_ORDER[i]) if key in table: return table[key], f"{_SEASON_ORDER[i]} {y} (lagged from {season} {year})" raise ValueError(f"No RONI value for {season} {year} or any prior season in table") def _lookup_dmi_with_lag( text: str, year: int, month: int ) -> Tuple[float, str]: y, m = year, month for _ in range(24): v = _parse_psl_monthly_ascii(text, y, m) if v is not None: label = f"{y}-{m:02d}" if (y, m) != (year, month): label += f" (lagged from {year}-{month:02d})" return v, label m -= 1 if m < 1: m = 12 y -= 1 raise ValueError(f"No DMI value for {year}-{month:02d} or any prior month") def _extract_solar_wind_speeds(payload: Any) -> List[float]: speeds: List[float] = [] def _from_row(row: Any) -> None: if not isinstance(row, dict): return for key in ("proton_speed", "speed", "wind_speed", "value"): sp = row.get(key) if sp is None: continue try: f = float(sp) except (TypeError, ValueError): continue if f > 0: speeds.append(f) return if isinstance(payload, list): for row in payload: if isinstance(row, dict): _from_row(row) elif isinstance(row, (list, tuple)) and len(row) >= 2: try: f = float(row[1]) if f > 0: speeds.append(f) except (TypeError, ValueError): continue elif isinstance(payload, dict): _from_row(payload) for v in payload.values(): if isinstance(v, list): speeds.extend(_extract_solar_wind_speeds(v)) return speeds def _fetch_swpc_helio(valid_date: datetime) -> Dict[str, Any]: out: Dict[str, Any] = {} if not REQUESTS_AVAILABLE: return out plasma_err: Optional[Exception] = None for url in _SWPC_PLASMA_URLS: try: resp = requests.get(url, timeout=_TIMEOUT_S) resp.raise_for_status() speeds = _extract_solar_wind_speeds(resp.json()) if speeds: out["solar_wind_speed_kms"] = speeds[-1] plasma_err = None break plasma_err = ValueError(f"no speed samples in {url}") except Exception as e: plasma_err = e continue if "solar_wind_speed_kms" not in out and plasma_err is not None: logger.warning( "_fetch_swpc_helio: plasma fetch failed (%s) for %s", plasma_err, valid_date.date(), ) try: resp = requests.get(_SWPC_KP_URL, timeout=_TIMEOUT_S) resp.raise_for_status() rows = resp.json() kps = [] for row in rows: try: kp = row.get("kp_index", row.get("kp")) if kp is not None: kps.append(float(kp)) except (TypeError, ValueError, AttributeError): continue if kps: out["kp_index"] = kps[-1] except Exception as e: logger.warning( "_fetch_swpc_helio: Kp fetch failed (%s) for %s", e, valid_date.date(), ) try: resp = requests.get(_SWPC_XRAY_URL, timeout=_TIMEOUT_S) resp.raise_for_status() rows = resp.json() fluxes = [] for row in rows: try: flux = row.get("flux") energy = str(row.get("energy", "")).lower() if flux is not None and float(flux) > 0: if "0.1-0.8" in energy or "long" in energy or not energy: fluxes.append(float(flux)) except (TypeError, ValueError, AttributeError): continue if fluxes: out["goes_xray_flux"] = fluxes[-1] except Exception as e: logger.warning( "_fetch_swpc_helio: GOES X-ray fetch failed (%s) for %s", e, valid_date.date(), ) if out: kp = float(out.get("kp_index", 2.0)) xray = float(out.get("goes_xray_flux", 1e-7)) out["helio_regime"] = derive_helio_regime(kp, xray) return out def fetch_basin_context( valid_date: datetime, config: Optional[ForecastConfig] = None, ) -> BasinContext: cfg = config or ForecastConfig() valid_date = _ensure_utc(valid_date) season = _SEASON_FOR_MONTH[valid_date.month] strict = bool(getattr(cfg, "require_real_basin_context", False)) if not REQUESTS_AVAILABLE: if strict: raise RuntimeError( "require_real_basin_context=True but the 'requests' package " "is not installed — cannot fetch live basin/helio indices." ) logger.info("requests not installed — basin context falls back to synthetic") return make_synthetic_basin_context(valid_date=valid_date) try: resp = requests.get(_NOAA_RONI_URL, timeout=_TIMEOUT_S) resp.raise_for_status() table = _parse_cpc_seasonal_ascii(resp.text) enso_oni, roni_label = _lookup_roni_with_lag(table, valid_date.year, season) if "lagged" in roni_label: logger.info( "fetch_basin_context: RONI using %s for %s", roni_label, valid_date.date(), ) except Exception as e: if strict: raise RuntimeError( f"require_real_basin_context=True: RONI fetch failed for " f"{valid_date.date()}: {e}" ) from e logger.warning( "fetch_basin_context: RONI fetch/parse failed (%s) — synthetic ENSO for %s", e, valid_date.date(), ) enso_oni = make_synthetic_basin_context( valid_date=valid_date, seed=_stable_seed(f"enso_{valid_date.date().isoformat()}") ).enso_oni iod_dmi: Optional[float] = None dmi_err: Optional[Exception] = None for dmi_url in _PSL_DMI_URLS: try: resp = requests.get(dmi_url, timeout=_TIMEOUT_S) resp.raise_for_status() iod_dmi, dmi_label = _lookup_dmi_with_lag( resp.text, valid_date.year, valid_date.month ) if "lagged" in dmi_label: logger.info( "fetch_basin_context: DMI using %s for %s (from %s)", dmi_label, valid_date.date(), dmi_url, ) dmi_err = None break except Exception as e: dmi_err = e continue if iod_dmi is None: if strict: raise RuntimeError( f"require_real_basin_context=True: DMI fetch failed for " f"{valid_date.date()}: {dmi_err}" ) logger.warning( "fetch_basin_context: DMI fetch/parse failed (%s) — synthetic IOD for %s", dmi_err, valid_date.date(), ) iod_dmi = make_synthetic_basin_context( valid_date=valid_date, seed=_stable_seed(f"iod_{valid_date.date().isoformat()}") ).iod_dmi helio = _fetch_swpc_helio(valid_date) required_helio = ("solar_wind_speed_kms", "kp_index", "goes_xray_flux") missing_helio = [k for k in required_helio if k not in helio] if strict and missing_helio: raise RuntimeError( f"require_real_basin_context=True: SWPC helio incomplete for " f"{valid_date.date()} — missing {missing_helio}. " f"Got keys: {sorted(helio.keys())}" ) return BasinContext( valid_date=valid_date, enso_oni=enso_oni, iod_dmi=iod_dmi, solar_wind_speed_kms=float(helio.get("solar_wind_speed_kms", 400.0)), kp_index=float(helio.get("kp_index", 2.0)), goes_xray_flux=float(helio.get("goes_xray_flux", 1e-7)), helio_regime=str(helio.get("helio_regime", "quiet")), source=DataSource.PUBLISHED_INDEX, ) def _select_source(cfg: ForecastConfig, rng: random.Random) -> DataSource: real_ratio = getattr(cfg, "real_data_ratio", 0.7) era5_ratio = getattr(cfg, "era5_ratio", 0.5) if rng.random() > real_ratio: return DataSource.SYNTHETIC satellite_options = [] if getattr(cfg, "use_satellite_precip", False): satellite_options.append(DataSource.SATELLITE_PRECIP) if getattr(cfg, "use_satellite_soil", False): satellite_options.append(DataSource.SATELLITE_SOIL) if satellite_options: return satellite_options[rng.randrange(len(satellite_options))] return ( DataSource.ERA5_REANALYSIS if rng.random() < era5_ratio else DataSource.OPENMETEO_LIVE ) def _fallback_chain(primary: DataSource) -> List[DataSource]: if primary in (DataSource.SATELLITE_PRECIP, DataSource.SATELLITE_SOIL): return [primary, DataSource.ERA5_REANALYSIS, DataSource.SYNTHETIC] return [primary, DataSource.SYNTHETIC] def fetch_zone_obs( zone_id: str, date_range: Tuple[datetime, datetime], config: Optional[ForecastConfig] = None, ) -> ZoneObs: cfg = config or ForecastConfig() seed = _stable_seed(zone_id + date_range[0].isoformat()) rng = random.Random(seed) if getattr(cfg, "force_data_source", None) is not None: primary = cfg.force_data_source else: primary = _select_source(cfg, rng) fetchers = { DataSource.OPENMETEO_LIVE: _fetch_openmeteo, DataSource.ERA5_REANALYSIS: _fetch_era5, DataSource.SATELLITE_PRECIP: _fetch_imerg, DataSource.SATELLITE_SOIL: _fetch_smap, DataSource.SYNTHETIC: _fetch_synthetic, } obs: Optional[ZoneObs] = None for source in _fallback_chain(primary): try: obs = fetchers[source](zone_id, date_range) ZoneObs.validate(obs, strict=True) break except Exception as e: logger.warning("%s failed for %s: %s", source.value, zone_id, e) obs = None if obs is None: raise RuntimeError( f"All data sources failed for zone '{zone_id}'. " f"Check zone registration and network connectivity." ) if getattr(cfg, "inject_noise", False): obs = _inject_noise(obs, rng, cfg.noise_scale) ZoneObs.validate(obs, strict=True) if getattr(cfg, "use_climatology_anomalies", False): try: _lat, _lon = _resolve_latlon(zone_id) except KeyError: logger.warning( "use_climatology_anomalies=True but zone '%s' is not " "registered -- cannot locate the zone for its climatology. " "Anomaly fields left at fetcher defaults. Call " "register_zone() first.", zone_id, ) else: from climatology import ( apply_climatology_anomalies, get_zone_climatology, ) _clim = get_zone_climatology( zone_id, _lat, _lon, years=getattr(cfg, "climatology_years", 10), end_year=date_range[0].year - 1, ) obs = apply_climatology_anomalies(obs, _clim) ZoneObs.validate(obs, strict=True) return obs def _build_context_forecast( zone_id: str, obs: ZoneObs, cfg: ForecastConfig, seed: int, ) -> "ForecastResult": from zone_observation import ForecastResult # noqa: F401 (type hint only) mode = getattr(cfg, "forecast_backend", "synthetic") def _synthetic() -> "ForecastResult": from zone_observation import make_synthetic_forecast_result return make_synthetic_forecast_result( zone_id=zone_id, valid_time=obs.valid_time, horizon_days=cfg.horizon_days, seed=seed, ) if mode == "synthetic": return _synthetic() if mode == "timesfm": logger.warning( "forecast_backend='timesfm' cannot be constructed inside " "fetch_episode_context (no checkpoint/sha available at this " "layer) -- falling back to synthetic for zone=%s. Build the " "backend via timesfm_wrapper.create_forecast_backend() and " "call backend.forecast(obs) directly if you need TimesFM.", zone_id, ) return _synthetic() try: from timesfm_wrapper import create_forecast_backend lat = lon = None if mode == "openmeteo": lat, lon = _resolve_latlon(zone_id) backend = create_forecast_backend( mode=mode, lat=lat, lon=lon, horizon_days=cfg.horizon_days, ) return backend.forecast(obs) except Exception as e: logger.warning( "forecast backend %r failed for zone=%s (%s) -- synthetic " "forecast fallback", mode, zone_id, e, ) return _synthetic() def fetch_episode_context( zone_id: str, date_range: Tuple[datetime, datetime], config: Optional[ForecastConfig] = None, ) -> EpisodeContext: cfg = config or ForecastConfig() obs = fetch_zone_obs(zone_id, date_range, cfg) seed = _stable_seed(zone_id + date_range[0].isoformat() + "_forecast") fcast = _build_context_forecast(zone_id, obs, cfg, seed) basin_context = None if getattr(cfg, "include_basin_context", False): strict_basin = bool(getattr(cfg, "require_real_basin_context", False)) try: basin_context = fetch_basin_context(date_range[0], cfg) except Exception as e: if strict_basin: raise logger.warning( "fetch_episode_context: basin context fetch failed (%s) — " "leaving basin_context=None for zone=%s", e, zone_id, ) basin_context = None return EpisodeContext( obs=obs, forecast=fcast, config=cfg, zone_ids=[zone_id], data_source=obs.source, basin_context=basin_context, ) try: import numpy as _np _NUMPY_FOR_DYNAMICS = True except ImportError: _NUMPY_FOR_DYNAMICS = False try: import torch as _torch _TORCH_FOR_DYNAMICS = True except ImportError: _TORCH_FOR_DYNAMICS = False def _forecast_result_to_arrays( forecast_result: Any, horizon_days: int, ) -> Tuple["np.ndarray", float]: import numpy as np precip_seq = list(forecast_result.precip_mm) if forecast_result.precip_mm else [] if len(precip_seq) < horizon_days: precip_seq = precip_seq + [0.0] * (horizon_days - len(precip_seq)) precip_arr = np.clip( np.array(precip_seq[:horizon_days], dtype=np.float32), 0.0, 500.0, ) uncertainty = 0.5 if forecast_result.precip_p90 and forecast_result.precip_p10: p90 = np.array(list(forecast_result.precip_p90)[:horizon_days], dtype=np.float32) p10 = np.array(list(forecast_result.precip_p10)[:horizon_days], dtype=np.float32) spread = np.clip( (p90 - p10) / np.maximum(np.abs(p90), 1e-6), 0.0, 1.0, ) uncertainty = float(np.mean(spread)) return precip_arr, uncertainty def _episode_context_to_state_tensor( contexts: List[Any], horizon_days: int, prior_belief: float = 0.5, ) -> "ZoneStateTensor": import numpy as np from physics_dynamics import ZoneStateTensor n_zones = len(contexts) precip_arr = np.zeros((1, n_zones, horizon_days), dtype=np.float32) uncert_arr = np.zeros((1, n_zones), dtype=np.float32) belief_arr = np.zeros((1, n_zones), dtype=np.float32) for zi, ctx in enumerate(contexts): p_arr, unc = _forecast_result_to_arrays(ctx.forecast, horizon_days) precip_arr[0, zi, :] = p_arr uncert_arr[0, zi] = unc try: signal = float(ctx.obs.composite_risk()) belief = 0.7 * prior_belief + 0.3 * signal except Exception: belief = prior_belief belief_arr[0, zi] = float(np.clip(belief, 0.0, 1.0)) import torch return ZoneStateTensor( precip=torch.from_numpy(precip_arr), uncertainty=torch.from_numpy(uncert_arr), belief=torch.from_numpy(belief_arr), ) def _build_synthetic_sequence( zone_ids: List[str], start_date: "datetime", n_days: int, horizon_days: int, config: Optional["ForecastConfig"] = None, ) -> List["ZoneStateTensor"]: from zone_observation import make_synthetic_episode_context, _stable_seed from physics_dynamics import ZoneStateTensor import numpy as np import torch cfg = config or ForecastConfig() sequence: List[ZoneStateTensor] = [] for day_offset in range(n_days): current_date = start_date + timedelta(days=day_offset) contexts = [] for zi, zid in enumerate(zone_ids): seed = _stable_seed( zid + current_date.isoformat() + str(day_offset) ) ctx = make_synthetic_episode_context(zone_id=zid, seed=seed) contexts.append(ctx) state = _episode_context_to_state_tensor( contexts, horizon_days, prior_belief=cfg.prior_belief ) sequence.append(state) return sequence def _build_era5_sequence( zone_ids: List[str], start_date: "datetime", n_days: int, horizon_days: int, config: Optional["ForecastConfig"] = None, ) -> List["ZoneStateTensor"]: cfg = config or ForecastConfig() sequence = [] for day_offset in range(n_days): current_date = _ensure_utc(start_date + timedelta(days=day_offset)) date_range = (current_date, current_date + timedelta(days=horizon_days)) contexts = [] for zid in zone_ids: try: ctx = fetch_episode_context(zid, date_range, cfg) except Exception as e: logger.warning( "_build_era5_sequence: failed for zone=%s date=%s: %s — using synthetic", zid, current_date.date().isoformat(), e, ) from zone_observation import make_synthetic_episode_context, _stable_seed seed = _stable_seed(zid + current_date.isoformat()) ctx = make_synthetic_episode_context(zone_id=zid, seed=seed) contexts.append(ctx) state = _episode_context_to_state_tensor( contexts, horizon_days, prior_belief=cfg.prior_belief ) sequence.append(state) return sequence def get_consecutive_pairs( zone_ids: List[str], start_date: "datetime", n_days: int = 365, horizon_days: int = 14, config: Optional["ForecastConfig"] = None, use_real_data: bool = True, synthetic_fallback: bool = True, ) -> List[Tuple["ZoneStateTensor", "ZoneStateTensor"]]: if not _NUMPY_FOR_DYNAMICS: raise ImportError("numpy required for get_consecutive_pairs(). pip install numpy") if not _TORCH_FOR_DYNAMICS: raise ImportError("torch required for get_consecutive_pairs(). pip install torch") if n_days < 2: raise ValueError(f"n_days must be >= 2 to produce at least one pair, got {n_days}") start_date = _ensure_utc(start_date) logger.info( "get_consecutive_pairs: zones=%s start=%s n_days=%d horizon=%d real=%s", zone_ids, start_date.date().isoformat(), n_days, horizon_days, use_real_data, ) if use_real_data: try: sequence = _build_era5_sequence( zone_ids, start_date, n_days, horizon_days, config ) except Exception as e: if not synthetic_fallback: raise RuntimeError( f"ERA5 sequence build failed and synthetic_fallback=False: {e}" ) from e logger.warning( "ERA5 sequence build failed (%s) — falling back to full synthetic sequence", e, ) sequence = _build_synthetic_sequence( zone_ids, start_date, n_days, horizon_days, config ) else: sequence = _build_synthetic_sequence( zone_ids, start_date, n_days, horizon_days, config ) pairs = [ (sequence[i], sequence[i + 1]) for i in range(len(sequence) - 1) ] logger.info( "get_consecutive_pairs: built %d pairs from %d-day sequence", len(pairs), n_days, ) return pairs def get_consecutive_pairs_multi_year( zone_ids: List[str], years: List[int], horizon_days: int = 14, config: Optional["ForecastConfig"] = None, use_real_data: bool = True, skip_on_failure: bool = True, ) -> List[Tuple["ZoneStateTensor", "ZoneStateTensor"]]: all_pairs: List[Tuple["ZoneStateTensor", "ZoneStateTensor"]] = [] for year in years: start = _ensure_utc(datetime(year, 1, 1, tzinfo=timezone.utc)) n_days = 366 if (year % 4 == 0 and (year % 100 != 0 or year % 400 == 0)) else 365 try: year_pairs = get_consecutive_pairs( zone_ids=zone_ids, start_date=start, n_days=n_days, horizon_days=horizon_days, config=config, use_real_data=use_real_data, synthetic_fallback=True, ) all_pairs.extend(year_pairs) logger.info("Year %d: added %d pairs (total=%d)", year, len(year_pairs), len(all_pairs)) except Exception as e: if not skip_on_failure: raise logger.warning("Year %d failed (%s) — skipped", year, e) if not all_pairs: raise RuntimeError( f"No pairs collected across years {years}. " "Check zone registration and data availability." ) logger.info( "get_consecutive_pairs_multi_year: %d total pairs from %d years", len(all_pairs), len(years), ) return all_pairs def compute_dataset_statistics( pairs: List[Tuple["ZoneStateTensor", "ZoneStateTensor"]], ) -> Dict[str, Any]: if not pairs: raise ValueError("pairs is empty") import numpy as np import torch all_precip = [] all_uncert = [] all_belief = [] for curr, nxt in pairs: for state in (curr, nxt): all_precip.append(state.precip.numpy().flatten()) all_uncert.append(state.uncertainty.numpy().flatten()) all_belief.append(state.belief.numpy().flatten()) precip_all = np.concatenate(all_precip) uncert_all = np.concatenate(all_uncert) belief_all = np.concatenate(all_belief) stats = { "precip_mean": float(np.mean(precip_all)), "precip_std": float(np.std(precip_all)) + 1e-8, "uncert_mean": float(np.mean(uncert_all)), "uncert_std": float(np.std(uncert_all)) + 1e-8, "belief_mean": float(np.mean(belief_all)), "belief_std": float(np.std(belief_all)) + 1e-8, "n_pairs": len(pairs), "n_zones": pairs[0][0].n_zones, "horizon_days": pairs[0][0].horizon_days, "precip_p95": float(np.percentile(precip_all, 95)), "precip_max": float(np.max(precip_all)), } logger.info( "Dataset stats: n_pairs=%d precip_mean=%.1f±%.1f mm " "uncert_mean=%.3f belief_mean=%.3f", stats["n_pairs"], stats["precip_mean"], stats["precip_std"], stats["uncert_mean"], stats["belief_mean"], ) return stats