""" timesfm_wrapper.py ============================= """ from __future__ import annotations import hashlib import logging import math from dataclasses import dataclass, field from datetime import datetime, timezone from pathlib import Path from typing import List, Optional, Protocol, Tuple, runtime_checkable import numpy as np import zone_observation as _zo assert _zo.SCHEMA_VERSION == 3, ( f"timesfm_wrapper: zone_observation schema mismatch " f"(expected 3, got {_zo.SCHEMA_VERSION})" ) from zone_observation import ( DataSource, ForecastResult, ZoneObs, _stable_seed, ) logger = logging.getLogger(__name__) try: import requests as _requests REQUESTS_AVAILABLE = True except ImportError: REQUESTS_AVAILABLE = False logger.info("requests not installed — OpenMeteoBackend unavailable") # --------------------------------------------------------------------------- # Protocol # --------------------------------------------------------------------------- @runtime_checkable class ForecastBackend(Protocol): def forecast(self, obs: ZoneObs) -> ForecastResult: ... # --------------------------------------------------------------------------- # Shared helpers # --------------------------------------------------------------------------- def _clip_prob(values: List[float]) -> Tuple[float, ...]: return tuple(max(0.0, min(1.0, v)) for v in values) def _make_temp_sequences( temp_mean_c: float, horizon_days: int, spread_c: float = 2.0, ) -> Tuple[Tuple[float, ...], Tuple[float, ...], Tuple[float, ...]]: mean_seq = tuple(float(temp_mean_c) for _ in range(horizon_days)) p10_seq = tuple(float(temp_mean_c - spread_c) for _ in range(horizon_days)) p90_seq = tuple(float(temp_mean_c + spread_c) for _ in range(horizon_days)) return mean_seq, p10_seq, p90_seq def _rh_to_prob_humidity(rh_max_pct: float) -> float: x = (rh_max_pct - 80.0) / 10.0 return float(max(0.0, min(1.0, 1.0 / (1.0 + math.exp(-x))))) # --------------------------------------------------------------------------- # Tier 1 — BaselineBackend # --------------------------------------------------------------------------- class BaselineBackend: """ Deterministic statistical forecast derived entirely from ZoneObs fields. No randomness — deterministic forecasts let the policy learn stable input-output mappings. Stochasticity enters through the observation pipeline, not the forecast. """ def __init__(self, horizon_days: int = 30) -> None: self.horizon_days = horizon_days def forecast(self, obs: ZoneObs) -> ForecastResult: h = self.horizon_days daily_avg = obs.precip_30d_mm / max(30.0, 1.0) decay = np.linspace(1.0, 0.3, h) precip = np.clip(daily_avg * decay, 0.0, 500.0) flood_m = 1.0 + 0.8 * obs.flood_signal() drought_m = 1.0 - 0.6 * obs.drought_signal() precip = np.clip(precip * flood_m * drought_m, 0.0, 500.0) lo_factor = np.linspace(0.75, 0.50, h) hi_factor = np.linspace(1.25, 1.50, h) p10 = np.clip(precip * lo_factor, 0.0, 500.0) p90 = np.clip(precip * hi_factor, 0.0, 500.0) prob_heavy = [float(min(1.0, v / 20.0)) for v in precip] prob_drought = [float(max(0.0, 1.0 - v / 2.0)) for v in precip] prob_humid = [_rh_to_prob_humidity(obs.rh_max_pct)] * h temp_mean, temp_p10, temp_p90 = _make_temp_sequences(obs.temp_mean_c, h) rh_seq = tuple(float(obs.rh_mean_pct) for _ in range(h)) return ForecastResult( zone_id=obs.zone_id, forecast_time=obs.valid_time, horizon_days=h, precip_mm=tuple(float(v) for v in precip), precip_p10=tuple(float(v) for v in p10), precip_p90=tuple(float(v) for v in p90), temp_mean_c=temp_mean, temp_p10=temp_p10, temp_p90=temp_p90, rh_mean_pct=rh_seq, prob_heavy_rain=_clip_prob(prob_heavy), prob_drought_day=_clip_prob(prob_drought), prob_high_humidity=_clip_prob(prob_humid), model_id="baseline-persistence-v1", source=DataSource.SYNTHETIC, ) # --------------------------------------------------------------------------- # Tier 2 — OpenMeteoBackend # --------------------------------------------------------------------------- class OpenMeteoBackend: """ Real probabilistic forecast via Open-Meteo free API. FIX vs original: lat/lon moved from forecast() into __init__() so the object satisfies the ForecastBackend Protocol (forecast(self, obs) only). Resolve lat/lon from the zone registry before constructing this object. Example: lat, lon = zone_registry[zone_id].centroid backend = OpenMeteoBackend(lat=lat, lon=lon, horizon_days=16) result = backend.forecast(obs) """ BASE_URL = "https://api.open-meteo.com/v1/forecast" TIMEOUT_S = 15 def __init__(self, lat: float, lon: float, horizon_days: int = 16) -> None: self.lat = lat self.lon = lon self.horizon_days = min(horizon_days, 16) if horizon_days > 16: logger.warning( "OpenMeteoBackend: horizon_days=%d capped at 16 " "(Open-Meteo free tier limit).", horizon_days ) def forecast(self, obs: ZoneObs) -> ForecastResult: if not REQUESTS_AVAILABLE: raise RuntimeError( "requests not installed — cannot use OpenMeteoBackend. " "Install with: pip install requests" ) h = self.horizon_days params = { "latitude": self.lat, "longitude": self.lon, "daily": ",".join([ "precipitation_sum", "temperature_2m_mean", "temperature_2m_max", "temperature_2m_min", "relative_humidity_2m_mean", "relative_humidity_2m_max", "precipitation_probability_mean", "et0_fao_evapotranspiration", ]), "forecast_days": h, "timezone": "UTC", } try: resp = _requests.get(self.BASE_URL, params=params, timeout=self.TIMEOUT_S) resp.raise_for_status() data = resp.json() except Exception as e: logger.warning( "OpenMeteoBackend: API call failed (%s) — falling back to BaselineBackend", e ) return BaselineBackend(horizon_days=h).forecast(obs) daily = data.get("daily", {}) def _safe(key: str, n: int, default: float = 0.0) -> List[float]: vals = daily.get(key, []) return [ float(vals[i]) if i < len(vals) and vals[i] is not None else default for i in range(n) ] precip_raw = _safe("precipitation_sum", h) temp_mean = _safe("temperature_2m_mean", h, obs.temp_mean_c) temp_max = _safe("temperature_2m_max", h, obs.temp_max_c) temp_min = _safe("temperature_2m_min", h, obs.temp_min_c) rh_mean = _safe("relative_humidity_2m_mean", h, obs.rh_mean_pct) rh_max_raw = _safe("relative_humidity_2m_max", h, obs.rh_max_pct) precip_prob = _safe("precipitation_probability_mean", h, 50.0) precip = [max(0.0, min(500.0, v)) for v in precip_raw] t_p10 = [min(mn, mx) for mn, mx in zip(temp_min, temp_mean)] t_p90 = [max(mx, mn) for mx, mn in zip(temp_max, temp_mean)] obs_spread = float(np.clip( (obs.precip_30d_mm - obs.precip_7d_mm * 4.0) / max(obs.precip_30d_mm, 1.0), 0.1, 0.6, )) unc = [obs_spread * (1.0 + 0.5 * i / h) for i in range(h)] p10 = [max(0.0, p * (1.0 - u * 0.8)) for p, u in zip(precip, unc)] p90 = [min(500.0, p * (1.0 + u)) for p, u in zip(precip, unc)] prob_heavy = [max(0.0, min(1.0, pp / 100.0)) for pp in precip_prob] prob_drought = [max(0.0, min(1.0, 1.0 - pp / 100.0)) for pp in precip_prob] prob_humid = [_rh_to_prob_humidity(rh) for rh in rh_max_raw] return ForecastResult( zone_id=obs.zone_id, forecast_time=obs.valid_time, horizon_days=h, precip_mm=tuple(precip), precip_p10=tuple(p10), precip_p90=tuple(p90), temp_mean_c=tuple(temp_mean), temp_p10=tuple(t_p10), temp_p90=tuple(t_p90), rh_mean_pct=tuple(rh_mean), prob_heavy_rain=_clip_prob(prob_heavy), prob_drought_day=_clip_prob(prob_drought), prob_high_humidity=_clip_prob(prob_humid), model_id="open-meteo-forecast-v1", source=DataSource.OPENMETEO_LIVE, ) # --------------------------------------------------------------------------- # Tier 3 — LocalTimesFMBackend # --------------------------------------------------------------------------- @dataclass class LocalTimesFMBackend: """ Local TimesFM-2.5-200m inference. SHA256-verified. No synthetic fallback. Download checkpoint from: https://huggingface.co/google/timesfm-2.5-200m-pytorch """ checkpoint_path: Path expected_sha256: str device: str = "cpu" horizon_days: int = 30 context_len: int = 512 _model: Optional[object] = field(default=None, init=False, repr=False) def __post_init__(self) -> None: self.checkpoint_path = Path(self.checkpoint_path) if not self.checkpoint_path.exists(): raise FileNotFoundError( f"TimesFM checkpoint not found: {self.checkpoint_path}" ) logger.info("Verifying TimesFM checkpoint SHA256...") computed = self._sha256_file(self.checkpoint_path) if computed != self.expected_sha256.lower().strip(): raise RuntimeError( f"TimesFM checkpoint SHA256 MISMATCH.\n" f"Expected: {self.expected_sha256}\nGot: {computed}" ) logger.info("TimesFM checkpoint verified OK") self._model = self._load_model() def _load_model(self) -> object: try: import timesfm # type: ignore except ImportError: raise ImportError("timesfm not installed. pip install timesfm") model = timesfm.TimesFm( hparams=timesfm.TimesFmHparams( backend=self.device, per_core_batch_size=32, horizon_len=self.horizon_days, context_len=self.context_len, ), checkpoint=timesfm.TimesFmCheckpoint( local_path=str(self.checkpoint_path) ), ) logger.info("TimesFM loaded: device=%s horizon=%d", self.device, self.horizon_days) return model def forecast(self, obs: ZoneObs) -> ForecastResult: if self._model is None: raise RuntimeError("TimesFM model not loaded.") import pandas as pd # type: ignore import datetime as _dt h = self.horizon_days daily_30d = obs.precip_30d_mm / 30.0 daily_7d = obs.precip_7d_mm / 7.0 daily_24h = obs.precip_24h_mm context = np.concatenate([ np.full(23, daily_30d), np.full(6, daily_7d), np.array([daily_24h]), ]).astype(np.float32) base_date = obs.valid_time.date() dates = [ base_date - _dt.timedelta(days=len(context) - i - 1) for i in range(len(context)) ] df_input = pd.DataFrame({ "unique_id": [obs.zone_id] * len(context), "ds": dates, "y": context, }) forecast_df, _ = self._model.forecast_on_df( inputs=df_input, freq="D", value_name="y", num_jobs=1, ) p50_col = "timesfm-q-0.5" if "timesfm-q-0.5" in forecast_df.columns else "timesfm" p10_col = "timesfm-q-0.1" if "timesfm-q-0.1" in forecast_df.columns else p50_col p90_col = "timesfm-q-0.9" if "timesfm-q-0.9" in forecast_df.columns else p50_col p50 = np.clip(forecast_df[p50_col].values[:h].astype(float), 0.0, 500.0) p10 = np.clip(forecast_df[p10_col].values[:h].astype(float), 0.0, 500.0) p90 = np.clip(forecast_df[p90_col].values[:h].astype(float), 0.0, 500.0) p10 = np.minimum(p10, p50) p90 = np.maximum(p90, p50) prob_heavy = [float(min(1.0, v / 20.0)) for v in p50] prob_drought = [float(max(0.0, 1.0 - v / 2.0)) for v in p50] prob_humid = [_rh_to_prob_humidity(obs.rh_max_pct)] * h temp_mean_seq, temp_p10_seq, temp_p90_seq = _make_temp_sequences( obs.temp_mean_c, h ) return ForecastResult( zone_id=obs.zone_id, forecast_time=obs.valid_time, horizon_days=h, precip_mm=tuple(float(v) for v in p50), precip_p10=tuple(float(v) for v in p10), precip_p90=tuple(float(v) for v in p90), temp_mean_c=temp_mean_seq, temp_p10=temp_p10_seq, temp_p90=temp_p90_seq, rh_mean_pct=tuple(float(obs.rh_mean_pct) for _ in range(h)), prob_heavy_rain=_clip_prob(prob_heavy), prob_drought_day=_clip_prob(prob_drought), prob_high_humidity=_clip_prob(prob_humid), model_id="timesfm-2.5-200m", source=DataSource.SYNTHETIC, # FIX: was ERA5_REANALYSIS (incorrect) ) @staticmethod def _sha256_file(path: Path, chunk_size: int = 1 << 20) -> str: h = hashlib.sha256() with open(path, "rb") as f: while chunk := f.read(chunk_size): h.update(chunk) return h.hexdigest() # --------------------------------------------------------------------------- # Factory # --------------------------------------------------------------------------- def create_forecast_backend( mode: str = "baseline", checkpoint_path: Optional[Path] = None, expected_sha256: Optional[str] = None, device: str = "cpu", horizon_days: int = 30, lat: Optional[float] = None, lon: Optional[float] = None, ) -> ForecastBackend: """ Create a ForecastBackend. mode="baseline" — deterministic statistical forecast (default, no deps) mode="openmeteo" — real Open-Meteo API forecast (requires requests, lat, lon) mode="timesfm" — local TimesFM inference (requires checkpoint + timesfm) """ if mode == "baseline": return BaselineBackend(horizon_days=horizon_days) if mode == "openmeteo": if not REQUESTS_AVAILABLE: logger.warning( "requests not installed — falling back to BaselineBackend." ) return BaselineBackend(horizon_days=horizon_days) if lat is None or lon is None: raise ValueError( "mode='openmeteo' requires lat and lon. " "Resolve from zone registry: lat, lon = registry[zone_id].centroid" ) return OpenMeteoBackend(lat=lat, lon=lon, horizon_days=horizon_days) if mode == "timesfm": if checkpoint_path is None or expected_sha256 is None: raise ValueError( "mode='timesfm' requires both checkpoint_path and expected_sha256." ) return LocalTimesFMBackend( checkpoint_path=Path(checkpoint_path), expected_sha256=expected_sha256, device=device, horizon_days=horizon_days, ) raise ValueError( f"Unknown backend mode: '{mode}'. Valid: 'baseline', 'openmeteo', 'timesfm'" )