from __future__ import annotations import logging import time from dataclasses import dataclass from datetime import datetime, timedelta, timezone from typing import Any, Iterator from zoneinfo import ZoneInfo import httpx import numpy as np import pandas as pd from tsfm_bench.data.interface import DatasetMetadata, TimeSeriesRecord logger = logging.getLogger(__name__) FORECAST_API = "https://api.open-meteo.com/v1/forecast" ARCHIVE_API = "https://archive-api.open-meteo.com/v1/archive" @dataclass(frozen=True) class _LocationSpec: name: str domain: str latitude: float longitude: float variables: tuple[str, ...] frequency: str terms: tuple[str, ...] timezone: str = "UTC" city: str = "" @dataclass(frozen=True) class _OpenMeteoSettings: past_days: int = 14 archive_days: int = 30 min_context_hours: int = 168 prediction_length: int = 24 request_timeout: float = 30.0 http_retries: int = 3 http_retry_delay: float = 2.0 class OpenMeteoDataSource: """Live weather time series from Open-Meteo (https://open-meteo.com/).""" def __init__( self, locations: tuple[_LocationSpec, ...], settings: _OpenMeteoSettings | None = None, ): self._locations = {spec.name: spec for spec in locations} self._settings = settings or _OpenMeteoSettings() @classmethod def from_config(cls, payload: dict[str, Any]) -> OpenMeteoDataSource: settings_raw = payload.get("settings", {}) settings = _OpenMeteoSettings( past_days=int(settings_raw.get("past_days", 14)), archive_days=int(settings_raw.get("archive_days", 30)), min_context_hours=int(settings_raw.get("min_context_hours", 168)), prediction_length=int(settings_raw.get("prediction_length", 24)), request_timeout=float(settings_raw.get("request_timeout", 30.0)), http_retries=int(settings_raw.get("http_retries", 3)), http_retry_delay=float(settings_raw.get("http_retry_delay", 2.0)), ) specs = tuple(_parse_location(entry) for entry in payload.get("locations", [])) if not specs: raise ValueError("open_meteo config requires at least one location") return cls(locations=specs, settings=settings) def list_datasets(self) -> list[str]: return sorted(self._locations.keys()) def get_metadata(self, name: str) -> DatasetMetadata: spec = self._require_spec(name) return DatasetMetadata( name=spec.name, domain=spec.domain, frequency=spec.frequency, num_variates=len(spec.variables), terms=spec.terms, source_type="open_meteo", ) def stream(self, name: str) -> Iterator[TimeSeriesRecord]: spec = self._require_spec(name) frame = self.fetch_location_frame(spec) if frame.empty: raise RuntimeError(f"No Open-Meteo data returned for {name}") values = frame[list(spec.variables)].to_numpy(dtype=np.float32).T mask = ~np.isnan(values).any(axis=0) values = values[:, mask] if values.shape[-1] < self._settings.min_context_hours + self._settings.prediction_length: raise RuntimeError( f"Insufficient history for {name}: got {values.shape[-1]} hours, " f"need at least {self._settings.min_context_hours + self._settings.prediction_length}" ) start_ts = frame.index[mask][0] freq = pd.tseries.frequencies.to_offset(spec.frequency) yield TimeSeriesRecord( item_id=spec.name.replace("/", "_"), target=values, start=str(pd.Period(start_ts, freq=freq)), freq=spec.frequency, ) def fetch_location_frame(self, spec: _LocationSpec) -> pd.DataFrame: """Fetch merged archive + recent forecast observations for one location.""" now = datetime.now(timezone.utc) archive_start = (now - timedelta(days=self._settings.archive_days)).date() archive_end = (now - timedelta(days=1)).date() frames: list[pd.DataFrame] = [] if archive_start <= archive_end: frames.append( self._fetch_hourly( ARCHIVE_API, spec, start_date=str(archive_start), end_date=str(archive_end), ) ) frames.append( self._fetch_hourly( FORECAST_API, spec, past_days=self._settings.past_days, forecast_days=1, ) ) merged = pd.concat(frames).sort_index() merged = merged[~merged.index.duplicated(keep="last")] merged = merged.dropna(how="all") return merged def fetched_at(self) -> str: return datetime.now(timezone.utc).isoformat() def _fetch_hourly( self, base_url: str, spec: _LocationSpec, **extra_params: Any, ) -> pd.DataFrame: params: dict[str, Any] = { "latitude": spec.latitude, "longitude": spec.longitude, "hourly": ",".join(spec.variables), "timezone": spec.timezone, **extra_params, } retries = max(1, self._settings.http_retries) delay = self._settings.http_retry_delay last_exc: Exception | None = None for attempt in range(1, retries + 1): try: with httpx.Client(timeout=self._settings.request_timeout, trust_env=False) as client: response = client.get(base_url, params=params) response.raise_for_status() payload = response.json() break except (httpx.TimeoutException, httpx.NetworkError) as exc: last_exc = exc except httpx.HTTPStatusError as exc: if exc.response.status_code < 500: raise last_exc = exc if attempt < retries: wait = delay * (2 ** (attempt - 1)) logger.warning( "Open-Meteo request failed (%s); retrying in %.1fs (attempt %d/%d)", last_exc, wait, attempt, retries, ) time.sleep(wait) else: raise RuntimeError( f"Open-Meteo request to {base_url} failed after {retries} attempts" ) from last_exc hourly = payload.get("hourly", {}) times = hourly.get("time", []) if not times: return pd.DataFrame() tz = ZoneInfo(spec.timezone) index = pd.DatetimeIndex( [datetime.fromisoformat(ts).replace(tzinfo=tz) for ts in times] ) data = {var: hourly.get(var, []) for var in spec.variables} frame = pd.DataFrame(data, index=index) for col in spec.variables: frame[col] = pd.to_numeric(frame[col], errors="coerce") return frame def _require_spec(self, name: str) -> _LocationSpec: if name not in self._locations: raise KeyError(f"Unknown Open-Meteo dataset: {name}") return self._locations[name] def _parse_location(entry: dict[str, Any]) -> _LocationSpec: variables = entry.get("variables", ["temperature_2m"]) if isinstance(variables, str): variables = [v.strip() for v in variables.split(",")] return _LocationSpec( name=entry["name"], domain=entry.get("domain", "Nature"), latitude=float(entry["latitude"]), longitude=float(entry["longitude"]), variables=tuple(variables), frequency=entry.get("frequency", "H"), terms=tuple(entry.get("terms", ["short"])), timezone=entry.get("timezone", "UTC"), city=entry.get("city", entry["name"].split("/")[0]), )