Spaces:
Running
Running
| 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" | |
| class _LocationSpec: | |
| name: str | |
| domain: str | |
| latitude: float | |
| longitude: float | |
| variables: tuple[str, ...] | |
| frequency: str | |
| terms: tuple[str, ...] | |
| timezone: str = "UTC" | |
| city: str = "" | |
| 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() | |
| 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]), | |
| ) | |