Spaces:
Sleeping
Sleeping
| """ | |
| ERA5 meteorological reanalysis data loader (SAFRAN alternative). | |
| Single Responsibility: Load and parse ERA5 data with spatial interpolation. | |
| """ | |
| import pandas as pd | |
| import xarray as xr | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import matplotlib.dates as mdates | |
| from pathlib import Path | |
| from typing import Optional, List | |
| from .base import BaseDataLoader | |
| class SAFRANLoader(BaseDataLoader): | |
| """ | |
| Loads ERA5 reanalysis data (used as SAFRAN alternative). | |
| Handles actual ERA5 NetCDF structure: | |
| - Gridded at 0.25° resolution | |
| - Variables: t2m, tp, pev, ssrd, u10, v10, sf, ro | |
| - Interpolates to station points | |
| """ | |
| # Variable mapping | |
| VAR_MAP = { | |
| 't2m': 'temp_2m_K', | |
| 'tp': 'precip_m', | |
| 'pev': 'evap_m', | |
| 'ssrd': 'solar_Jm2', | |
| 'u10': 'wind_u_ms', | |
| 'v10': 'wind_v_ms', | |
| 'sf': 'snow_m', | |
| 'ro': 'runoff_m' | |
| } | |
| def __init__( | |
| self, | |
| data_path: Path, | |
| station_coords: Optional[pd.DataFrame] = None, | |
| interp_method: str = "nearest" | |
| ): | |
| """ | |
| Initialize ERA5 loader. | |
| Args: | |
| data_path: Path to datasets/safran directory | |
| station_coords: DataFrame with [station_code, lat, lon] | |
| interp_method: "nearest" or "linear" | |
| """ | |
| super().__init__(data_path) | |
| self.station_coords = station_coords | |
| self.interp_method = interp_method | |
| def load(self) -> pd.DataFrame: | |
| """ | |
| Load ERA5 data and interpolate to stations. | |
| Returns: | |
| DataFrame with [date, station_code, temp_2m_K, precip_m, ...] | |
| """ | |
| nc_files = sorted(Path(self.data_path).glob("era5_*.nc")) | |
| if not nc_files: | |
| raise FileNotFoundError(f"No ERA5 files in {self.data_path}") | |
| all_dfs = [] | |
| for nc_file in nc_files: | |
| df = self._load_and_interpolate(nc_file) | |
| all_dfs.append(df) | |
| combined = pd.concat(all_dfs, ignore_index=True) | |
| return combined.sort_values(['date', 'station_code']).reset_index(drop=True) | |
| def load_grid(self, bbox: Optional[tuple] = None) -> pd.DataFrame: | |
| """ | |
| Load ERA5 at its NATIVE grid resolution (0.25 degrees) -- no | |
| per-station interpolation at all, unlike load(). Confirmed | |
| against real data that station-level interpolation duplicates | |
| signal: multiple real stations landed on the exact identical | |
| grid cell (several stations sharing bit-identical wind_speed_ms | |
| values across the whole real climatology test), since ERA5's | |
| 0.25 deg cells are coarse relative to how tightly some real | |
| gauges cluster. Modeling at grid-cell resolution avoids this | |
| redundancy directly, rather than training separate models on | |
| what's provably the same input. | |
| Args: | |
| bbox: optional (min_lon, min_lat, max_lon, max_lat) to | |
| subset the grid before loading -- the whole point is to | |
| avoid loading the full national/continental ERA5 extent | |
| when only a couple dozen cells actually cover the study | |
| area. Defaults to no subsetting (the full extent | |
| already present in the downloaded files, which | |
| download_era5_full.py already bounded at fetch time). | |
| Returns: | |
| DataFrame with [date, grid_lat, grid_lon, temp_2m_K, | |
| precip_m, ...] -- one row per (date, cell), long format | |
| (not yet pivoted to wide per-variable, since the number of | |
| variables and their names come from VAR_MAP and shouldn't | |
| be hardcoded here). | |
| """ | |
| nc_files = sorted(Path(self.data_path).glob("era5_*.nc")) | |
| if not nc_files: | |
| raise FileNotFoundError(f"No ERA5 files in {self.data_path}") | |
| all_dfs = [] | |
| for nc_file in nc_files: | |
| df = self._load_grid_single_file(nc_file, bbox) | |
| all_dfs.append(df) | |
| combined = pd.concat(all_dfs, ignore_index=True) | |
| return combined.sort_values(['date', 'grid_lat', 'grid_lon']).reset_index(drop=True) | |
| def _load_grid_single_file(self, nc_file: Path, bbox: Optional[tuple]) -> pd.DataFrame: | |
| """Load a single NetCDF file's full grid (optionally bbox-subset), | |
| no station interpolation. Every (lat, lon, date) triple becomes | |
| one row -- this is real gridded data, not point samples.""" | |
| ds = xr.open_dataset(nc_file) | |
| time_coord = 'valid_time' if 'valid_time' in ds.coords else 'time' | |
| if bbox is not None: | |
| min_lon, min_lat, max_lon, max_lat = bbox | |
| # ERA5 latitude is typically stored descending (90 to -90); | |
| # slice() requires matching that order or it silently | |
| # returns nothing, which is exactly the kind of thing that | |
| # produced a wrong-axis-order failure earlier in this | |
| # project (BD TOPO's WFS) -- check the actual coordinate | |
| # direction rather than assume one. | |
| lat_vals = ds['latitude'].values | |
| if lat_vals[0] > lat_vals[-1]: | |
| ds = ds.sel(latitude=slice(max_lat, min_lat), longitude=slice(min_lon, max_lon)) | |
| else: | |
| ds = ds.sel(latitude=slice(min_lat, max_lat), longitude=slice(min_lon, max_lon)) | |
| if ds.sizes.get('latitude', 0) == 0 or ds.sizes.get('longitude', 0) == 0: | |
| raise ValueError( | |
| f"bbox {bbox} selected zero grid cells from {nc_file.name} -- check the bbox " | |
| f"is within this file's real coordinate range: " | |
| f"lat {ds['latitude'].values.min() if ds.sizes.get('latitude',1) else 'N/A'}.." | |
| f"{ds['latitude'].values.max() if ds.sizes.get('latitude',1) else 'N/A'}, " | |
| f"lon similarly." | |
| ) | |
| df = ds.to_dataframe().reset_index() | |
| df = df.rename(columns={time_coord: 'date', 'latitude': 'grid_lat', 'longitude': 'grid_lon', **self.VAR_MAP}) | |
| keep_cols = ['date', 'grid_lat', 'grid_lon'] + [v for v in self.VAR_MAP.values() if v in df.columns] | |
| df = df[keep_cols].copy() | |
| df['date'] = pd.to_datetime(df['date']).dt.tz_localize(None) | |
| ds.close() | |
| return df | |
| def _load_and_interpolate(self, nc_file: Path) -> pd.DataFrame: | |
| """ | |
| Load single NetCDF and interpolate to ALL stations in one | |
| vectorized xarray selection, not one .sel() + .to_dataframe() | |
| call per station. | |
| The original per-station loop does real work with meaningful | |
| per-call overhead (an index lookup, then a full DataFrame | |
| conversion) once per station PER FILE. At 27 stations across | |
| ~67 year-files (1960-2026) that's ~1,800 calls -- slow but | |
| tolerable. At the reach graph's ~2,900 nodes it's ~193,000 | |
| individual calls, which is what was actually hanging, not a | |
| proportionally-worse runtime. | |
| Fix: xarray's vectorized ("pointwise") indexing -- when the | |
| latitude/longitude indexers are DataArrays sharing a common | |
| dimension name ("station" here), a single .sel() call looks up | |
| every point at once instead of one point per call, producing a | |
| result with a "station" dimension instead of separate | |
| latitude/longitude dimensions. This turns ~193,000 calls into | |
| ~67 (one per file). | |
| Reasoned from documented xarray API behavior, not verified | |
| against a real NetCDF file with real execution -- test against | |
| a small subset of the era5_*.nc files first, and confirm the | |
| output matches the old per-station-loop version's shape/values | |
| before trusting it for real training data. | |
| """ | |
| ds = xr.open_dataset(nc_file) | |
| time_coord = 'valid_time' if 'valid_time' in ds.coords else 'time' | |
| station_lats = xr.DataArray(self.station_coords['lat'].values, dims='station') | |
| station_lons = xr.DataArray(self.station_coords['lon'].values, dims='station') | |
| station_codes = self.station_coords['station_code'].values | |
| point_ds = ds.sel(latitude=station_lats, longitude=station_lons, method=self.interp_method) | |
| # point_ds now has a 'station' dimension (length n_stations) instead | |
| # of separate latitude/longitude dimensions -- attach real station | |
| # codes as a coordinate so they survive the to_dataframe() below, | |
| # rather than staying as anonymous integer positions. | |
| point_ds = point_ds.assign_coords(station=station_codes) | |
| df = point_ds.to_dataframe().reset_index() | |
| df = df.rename(columns={time_coord: 'date', 'station': 'station_code', **self.VAR_MAP}) | |
| keep_cols = ['date', 'station_code'] + [v for v in self.VAR_MAP.values() if v in df.columns] | |
| df = df[keep_cols].copy() | |
| df['date'] = pd.to_datetime(df['date']).dt.tz_localize(None) | |
| ds.close() | |
| return df | |
| def convert_units(df: pd.DataFrame) -> pd.DataFrame: | |
| """Convert ERA5 units to standard (°C, mm, W/m²).""" | |
| df = df.copy() | |
| # K → °C | |
| if 'temp_2m_K' in df.columns: | |
| df['temp_C'] = df['temp_2m_K'] - 273.15 | |
| # m → mm | |
| for var in ['precip_m', 'evap_m', 'snow_m', 'runoff_m']: | |
| if var in df.columns: | |
| df[var.replace('_m', '_mm')] = df[var] * 1000 | |
| # J/m² → W/m² | |
| if 'solar_Jm2' in df.columns: | |
| df['solar_Wm2'] = df['solar_Jm2'] / 86400 | |
| # Wind speed | |
| if 'wind_u_ms' in df.columns and 'wind_v_ms' in df.columns: | |
| df['wind_speed_ms'] = np.sqrt(df['wind_u_ms']**2 + df['wind_v_ms']**2) | |
| return df | |
| def get_metadata(self) -> dict: | |
| """Get ERA5 metadata.""" | |
| meta = super().get_metadata() | |
| meta.update({ | |
| "data_type": "era5_reanalysis", | |
| "source": "Copernicus CDS", | |
| "resolution": "0.25 degrees", | |
| "interp_method": self.interp_method | |
| }) | |
| return meta | |
| def plot_temperature( | |
| self, | |
| df: Optional[pd.DataFrame] = None, | |
| stations: Optional[List[str]] = None, | |
| figsize: tuple = (12, 5), | |
| save_path: Optional[Path] = None, | |
| ) -> plt.Axes: | |
| """ | |
| Plot 2m air temperature time series, one line per station. | |
| Args: | |
| df: Optional pre-loaded DataFrame (raw or already converted). If not | |
| provided, data is loaded and unit-converted from disk. | |
| stations: Optional list of station codes to plot. Defaults to all. | |
| figsize: Figure size in inches. | |
| save_path: If provided, saves the figure to this path. | |
| Returns: | |
| The matplotlib Axes object. | |
| """ | |
| if df is None: | |
| df = self.convert_units(self.load()) | |
| elif "temp_C" not in df.columns and "temp_2m_K" in df.columns: | |
| df = self.convert_units(df) | |
| plot_df = df.dropna(subset=["temp_C"]) | |
| if stations: | |
| plot_df = plot_df[plot_df["station_code"].isin(stations)] | |
| if plot_df.empty: | |
| raise ValueError("No temperature data available to plot") | |
| fig, ax = plt.subplots(figsize=figsize) | |
| for station_code, group in plot_df.groupby("station_code"): | |
| group = group.sort_values("date") | |
| ax.plot(group["date"], group["temp_C"], linewidth=1, label=station_code) | |
| ax.set_title("2m Air Temperature") | |
| ax.set_xlabel("Date") | |
| ax.set_ylabel("Temperature (°C)") | |
| ax.xaxis.set_major_locator(mdates.AutoDateLocator()) | |
| ax.xaxis.set_major_formatter(mdates.ConciseDateFormatter(ax.xaxis.get_major_locator())) | |
| ax.legend(title="Station", fontsize=8) | |
| ax.grid(True, alpha=0.3) | |
| plt.tight_layout() | |
| if save_path: | |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") | |
| print(f"Plot saved to {save_path}") | |
| return ax | |
| def plot_precipitation( | |
| self, | |
| df: Optional[pd.DataFrame] = None, | |
| stations: Optional[List[str]] = None, | |
| cumulative: bool = True, | |
| figsize: tuple = (12, 5), | |
| save_path: Optional[Path] = None, | |
| ) -> plt.Axes: | |
| """ | |
| Plot precipitation, one line per station. Cumulative by default, | |
| which makes it easy to compare total rainfall received over the | |
| record and spot where stations diverge. | |
| Args: | |
| df: Optional pre-loaded DataFrame (raw or already converted). If not | |
| provided, data is loaded and unit-converted from disk. | |
| stations: Optional list of station codes to plot. Defaults to all. | |
| cumulative: If True, plot the running cumulative total (mm). | |
| If False, plot the raw daily/timestep values. | |
| figsize: Figure size in inches. | |
| save_path: If provided, saves the figure to this path. | |
| Returns: | |
| The matplotlib Axes object. | |
| """ | |
| if df is None: | |
| df = self.convert_units(self.load()) | |
| elif "precip_mm" not in df.columns and "precip_m" in df.columns: | |
| df = self.convert_units(df) | |
| plot_df = df.dropna(subset=["precip_mm"]) | |
| if stations: | |
| plot_df = plot_df[plot_df["station_code"].isin(stations)] | |
| if plot_df.empty: | |
| raise ValueError("No precipitation data available to plot") | |
| fig, ax = plt.subplots(figsize=figsize) | |
| for station_code, group in plot_df.groupby("station_code"): | |
| group = group.sort_values("date") | |
| values = group["precip_mm"].cumsum() if cumulative else group["precip_mm"] | |
| ax.plot(group["date"], values, linewidth=1, label=station_code) | |
| ax.set_title("Cumulative Precipitation" if cumulative else "Precipitation") | |
| ax.set_xlabel("Date") | |
| ax.set_ylabel("Cumulative precipitation (mm)" if cumulative else "Precipitation (mm)") | |
| ax.xaxis.set_major_locator(mdates.AutoDateLocator()) | |
| ax.xaxis.set_major_formatter(mdates.ConciseDateFormatter(ax.xaxis.get_major_locator())) | |
| ax.legend(title="Station", fontsize=8) | |
| ax.grid(True, alpha=0.3) | |
| plt.tight_layout() | |
| if save_path: | |
| plt.savefig(save_path, dpi=150, bbox_inches="tight") | |
| print(f"Plot saved to {save_path}") | |
| return ax |