Spaces:
Running
Running
File size: 8,027 Bytes
e317359 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 | 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]),
)
|