diff --git a/pyproject.toml b/pyproject.toml index 76d00d9..403b6fd 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -8,7 +8,7 @@ version = "1.0.0" description = "Eurus Climate Agent - Access ERA5 reanalysis data through Model Context Protocol" readme = "README.md" license = {text = "MIT"} -requires-python = ">=3.10" +requires-python = ">=3.11" authors = [ {name = "Eurus Team", email = "eurus@example.com"} ] @@ -30,9 +30,9 @@ classifiers = [ "License :: OSI Approved :: MIT License", "Operating System :: OS Independent", "Programming Language :: Python :: 3", - "Programming Language :: Python :: 3.10", "Programming Language :: Python :: 3.11", "Programming Language :: Python :: 3.12", + "Programming Language :: Python :: 3.13", "Topic :: Scientific/Engineering :: Atmospheric Science", "Topic :: Scientific/Engineering :: GIS", ] @@ -42,6 +42,7 @@ dependencies = [ "arraylake>=0.10.0", "xarray>=2024.10.0", "zarr>=3.0.0", + "numcodecs[pcodec]>=0.16.0", "pandas>=2.0.0", "numpy>=1.24.0", "pydantic>=2.0.0", @@ -106,7 +107,7 @@ include = [ [tool.black] line-length = 100 -target-version = ['py310', 'py311', 'py312'] +target-version = ['py311', 'py312', 'py313'] [tool.ruff] line-length = 100 @@ -125,7 +126,7 @@ ignore = [ ] [tool.mypy] -python_version = "3.10" +python_version = "3.11" warn_return_any = true warn_unused_configs = true disallow_untyped_defs = true diff --git a/src/eurus/config.py b/src/eurus/config.py index 22c0867..d202785 100644 --- a/src/eurus/config.py +++ b/src/eurus/config.py @@ -54,13 +54,16 @@ class ERA5Variable: category: str typical_range: tuple[float | None, float | None] = (None, None) colormap: str = "viridis" + zarr_name: Optional[str] = None # array name in the store, if it differs from short_name def __str__(self) -> str: return f"{self.short_name}: {self.long_name} ({self.units})" -# Comprehensive ERA5 variable mapping — ALL 22 Arraylake variables -# Source: earthmover-public/era5-surface-aws Icechunk store +# Comprehensive ERA5 variable mapping — surface/single-level variables from +# the earthmover-public/era5 Icechunk store (group "single"). +# `zarr_name` is set when the array name in the store differs from our +# catalog `short_name` (kept stable for backwards-compatible outputs/filenames). ERA5_VARIABLES: Dict[str, ERA5Variable] = { # ── Ocean ────────────────────────────────────────────────────────────── "sst": ERA5Variable( @@ -80,7 +83,8 @@ def __str__(self) -> str: description="Air temperature at 2 meters above the surface", category="atmosphere", typical_range=(220, 330), - colormap="RdYlBu_r" + colormap="RdYlBu_r", + zarr_name="t2m" ), "d2": ERA5Variable( short_name="d2", @@ -89,7 +93,8 @@ def __str__(self) -> str: description="Temperature to which air at 2m must cool to reach saturation; indicates humidity", category="atmosphere", typical_range=(220, 310), - colormap="RdYlBu_r" + colormap="RdYlBu_r", + zarr_name="d2m" ), "skt": ERA5Variable( short_name="skt", @@ -155,7 +160,8 @@ def __str__(self) -> str: description="Atmospheric pressure reduced to mean sea level", category="atmosphere", typical_range=(96000, 105000), - colormap="viridis" + colormap="viridis", + zarr_name="msl" ), # ── Boundary Layer ───────────────────────────────────────────────────── "blh": ERA5Variable( @@ -279,6 +285,145 @@ def __str__(self) -> str: typical_range=(0, 0.5), colormap="YlGnBu" ), + "stl2": ERA5Variable( + short_name="stl2", + long_name="Soil Temperature Level 2", + units="K", + description="Temperature of the second soil layer (7-28 cm depth)", + category="land_surface", + typical_range=(220, 330), + colormap="RdYlBu_r" + ), + "stl3": ERA5Variable( + short_name="stl3", + long_name="Soil Temperature Level 3", + units="K", + description="Temperature of the third soil layer (28-100 cm depth)", + category="land_surface", + typical_range=(220, 330), + colormap="RdYlBu_r" + ), + "stl4": ERA5Variable( + short_name="stl4", + long_name="Soil Temperature Level 4", + units="K", + description="Temperature of the deepest soil layer (100-289 cm depth)", + category="land_surface", + typical_range=(220, 330), + colormap="RdYlBu_r" + ), + "fsr": ERA5Variable( + short_name="fsr", + long_name="Forecast Surface Roughness", + units="m", + description="Aerodynamic roughness length of the surface", + category="land_surface", + typical_range=(0, 3), + colormap="YlGnBu" + ), + # ── Wind (additional) ───────────────────────────────────────────────── + "fg10": ERA5Variable( + short_name="fg10", + long_name="10m Wind Gust", + units="m/s", + description="Maximum 3-second wind gust at 10 meters since the previous post-processing", + category="atmosphere", + typical_range=(0, 60), + colormap="viridis" + ), + "zust": ERA5Variable( + short_name="zust", + long_name="Friction Velocity", + units="m/s", + description="Turbulent surface stress expressed as a velocity scale", + category="atmosphere", + typical_range=(0, 3), + colormap="viridis" + ), + # ── Cloud cover (additional) ─────────────────────────────────────────── + "hcc": ERA5Variable( + short_name="hcc", + long_name="High Cloud Cover", + units="fraction (0-1)", + description="Fraction of sky covered by high-altitude clouds", + category="atmosphere", + typical_range=(0, 1), + colormap="gray_r" + ), + "mcc": ERA5Variable( + short_name="mcc", + long_name="Medium Cloud Cover", + units="fraction (0-1)", + description="Fraction of sky covered by mid-altitude clouds", + category="atmosphere", + typical_range=(0, 1), + colormap="gray_r" + ), + "lcc": ERA5Variable( + short_name="lcc", + long_name="Low Cloud Cover", + units="fraction (0-1)", + description="Fraction of sky covered by low-altitude clouds", + category="atmosphere", + typical_range=(0, 1), + colormap="gray_r" + ), + # ── Precipitation (additional) ───────────────────────────────────────── + "sf": ERA5Variable( + short_name="sf", + long_name="Snowfall", + units="m water equiv.", + description="Accumulated snowfall expressed as meters of water equivalent", + category="precipitation", + typical_range=(0, 0.05), + colormap="Blues" + ), + # ── Radiation & Heat Flux (additional) ───────────────────────────────── + "fdir": ERA5Variable( + short_name="fdir", + long_name="Total Sky Direct Solar Radiation", + units="J/m²", + description="Direct (unscattered) shortwave radiation reaching the surface", + category="radiation", + typical_range=(0, 3.5e7), + colormap="YlOrRd" + ), + "tisr": ERA5Variable( + short_name="tisr", + long_name="TOA Incident Solar Radiation", + units="J/m²", + description="Solar radiation incident at the top of the atmosphere", + category="radiation", + typical_range=(0, 4e7), + colormap="YlOrRd" + ), + "tsr": ERA5Variable( + short_name="tsr", + long_name="Top Net Solar Radiation", + units="J/m²", + description="Net shortwave radiation balance at the top of the atmosphere", + category="radiation", + typical_range=(0, 4e7), + colormap="YlOrRd" + ), + "slhf": ERA5Variable( + short_name="slhf", + long_name="Surface Latent Heat Flux", + units="J/m²", + description="Energy transferred between the surface and atmosphere through evaporation/condensation", + category="radiation", + typical_range=(-3e6, 1e6), + colormap="RdBu_r" + ), + "ie": ERA5Variable( + short_name="ie", + long_name="Instantaneous Moisture Flux", + units="kg m⁻² s⁻¹", + description="Instantaneous surface evaporation/condensation flux", + category="atmosphere", + typical_range=(-0.002, 0.002), + colormap="RdBu_r" + ), } # Aliases for long variable names → short names @@ -321,6 +466,26 @@ def __str__(self) -> str: "soil_temperature_level_1": "stl1", "soil_moisture": "swvl1", "volumetric_soil_water_layer_1": "swvl1", + "soil_temperature_level_2": "stl2", + "soil_temperature_level_3": "stl3", + "soil_temperature_level_4": "stl4", + "forecast_surface_roughness": "fsr", + # Wind (additional) + "10m_wind_gust_since_previous_post_processing": "fg10", + "wind_gust": "fg10", + "friction_velocity": "zust", + # Cloud cover (additional) + "high_cloud_cover": "hcc", + "medium_cloud_cover": "mcc", + "low_cloud_cover": "lcc", + # Precipitation (additional) + "snowfall": "sf", + # Radiation & heat flux (additional) + "total_sky_direct_solar_radiation_at_surface": "fdir", + "toa_incident_solar_radiation": "tisr", + "top_net_solar_radiation": "tsr", + "surface_latent_heat_flux": "slhf", + "instantaneous_moisture_flux": "ie", } @@ -334,7 +499,7 @@ def get_variable_info(variable_id: str) -> Optional[ERA5Variable]: def get_short_name(variable_id: str) -> str: - """Get the short name for a variable (for dataset access).""" + """Get the catalog short name for a variable (stable across dataset versions).""" key = variable_id.lower() # Check aliases first if key in VARIABLE_ALIASES: @@ -345,6 +510,18 @@ def get_short_name(variable_id: str) -> str: return key +def get_zarr_name(variable_id: str) -> str: + """Get the array name to use when indexing into the Zarr store. + + Usually equal to the catalog short name, except where the store's array + name differs (e.g. catalog "t2" maps to the store's "t2m"). + """ + var_info = get_variable_info(variable_id) + if var_info: + return var_info.zarr_name or var_info.short_name + return get_short_name(variable_id) + + def list_available_variables() -> str: """Return a formatted list of available variables.""" seen: set[str] = set() @@ -504,7 +681,11 @@ class AgentConfig: max_tokens: int = 4096 # Data Settings - data_source: str = "earthmover-public/era5-surface-aws" + data_source: str = "earthmover-public/era5" + # Zarr group holding surface/single-level variables. Queries resolve to + # f"{data_group}/{query_type}", so this travels with data_source — the two + # describe the same store and must be changed together. + data_group: str = "single" default_query_type: str = "temporal" max_download_size_gb: float = 15.0 @@ -592,9 +773,9 @@ class AgentConfig: **⚠️ CRITICAL:** When `calculate_maritime_route` returns a bounding box, USE THOSE EXACT VALUES for min/max longitude. Do NOT convert to 0-360! -**DATA AVAILABILITY:** 1975 to present (updated regularly) +**DATA AVAILABILITY:** 1940-01-01 to 2025-12-31, updated quarterly -**Available Variables (22 total):** +**Available Variables (38 total):** | Variable | Description | Units | Category | |----------|-------------|-------|----------| | sst | Sea Surface Temperature | K | Ocean | @@ -605,21 +786,47 @@ class AgentConfig: | v10 | 10m V-Wind (Northward) | m/s | Wind | | u100 | 100m U-Wind (Eastward) | m/s | Wind | | v100 | 100m V-Wind (Northward) | m/s | Wind | +| fg10 | 10m Wind Gust | m/s | Wind | +| zust | Friction Velocity | m/s | Wind | | sp | Surface Pressure | Pa | Pressure | | mslp | Mean Sea Level Pressure | Pa | Pressure | | blh | Boundary Layer Height | m | Atmosphere | | cape | Convective Available Potential Energy | J/kg | Atmosphere | | tcc | Total Cloud Cover | 0-1 | Cloud | +| hcc | High Cloud Cover | 0-1 | Cloud | +| mcc | Medium Cloud Cover | 0-1 | Cloud | +| lcc | Low Cloud Cover | 0-1 | Cloud | | cp | Convective Precipitation | m | Precipitation | | lsp | Large-scale Precipitation | m | Precipitation | | tp | Total Precipitation | m | Precipitation | +| sf | Snowfall | m water eq. | Precipitation | | ssr | Surface Net Solar Radiation | J/m² | Radiation | | ssrd | Surface Solar Radiation Downwards | J/m² | Radiation | +| fdir | Total Sky Direct Solar Radiation | J/m² | Radiation | +| tisr | TOA Incident Solar Radiation | J/m² | Radiation | +| tsr | Top Net Solar Radiation | J/m² | Radiation | +| slhf | Surface Latent Heat Flux | J/m² | Radiation | +| ie | Instantaneous Moisture Flux | kg/m²/s | Atmosphere | | tcw | Total Column Water | kg/m² | Moisture | | tcwv | Total Column Water Vapour | kg/m² | Moisture | | sd | Snow Depth | m water eq. | Land | | stl1 | Soil Temperature Level 1 | K | Land | +| stl2 | Soil Temperature Level 2 | K | Land | +| stl3 | Soil Temperature Level 3 | K | Land | +| stl4 | Soil Temperature Level 4 | K | Land | | swvl1 | Volumetric Soil Water Layer 1 | m³/m³ | Land | +| fsr | Forecast Surface Roughness | m | Land | + +⚠️ Pressure-level variables (temperature, wind, humidity, geopotential at 13 +levels from 1000-50 hPa) exist in the archive but are NOT yet wired up to +`retrieve_era5_data` — surface/single-level variables only for now. + +**SHAPE OF THE DOWNLOADED FILE:** +- Dimensions: `(time, latitude, longitude)`. Latitude descends 90 → -90. +- Coordinates: `time`, `latitude`, `longitude`. Nothing else. +- Data variable: named by the short name you requested (e.g. `t2`, not `t2m`). +- No land-sea mask is provided. `sst` is already NaN over land, so mask ocean + points with `ds.t2.where(ds.sst.notnull())` if you need one. ### 2. CUSTOM ANALYSIS: `python_repl` Persistent Python kernel for custom analysis and visualization. diff --git a/src/eurus/retrieval.py b/src/eurus/retrieval.py index 554b97f..fadfa63 100644 --- a/src/eurus/retrieval.py +++ b/src/eurus/retrieval.py @@ -10,10 +10,11 @@ import json import logging import os +import re import shutil import threading import time -from datetime import datetime, timedelta +from datetime import datetime from pathlib import Path from typing import Optional from urllib.request import Request, urlopen @@ -24,6 +25,7 @@ get_region, get_short_name, get_variable_info, + get_zarr_name, list_available_variables, ) from eurus.memory import get_memory @@ -33,6 +35,7 @@ def _arraylake_snippet( variable: str, + zarr_variable: str, query_type: str, start_date: str, end_date: str, @@ -45,6 +48,7 @@ def _arraylake_snippet( # Convert negative lons to 0-360 for ERA5 era5_min = min_lon % 360 if min_lon < 0 else min_lon era5_max = max_lon % 360 if max_lon < 0 else max_lon + group = f"{CONFIG.data_group}/{query_type}" return ( f"\n📦 Reproduce this download yourself (copy-paste into Jupyter):\n" f"```python\n" @@ -57,15 +61,18 @@ def _arraylake_snippet( f"\n" f"ds = xr.open_dataset(session.store, engine='zarr',\n" f" consolidated=False, zarr_format=3,\n" - f" chunks=None, group='{query_type}')\n" + f" chunks=None, group='{group}')\n" + f"ds = ds.rename({{'valid_time': 'time'}}) # store names the time dim 'valid_time'\n" f"\n" - f"subset = ds['{variable}'].sel(\n" + f"subset = ds['{zarr_variable}'].sel(\n" f" time=slice('{start_date}', '{end_date}'),\n" f" latitude=slice({max_lat}, {min_lat}), # ERA5: descending\n" f" longitude=slice({era5_min}, {era5_max}),\n" f")\n" f"\n" - f"subset.load().to_dataset(name='{variable}').to_zarr('my_data.zarr', mode='w')\n" + f"out = subset.load().to_dataset(name='{variable}')\n" + f"out = out.drop_vars('lsm', errors='ignore') # store ships lsm unwritten (all-NaN)\n" + f"out.to_zarr('my_data.zarr', mode='w')\n" f"```" ) @@ -77,6 +84,14 @@ def _format_coord(value: float) -> str: return f"{value:.2f}" +def source_tag(source: str) -> str: + """Reduce an Arraylake repo name to a filename-safe cache discriminator. + + e.g. "earthmover-public/era5" -> "earthmover-public-era5" + """ + return re.sub(r"[^0-9a-z]+", "-", source.lower()).strip("-") + + def generate_filename( variable: str, query_type: str, @@ -87,11 +102,18 @@ def generate_filename( min_longitude: float, max_longitude: float, region: Optional[str] = None, + source: Optional[str] = None, ) -> str: - """Generate a descriptive filename for the dataset.""" + """Generate a descriptive filename for the dataset. + + The source repo is part of the name: identical queries against different + ERA5 stores return different arrays (renamed variables, different dim + order, different time coverage), so they must not share a cache entry. + """ clean_var = variable.replace("_", "") clean_start = start.replace("-", "") clean_end = end.replace("-", "") + src = source_tag(source or CONFIG.data_source) if region: region_tag = region.lower() else: @@ -99,7 +121,7 @@ def generate_filename( f"lat{_format_coord(min_latitude)}_{_format_coord(max_latitude)}" f"_lon{_format_coord(min_longitude)}_{_format_coord(max_longitude)}" ) - return f"era5_{clean_var}_{query_type}_{clean_start}_{clean_end}_{region_tag}.zarr" + return f"{src}_{clean_var}_{query_type}_{clean_start}_{clean_end}_{region_tag}.zarr" def format_file_size(size_bytes: int) -> str: @@ -215,6 +237,17 @@ def retrieve_era5_data( """ memory = get_memory() + # Pure input validation first — no credentials or network needed. + # Only genuinely impossible dates are rejected here. The archive's coverage + # end advances with every quarterly update, so it is read from the store + # further down rather than duplicated as a constant that would go stale. + req_start = datetime.strptime(start_date, '%Y-%m-%d') + if req_start > datetime.now(): + return ( + f"Error: Requested start date ({start_date}) is in the future.\n" + f"ERA5 is a reanalysis of the past, refreshed quarterly." + ) + # Get API key: prefer explicit parameter, fall back to env var api_key = api_key or os.environ.get("ARRAYLAKE_API_KEY") if not api_key: @@ -255,18 +288,13 @@ def retrieve_era5_data( else: logger.warning(f"Unknown region '{region}', using provided coordinates") - # Resolve variable name + # Resolve variable name (catalog short_var for display/filenames, + # zarr_var for indexing into the store — they differ for a few + # renamed variables, e.g. "t2" -> "t2m") short_var = get_short_name(variable_id) + zarr_var = get_zarr_name(variable_id) var_info = get_variable_info(variable_id) - # Check for future / too-recent dates (ERA5T has a ~5-day processing lag) - req_start = datetime.strptime(start_date, '%Y-%m-%d') - if req_start > datetime.now() - timedelta(days=5): - return ( - f"Error: Requested start date ({start_date}) is too recent or in the future.\n" - f"ERA5 data has a ~5-day processing lag. Please request dates at least 5 days ago." - ) - # Setup paths output_dir = get_data_dir() filename = generate_filename( @@ -347,30 +375,30 @@ def retrieve_era5_data( repo = client.get_repo(CONFIG.data_source) session = repo.readonly_session("main") - logger.info(f"Opening {query_type} dataset...") + group = f"{CONFIG.data_group}/{query_type}" + logger.info(f"Opening {group} dataset...") ds = xr.open_dataset( session.store, engine="zarr", consolidated=False, zarr_format=3, chunks=None, - group=query_type, + group=group, ) + # Store's time coordinate is "valid_time" — normalize to "time" + # so the rest of this function (and downstream file outputs) + # stay unchanged. + if "valid_time" in ds.dims: + ds = ds.rename({"valid_time": "time"}) # Validate variable exists - # Auto-compute tp = cp + lsp if tp is not directly available - compute_tp = False - if short_var not in ds: - if short_var == "tp" and "cp" in ds and "lsp" in ds: - logger.info("Variable 'tp' not in store — will compute tp = cp + lsp") - compute_tp = True - else: - available = list(ds.data_vars) - return ( - f"Error: Variable '{short_var}' not found in dataset.\n" - f"Available variables: {', '.join(available)}\n\n" - f"Variable reference:\n{list_available_variables()}" - ) + if zarr_var not in ds: + available = list(ds.data_vars) + return ( + f"Error: Variable '{short_var}' not found in dataset.\n" + f"Available variables: {', '.join(available)}\n\n" + f"Variable reference:\n{list_available_variables()}" + ) # ERA5 latitude is stored 90 -> -90 (descending) lat_slice = slice(max_latitude, min_latitude) @@ -413,67 +441,59 @@ def retrieve_era5_data( # Subset both portions logger.info("Subsetting data (two-part: west + east of prime meridian)...") - fetch_vars = ["cp", "lsp"] if compute_tp else [short_var] - subsets_all = [] - for fv in fetch_vars: - subset_west = ds[fv].sel( - time=slice(start_date, end_date), - latitude=lat_slice, - longitude=west_slice, - ) - subset_east = ds[fv].sel( - time=slice(start_date, end_date), - latitude=lat_slice, - longitude=east_slice, - ) - - # Convert western longitudes from 360+ to negative (for -180/+180 output) - # e.g., 359.1 -> -0.9 - subset_west = subset_west.assign_coords( - longitude=subset_west.longitude - 360 - ) - - # Concatenate along longitude - subsets_all.append(xr.concat([subset_west, subset_east], dim='longitude')) - - if compute_tp: - subset = (subsets_all[0] + subsets_all[1]).rename("tp") - else: - subset = subsets_all[0] + subset_west = ds[zarr_var].sel( + time=slice(start_date, end_date), + latitude=lat_slice, + longitude=west_slice, + ) + subset_east = ds[zarr_var].sel( + time=slice(start_date, end_date), + latitude=lat_slice, + longitude=east_slice, + ) + + # Convert western longitudes from 360+ to negative (for -180/+180 output) + # e.g., 359.1 -> -0.9 + subset_west = subset_west.assign_coords( + longitude=subset_west.longitude - 360 + ) + + # Concatenate along longitude + subset = xr.concat([subset_west, subset_east], dim='longitude') else: # Normal case - no prime meridian crossing lon_slice = slice(req_min, req_max) # Subset the data logger.info("Subsetting data...") - fetch_vars = ["cp", "lsp"] if compute_tp else [short_var] - subsets_all = [] - for fv in fetch_vars: - subsets_all.append(ds[fv].sel( - time=slice(start_date, end_date), - latitude=lat_slice, - longitude=lon_slice, - )) - - if compute_tp: - subset = (subsets_all[0] + subsets_all[1]).rename("tp") - else: - subset = subsets_all[0] + subset = ds[zarr_var].sel( + time=slice(start_date, end_date), + latitude=lat_slice, + longitude=lon_slice, + ) # Convert to dataset ds_out = subset.to_dataset(name=short_var) + # The store declares an `lsm` land-sea-mask coordinate on every + # variable, but ships it unwritten — reads come back entirely NaN. + # Drop it so downloads don't carry an empty array that looks usable. + ds_out = ds_out.drop_vars("lsm", errors="ignore") + # Check for empty time dimension (no data in requested range) if ds_out.dims.get('time', 0) == 0: # Get actual data availability + time_min = ds['time'].min().values time_max = ds['time'].max().values import numpy as np + first_available = str(np.datetime_as_string(time_min, unit='D')) last_available = str(np.datetime_as_string(time_max, unit='D')) return ( f"Error: No data available for the requested time range.\n" f"Requested: {start_date} to {end_date}\n" - f"ERA5 data on Arraylake is available until {last_available}.\n\n" - f"Please request dates up to {last_available}." + f"This ERA5 archive covers {first_available} to {last_available} " + f"and is extended quarterly.\n\n" + f"Please request dates within that range." ) # Check for empty data (all NaNs) — only check 1st timestep @@ -484,15 +504,14 @@ def retrieve_era5_data( f"Error: The downloaded data for '{short_var}' is entirely empty (NaNs).\n" f"Possible causes:\n" f"1. The requested date/region has no data (e.g., SST over land).\n" - f"2. The request is too recent (ERA5T has a 5-day delay).\n" - f"3. Region bounds might be invalid or cross the prime meridian incorrectly." + f"2. Region bounds might be invalid or cross the prime meridian incorrectly." ) # Size guard — prevent downloading datasets larger than the configured limit estimated_gb = ds_out.nbytes / (1024 ** 3) if estimated_gb > CONFIG.max_download_size_gb: snippet = _arraylake_snippet( - short_var, query_type, start_date, end_date, + short_var, zarr_var, query_type, start_date, end_date, min_latitude, max_latitude, min_longitude if min_longitude >= 0 else min_longitude % 360, max_longitude if max_longitude >= 0 else max_longitude % 360, @@ -579,7 +598,7 @@ def retrieve_era5_data( time.sleep(wait_time) else: snippet = _arraylake_snippet( - short_var, query_type, start_date, end_date, + short_var, zarr_var, query_type, start_date, end_date, min_latitude, max_latitude, min_longitude if min_longitude >= 0 else min_longitude % 360, max_longitude if max_longitude >= 0 else max_longitude % 360, diff --git a/src/eurus/tools/era5.py b/src/eurus/tools/era5.py index 2b82049..ccb6db9 100644 --- a/src/eurus/tools/era5.py +++ b/src/eurus/tools/era5.py @@ -34,17 +34,21 @@ class ERA5RetrievalArgs(BaseModel): variable_id: str = Field( description=( - "ERA5 variable short name. Available variables (22 total):\n" + "ERA5 variable short name. Available variables (38 total):\n" "Ocean: sst (Sea Surface Temperature)\n" "Temperature: t2 (2m Air Temp), d2 (2m Dewpoint), skt (Skin Temp)\n" - "Wind 10m: u10 (Eastward), v10 (Northward)\n" + "Wind 10m: u10 (Eastward), v10 (Northward), fg10 (Gust)\n" "Wind 100m: u100 (Eastward), v100 (Northward)\n" + "Wind (other): zust (Friction Velocity)\n" "Pressure: sp (Surface), mslp (Mean Sea Level)\n" "Boundary Layer: blh (BL Height), cape (CAPE)\n" - "Cloud/Precip: tcc (Cloud Cover), cp (Convective), lsp (Large-scale), tp (Total Precip)\n" - "Radiation: ssr (Net Solar), ssrd (Solar Downwards)\n" + "Cloud: tcc (Total), hcc (High), mcc (Medium), lcc (Low)\n" + "Precip: cp (Convective), lsp (Large-scale), tp (Total), sf (Snowfall)\n" + "Radiation: ssr (Net Solar), ssrd (Solar Downwards), fdir (Direct Solar), " + "tisr (TOA Incident), tsr (TOA Net), slhf (Latent Heat Flux), ie (Moisture Flux)\n" "Moisture: tcw (Total Column Water), tcwv (Water Vapour)\n" - "Land: sd (Snow Depth), stl1 (Soil Temp L1), swvl1 (Soil Water L1)" + "Land: sd (Snow Depth), stl1-4 (Soil Temp L1-L4), swvl1 (Soil Water L1), " + "fsr (Surface Roughness)" ) ) @@ -194,11 +198,12 @@ def retrieve_era5_data( "Retrieves ERA5 climate reanalysis data from Earthmover's cloud archive.\n\n" "⚠️ query_type is AUTO-DETECTED - you don't need to specify it!\n\n" "Just provide:\n" - "- variable_id: one of 22 ERA5 variables (sst, t2, d2, skt, u10, v10, u100, v100, " - "sp, mslp, blh, cape, tcc, cp, lsp, tp, ssr, ssrd, tcw, tcwv, sd, stl1, swvl1)\n" + "- variable_id: one of 38 ERA5 surface variables (sst, t2, d2, skt, u10, v10, u100, v100, " + "fg10, zust, sp, mslp, blh, cape, tcc, hcc, mcc, lcc, cp, lsp, tp, sf, ssr, ssrd, fdir, " + "tisr, tsr, slhf, ie, tcw, tcwv, sd, stl1, stl2, stl3, stl4, swvl1, fsr)\n" "- start_date, end_date: YYYY-MM-DD format\n" "- lat/lon bounds: Use values from maritime route bounding box!\n\n" - "DATA: 1975-2024.\n" + "DATA: 1940-2025 (quarterly updates).\n" "Returns file path. Load with: xr.open_zarr('PATH')" ) diff --git a/tests/test_config.py b/tests/test_config.py index 554fd99..11d2963 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -1,27 +1,98 @@ import pytest -from eurus.config import ERA5_VARIABLES, VARIABLE_ALIASES, get_variable_info, get_short_name +from eurus.config import ( + ERA5_VARIABLES, + VARIABLE_ALIASES, + get_variable_info, + get_short_name, + get_zarr_name, +) -# All 22 variables in the Arraylake dataset -ALL_ARRAYLAKE_VARS = [ - "blh", "cape", "cp", "d2", "lsp", "mslp", "sd", "skt", "sp", - "ssr", "ssrd", "sst", "stl1", "swvl1", "t2", "tcc", "tcw", - "tcwv", "u10", "u100", "v10", "v100", -] +# The 38 array names actually present in the "single" group of the +# earthmover-public/era5 Icechunk store. Catalog short names map onto these +# via ERA5Variable.zarr_name; most are identity, a few are renamed. +STORE_ARRAY_NAMES = { + "blh", "cape", "cp", "d2m", "fdir", "fg10", "fsr", "hcc", "ie", "lcc", + "lsp", "mcc", "msl", "sd", "sf", "skt", "slhf", "sp", "ssr", "ssrd", + "sst", "stl1", "stl2", "stl3", "stl4", "swvl1", "t2m", "tcc", "tcw", + "tcwv", "tisr", "tp", "tsr", "u10", "u100", "v10", "v100", "zust", +} -# tp is a derived/accumulated variable kept for convenience -ALL_CATALOG_VARS = sorted(ALL_ARRAYLAKE_VARS + ["tp"]) +# Catalog short names whose store array name differs. These are the renames +# that would silently break retrieval if zarr_name were wrong. +RENAMED_VARS = {"t2": "t2m", "d2": "d2m", "mslp": "msl"} +ALL_CATALOG_VARS = sorted(ERA5_VARIABLES) -def test_variable_catalog_has_all_22(): - """Every Arraylake variable must appear in ERA5_VARIABLES.""" - for var in ALL_ARRAYLAKE_VARS: - assert var in ERA5_VARIABLES, f"Missing variable: {var}" + +def test_catalog_covers_every_store_array(): + """Every array in the store must be reachable from some catalog entry.""" + reachable = {v.zarr_name or v.short_name for v in ERA5_VARIABLES.values()} + assert reachable == STORE_ARRAY_NAMES def test_total_variable_count(): - """Catalog should contain at least 22 variables (22 Arraylake + tp).""" - assert len(ERA5_VARIABLES) >= 22 + """Catalog should expose all 38 single-level store variables.""" + assert len(ERA5_VARIABLES) == len(STORE_ARRAY_NAMES) == 38 + + +@pytest.mark.parametrize("short_name,expected_zarr", sorted(RENAMED_VARS.items())) +def test_renamed_variables_resolve_to_store_names(short_name, expected_zarr): + """Renamed vars must keep their catalog name but index the store correctly.""" + assert get_short_name(short_name) == short_name + assert get_zarr_name(short_name) == expected_zarr + + +@pytest.mark.parametrize("short_name", ALL_CATALOG_VARS) +def test_zarr_name_exists_in_store(short_name): + """get_zarr_name must always return a real array name in the store.""" + assert get_zarr_name(short_name) in STORE_ARRAY_NAMES + + +@pytest.mark.parametrize("short_name", ALL_CATALOG_VARS) +def test_zarr_name_only_set_when_it_differs(short_name): + """Don't carry a redundant zarr_name equal to the short name.""" + var = ERA5_VARIABLES[short_name] + assert var.zarr_name != var.short_name, ( + f"{short_name}: zarr_name is redundant, leave it as None" + ) + + +def test_zarr_name_defaults_to_short_name(): + """Variables without a rename index the store under their short name.""" + assert get_zarr_name("sst") == "sst" + assert get_zarr_name("u10") == "u10" + assert get_zarr_name("tp") == "tp" + + +def test_zarr_name_resolves_through_aliases(): + """Aliases must resolve all the way to the store array name.""" + assert get_zarr_name("2m_temperature") == "t2m" + assert get_zarr_name("dewpoint") == "d2m" + assert get_zarr_name("mean_sea_level_pressure") == "msl" + + +def test_zarr_name_passes_through_unknown_variables(): + """Unknown names fall through unchanged rather than raising.""" + assert get_zarr_name("not_a_real_var") == "not_a_real_var" + + +@pytest.mark.parametrize("short_name", ALL_CATALOG_VARS) +def test_diverging_colormap_only_on_signed_variables(short_name): + """RdBu_r centres on zero, so it's wrong for one-sided quantities. + + Wind components and fluxes are signed and diverge meaningfully; gust + speed, roughness and radiation totals are non-negative and must not. + """ + var = ERA5_VARIABLES[short_name] + if var.colormap != "RdBu_r": + return + low, high = var.typical_range + assert low is not None and high is not None, f"{short_name}: RdBu_r needs a range" + assert low < 0 < high, ( + f"{short_name}: diverging RdBu_r but typical_range {var.typical_range} " + f"does not straddle zero — use a sequential colormap" + ) def test_variable_loading(): @@ -96,10 +167,14 @@ def test_agent_prompt_branding(): assert "PANGAEA" not in AGENT_SYSTEM_PROMPT -def test_agent_prompt_lists_all_variables(): - """System prompt should mention all 22 Arraylake variable short names.""" +@pytest.mark.parametrize("var", ALL_CATALOG_VARS) +def test_agent_prompt_lists_all_variables(var): + """System prompt's variable table should have a row for every catalog entry. + + Matches the table row rather than a bare substring — short names like + "sd", "sp" and "ie" occur inside ordinary prompt prose. + """ from eurus.config import AGENT_SYSTEM_PROMPT - for var in ALL_ARRAYLAKE_VARS: - assert var in AGENT_SYSTEM_PROMPT, ( - f"System prompt missing variable: {var}" + assert f"| {var} |" in AGENT_SYSTEM_PROMPT, ( + f"System prompt variable table missing row for: {var}" ) diff --git a/tests/test_edge_cases.py b/tests/test_edge_cases.py index 9d22b2b..dbc2545 100644 --- a/tests/test_edge_cases.py +++ b/tests/test_edge_cases.py @@ -10,6 +10,7 @@ import os import pytest +from datetime import datetime, timedelta from pathlib import Path from dotenv import load_dotenv @@ -45,6 +46,59 @@ def test_region_tag_overrides_coords(self): assert "mediterranean" in name assert "lat" not in name # region tag replaces coord string + def test_filename_includes_data_source(self): + """Current source must appear in the name so caches can't collide.""" + from eurus.retrieval import generate_filename + from eurus.config import CONFIG + name = generate_filename( + "sst", "temporal", "2023-01-01", "2023-01-31", + min_latitude=0, max_latitude=10, + min_longitude=0, max_longitude=10, + ) + assert name.startswith("earthmover-public-era5_") + assert CONFIG.data_source == "earthmover-public/era5" + + def test_different_sources_do_not_share_a_cache_entry(self): + """The same query against the old and new store must not collide.""" + from eurus.retrieval import generate_filename + kwargs = dict( + variable="sst", query_type="temporal", + start="2023-01-01", end="2023-01-31", + min_latitude=0, max_latitude=10, + min_longitude=0, max_longitude=10, + ) + new = generate_filename(**kwargs, source="earthmover-public/era5") + old = generate_filename(**kwargs, source="earthmover-public/era5-surface-aws") + assert new != old + + def test_source_tag_is_filename_safe(self): + from eurus.retrieval import source_tag + assert source_tag("earthmover-public/era5") == "earthmover-public-era5" + assert source_tag("Org/Repo_Name.v2") == "org-repo-name-v2" + assert "/" not in source_tag("a/b/c") + + +class TestGroupPath: + """The zarr group path must track CONFIG, not be hardcoded.""" + + def test_snippet_group_follows_config(self, monkeypatch): + from eurus.config import CONFIG + from eurus.retrieval import _arraylake_snippet + monkeypatch.setattr(CONFIG, "data_group", "pressure") + snippet = _arraylake_snippet( + "t", "t", "spatial", "2020-01-01", "2020-01-02", 0, 10, 0, 10, + ) + assert "group='pressure/spatial'" in snippet + + def test_snippet_group_defaults_to_single(self): + from eurus.config import CONFIG + from eurus.retrieval import _arraylake_snippet + assert CONFIG.data_group == "single" + snippet = _arraylake_snippet( + "t2", "t2m", "temporal", "2020-01-01", "2020-01-02", 0, 10, 0, 10, + ) + assert "group='single/temporal'" in snippet + def test_format_coord_near_zero(self): from eurus.retrieval import _format_coord assert _format_coord(0.003) == "0.00" @@ -55,7 +109,9 @@ def test_format_coord_near_zero(self): class TestFutureDateRejection: """Ensure retrieval rejects future start dates without touching the API.""" - def test_future_date_returns_error(self): + def test_future_date_returns_error(self, monkeypatch): + """Runs before the credential check, so it needs no API key.""" + monkeypatch.delenv("ARRAYLAKE_API_KEY", raising=False) from eurus.retrieval import retrieve_era5_data result = retrieve_era5_data( query_type="temporal", @@ -67,6 +123,28 @@ def test_future_date_returns_error(self): ) assert "future" in result.lower() assert "Error" in result + assert "ARRAYLAKE_API_KEY" not in result + + def test_recent_past_date_is_not_rejected(self, monkeypatch): + """A date days old must reach the store, not trip a stale lag guard. + + The old guard rejected anything newer than now-5d for an "ERA5T + processing lag" that does not apply to this quarterly archive. + """ + monkeypatch.delenv("ARRAYLAKE_API_KEY", raising=False) + from eurus.retrieval import retrieve_era5_data + recent = (datetime.now() - timedelta(days=2)).strftime("%Y-%m-%d") + result = retrieve_era5_data( + query_type="temporal", + variable_id="sst", + start_date=recent, + end_date=recent, + min_latitude=0, max_latitude=10, + min_longitude=250, max_longitude=260, + ) + # Validation passed; it got as far as needing credentials. + assert "future" not in result.lower() + assert "ARRAYLAKE_API_KEY" in result # ============================================================================ diff --git a/tests/test_server_integration.py b/tests/test_server_integration.py index 7b77633..6be201f 100644 --- a/tests/test_server_integration.py +++ b/tests/test_server_integration.py @@ -94,7 +94,7 @@ def test_ensure_aws_region_sets_env_from_repo_metadata(self, monkeypatch): context_manager.__enter__.return_value = response with patch("eurus.retrieval.urlopen", return_value=context_manager) as mock_urlopen: - _ensure_aws_region("token", "earthmover-public/era5-surface-aws") + _ensure_aws_region("token", "earthmover-public/era5") assert os.environ["AWS_REGION"] == "eu-north-1" assert os.environ["AWS_DEFAULT_REGION"] == "eu-north-1" @@ -102,7 +102,7 @@ def test_ensure_aws_region_sets_env_from_repo_metadata(self, monkeypatch): assert os.environ["AWS_S3_ENDPOINT"] == "https://s3.eu-north-1.amazonaws.com" req = mock_urlopen.call_args.args[0] - assert req.full_url == "https://api.earthmover.io/repos/earthmover-public/era5-surface-aws" + assert req.full_url == "https://api.earthmover.io/repos/earthmover-public/era5" def test_ensure_aws_region_does_not_override_existing_env(self, monkeypatch): """Keep explicit user-provided AWS endpoint config untouched."""