Dynamical.org Data sources#976
Conversation
Add analysis and forecast data sources driven directly off the public dynamical.org STAC catalog (stac.dynamical.org), reading anonymous Icechunk repositories. - Dynamical: analysis DataSource ([time, variable, lat, lon]) - DynamicalForecast: forecast ForecastSource ([time, lead_time, variable, lat, lon]), with ensemble-member selection for ensemble collections - DynamicalLexicon maps Earth2Studio variable ids to dynamical.org names; unit conversions (C->K, geopotential height->geopotential, %->fraction) are derived from the STAC `unit` field, with native-name pass-through - Collection selected by STAC id; projected grids (HRRR/MRMS) raise a clear error; longitude normalized to 0-360 - Reads use xr.open_zarr(..., chunks=None) with lazy per-chunk indexing - Adds icechunk to the `data` optional-dependency extra Signed-off-by: mrsll <m@mrshll.com>
Close create-datasource skill gaps: add available() (checked against each collection's STAC temporal extent), wire verbose to a tqdm progress bar, and factor the temporal dimension name into a shared _TIME_DIMENSION attr. Add availability tests. Signed-off-by: mrsll <m@mrshll.com>
…alAnalysis Address PR review feedback by exposing concrete, discoverable per-collection data sources alongside the generic base classes: - Rename Dynamical -> DynamicalAnalysis for symmetry with DynamicalForecast; both remain public as generic (collection-id) escape hatches. - Add concrete global sources: DynamicalGFSAnalysis, DynamicalGEFSAnalysis, DynamicalGFSForecast, DynamicalGEFSForecast, DynamicalIFSENSForecast, DynamicalAIFSForecast, DynamicalAIFSENSForecast (ensemble sources expose member). - Move @check_optional_dependencies() to _DynamicalBase so all subclasses inherit the icechunk guard. - Fix availability against open-ended STAC extents: fall back to the opened store's last coordinate for the upper bound so available() is accurate and out-of-range fetches raise a clear ValueError instead of a raw KeyError. - Fix forecast badge dataclass:forecast -> dataclass:simulation (repo convention). - Update exports, docs, tests, and CHANGELOG.
Address Greptile review feedback: - _time_extent guards against truncated/empty STAC extent lists - document icechunk region auto-detect fallback when unadvertised
|
/blossom-ci |
Sanity-Check Validation — Dynamical.org Data SourcesData Sources Validated
Validation ChecksFor each source, the sanity-check script verifies:
Key Findings
Sanity-Check Plot
Full validation script"""Sanity check script for all dynamical.org data sources in Earth2Studio.
Fetches a small sample from each dynamical data source (both analysis and
forecast) and produces a multi-panel plot showing 2m temperature fields.
Validates that data is non-NaN, coordinates are correctly oriented, and
unit conversions are applied.
Usage:
uv run python dynamical_sanity_check.py
uv run python dynamical_sanity_check.py --output my_plot.png
uv run python dynamical_sanity_check.py --variable u10m
"""
import argparse
import sys
from datetime import datetime, timedelta
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from earth2studio.data import (
DynamicalAIFS_FX,
DynamicalAIFSENS_FX,
DynamicalGEFS,
DynamicalGEFS_FX,
DynamicalGFS,
DynamicalGFS_FX,
DynamicalIFSENS_FX,
)
# Analysis sources: called with (time, variable)
ANALYSIS_SOURCES = {
"GFS": DynamicalGFS,
"GEFS": DynamicalGEFS,
}
# Forecast sources: called with (time, lead_time, variable)
FORECAST_SOURCES = {
"GFS_FX": DynamicalGFS_FX,
"GEFS_FX": DynamicalGEFS_FX,
"IFSENS_FX": DynamicalIFSENS_FX,
"AIFS_FX": DynamicalAIFS_FX,
"AIFSENS_FX": DynamicalAIFSENS_FX,
}
def check_analysis(name, klass, time, variable):
"""Fetch from an analysis source and run basic sanity checks."""
print(f"\n{chr(61)*60}")
print(f" {name}")
print(f"{chr(61)*60}")
ds = klass(verbose=True)
data = ds([time], [variable])
print(f" Shape: {data.shape}")
print(f" Dims: {list(data.dims)}")
print(f" Lat range: [{float(data.coords[lat].min()):.2f}, "
f"{float(data.coords[lat].max()):.2f}]")
print(f" Lon range: [{float(data.coords[lon].min()):.2f}, "
f"{float(data.coords[lon].max()):.2f}]")
values = data.sel(variable=variable).values
print(f" Value range: [{np.nanmin(values):.4f}, {np.nanmax(values):.4f}]")
print(f" Mean: {np.nanmean(values):.4f}")
print(f" NaN count: {np.isnan(values).sum()}")
# Sanity checks
assert data.coords["lat"].values[0] > data.coords["lat"].values[-1], \
"Latitude should be descending (90 -> -90)"
assert (data.coords["lon"].values >= 0).all(), \
"Longitude should be in [0, 360)"
assert not np.isnan(values).all(), "All values are NaN!"
if variable == "t2m":
assert np.nanmin(values) > 180.0, f"t2m min too low: {np.nanmin(values)}"
assert np.nanmax(values) < 350.0, f"t2m max too high: {np.nanmax(values)}"
print(" PASSED")
return data
def check_forecast(name, klass, time, lead_time, variable):
"""Fetch from a forecast source and run basic sanity checks."""
print(f"\n{chr(61)*60}")
print(f" {name}")
print(f"{chr(61)*60}")
# Ensemble sources accept a member kwarg
if klass in (DynamicalGEFS_FX, DynamicalIFSENS_FX,
DynamicalAIFSENS_FX):
ds = klass(member=0, verbose=True)
else:
ds = klass(verbose=True)
data = ds([time], [lead_time], [variable])
print(f" Shape: {data.shape}")
print(f" Dims: {list(data.dims)}")
print(f" Lat range: [{float(data.coords[lat].min()):.2f}, "
f"{float(data.coords[lat].max()):.2f}]")
print(f" Lon range: [{float(data.coords[lon].min()):.2f}, "
f"{float(data.coords[lon].max()):.2f}]")
values = data.sel(variable=variable).values
print(f" Value range: [{np.nanmin(values):.4f}, {np.nanmax(values):.4f}]")
print(f" Mean: {np.nanmean(values):.4f}")
print(f" NaN count: {np.isnan(values).sum()}")
# Sanity checks
assert list(data.dims) == ["time", "lead_time", "variable", "lat", "lon"], \
f"Unexpected dims: {list(data.dims)}"
assert data.coords["lat"].values[0] > data.coords["lat"].values[-1], \
"Latitude should be descending (90 -> -90)"
assert (data.coords["lon"].values >= 0).all(), \
"Longitude should be in [0, 360)"
assert not np.isnan(values).all(), "All values are NaN!"
if variable == "t2m":
assert np.nanmin(values) > 180.0, f"t2m min too low: {np.nanmin(values)}"
assert np.nanmax(values) < 350.0, f"t2m max too high: {np.nanmax(values)}"
print(" PASSED")
return data
def make_plot(results, variable, output_path):
"""Create a multi-panel plot of all fetched fields."""
n = len(results)
ncols = min(3, n)
nrows = (n + ncols - 1) // ncols
fig, axes = plt.subplots(nrows, ncols, figsize=(7 * ncols, 4 * nrows),
subplot_kw={"projection": None})
if n == 1:
axes = np.array([axes])
axes = np.atleast_2d(axes)
for idx, (name, data) in enumerate(results.items()):
row, col = divmod(idx, ncols)
ax = axes[row, col]
# Extract the 2D field (first time, first lead_time if present)
field = data.sel(variable=variable).isel(time=0)
if "lead_time" in field.dims:
field = field.isel(lead_time=0)
lat = data.coords["lat"].values
lon = data.coords["lon"].values
im = ax.pcolormesh(lon, lat, field.values, cmap="RdYlBu_r", shading="auto")
ax.set_title(name, fontsize=9)
ax.set_xlabel("Longitude")
ax.set_ylabel("Latitude")
fig.colorbar(im, ax=ax, shrink=0.8, label=variable)
# Hide unused axes
for idx in range(n, nrows * ncols):
row, col = divmod(idx, ncols)
axes[row, col].set_visible(False)
fig.suptitle(f"Dynamical.org Data Sources - {variable}", fontsize=12, y=1.02)
fig.tight_layout()
fig.savefig(output_path, dpi=150, bbox_inches="tight")
print(f"\nPlot saved to: {output_path}")
def main():
parser = argparse.ArgumentParser(
description="Sanity check for Earth2Studio dynamical.org data sources"
)
parser.add_argument(
"--time", default="2025-01-01T00:00:00",
help="Datetime to fetch (ISO format, default: 2025-01-01T00:00:00)"
)
parser.add_argument(
"--lead-time", type=int, default=0,
help="Lead time in hours for forecast sources (default: 24)"
)
parser.add_argument(
"--variable", default="t2m",
help="Variable to fetch (default: t2m)"
)
parser.add_argument(
"--output", default="dynamical_sanity_check.png",
help="Output plot filename (default: dynamical_sanity_check.png)"
)
parser.add_argument(
"--sources", nargs="+", default=None,
help="Subset of sources to check (default: all). "
"Use names like GFS, GEFS, GFS_FX, AIFS_FX, etc."
)
args = parser.parse_args()
time = datetime.fromisoformat(args.time)
lead_time = timedelta(hours=args.lead_time)
variable = args.variable
print(f"Dynamical.org Data Source Sanity Check")
print(f" Time: {time}")
print(f" Lead time: {lead_time} (forecast sources)")
print(f" Variable: {variable}")
print(f" Output: {args.output}")
results = {}
failures = []
# Check analysis sources
for name, klass in ANALYSIS_SOURCES.items():
if args.sources and name not in args.sources:
continue
try:
data = check_analysis(name, klass, time, variable)
results[name] = data
except Exception as e:
print(f" FAILED: {e}")
failures.append((name, str(e)))
# Check forecast sources
for name, klass in FORECAST_SOURCES.items():
if args.sources and name not in args.sources:
continue
try:
data = check_forecast(name, klass, time, lead_time, variable)
results[name] = data
except Exception as e:
print(f" FAILED: {e}")
failures.append((name, str(e)))
# Summary
print(f"\n{chr(61)*60}")
print(f" SUMMARY")
print(f"{chr(61)*60}")
print(f" Passed: {len(results)}/{len(results) + len(failures)}")
if failures:
print(f" Failures:")
for name, err in failures:
print(f" - {name}: {err}")
# Generate plot if we have any successful results
if results:
make_plot(results, variable, args.output)
else:
print("\nNo successful fetches -- no plot generated.")
sys.exit(1)
if failures:
sys.exit(1)
if __name__ == "__main__":
main() |
Greptile SummaryThis PR adds seven new data sources (
|
| Filename | Overview |
|---|---|
| earth2studio/data/dynamical.py | New 880-line data source module implementing lazy Icechunk-backed access to dynamical.org STAC collections; grid normalization, time validation, ensemble member selection, and the analysis/forecast dual-path split are all correct. The has_lead_time=False branch omits the ensemble_member isel step the forecast branch performs — benign for current best-estimate analysis collections, but would break transpose if a future analysis collection exposed an ensemble_member dimension. |
| earth2studio/lexicon/dynamical.py | Lexicon maps E2S variable ids to dynamical.org names and applies unit modifiers. The r2m entry falls into the identity branch despite dynamical.org serving relative humidity in percent (0-100), consistent with tcc which correctly converts percent to fraction. This silently returns values 100x too large for r2m requests. |
| test/data/test_dynamical.py | Comprehensive mock-based test suite covering grid normalization, unit conversions, native pass-through, ensemble member selection, analysis/forecast paths, time validation errors, projected-grid rejection, collection-id baking, and available(). Live-network slow/xfail tests also included. No issues found. |
| pyproject.toml | Adds icechunk>=2.0.0 constrained to Python >= 3.12 in the data extra; the previous icechunk 1.x/2.x compatibility concern is correctly resolved by the version floor. |
| earth2studio/lexicon/init.py | Exports the new DynamicalLexicon — change is correct and consistent with the module export pattern. |
| earth2studio/data/init.py | Exports all seven new Dynamical* classes — change is correct and alphabetically ordered. |
Reviews (2): Last reviewed commit: "Feedback" | Re-trigger Greptile
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
|
coverage: |
|
/blossom-ci |
negin513
left a comment
negin513
left a comment
There was a problem hiding this comment.
Nice work on this — the STAC resolution flow, grid normalization, and mock test coverage are all well done. Left a few inline comments below, mostly around a unit conversion issue and a couple of smaller things worth addressing before merge.
| import numpy as np | ||
| import xarray as xr | ||
| from loguru import logger | ||
| from tqdm.asyncio import tqdm |
There was a problem hiding this comment.
tqdm.asyncio.tqdm is the async-iterator variant meant for async for loops. fetch is declared async but never actually awaits anything — all the zarr reads are synchronous and block the event loop. Plain from tqdm import tqdm is the right import here.
| xr.Dataset | ||
| The grid-normalized, lazily opened dataset. | ||
| """ | ||
| if self._ds is not None: |
There was a problem hiding this comment.
Race condition: if two threads or coroutines call fetch() on the same instance at the same time, both can slip past this None check before either one sets self._ds, ending up with duplicate Icechunk sessions. Wrapping the lazy-init block in a threading.Lock would close the gap.
| xr.Dataset | ||
| Lazily opened dataset backed by the Icechunk session store. | ||
| """ | ||
| href = asset["href"] |
There was a problem hiding this comment.
asset["href"] can raise a bare KeyError if the STAC asset is malformed — the outer check only confirms "icechunk" is in assets, not that the asset dict itself has "href". A small guard with a ValueError here would give a much friendlier error message.
| # manages its own metadata, so zarr consolidated metadata does not apply. | ||
| return xr.open_zarr(session.store, consolidated=False, chunks=None) | ||
|
|
||
| def _setup_grid(self, ds: xr.Dataset) -> None: |
There was a problem hiding this comment.
Since fetch calls normalize_lat_lon on each slice after .sel(), _setup_grid exists only to populate self._lat and self._lon for the output array. Could we go one step further and call normalize_lat_lon on the full dataset once at open time instead? With chunks=None (no dask), sortby on a lazy store only reorders metadata — it doesn't load data. That would let us drop _setup_grid, self._lat, and self._lon entirely, and remove the normalize_lat_lon call from the fetch loop since the dataset would already be normalized.
| this collection. | ||
| """ | ||
| if variable in DynamicalLexicon.VOCAB: | ||
| dynamical_name, modifier = DynamicalLexicon.get_item(variable) |
There was a problem hiding this comment.
The PR description says "Unit conversions are driven by each collection's STAC-reported unit at fetch time", but self._cube_variables[dynamical_name]["unit"] is never actually read — the modifier comes entirely from DynamicalLexicon.get_item, which is hardcoded by variable name. In practice that's probably fine since dynamical.org is consistent, but if a collection ever serves temperature_2m in Kelvin rather than Celsius the code would silently add 273.15 on top. Worth either wiring up the STAC unit or clarifying in the description that conversions are static.
| if value is None: | ||
| return None | ||
| return datetime.fromisoformat(value.replace("Z", "+00:00")).replace( | ||
| tzinfo=None |
There was a problem hiding this comment.
Stripping tzinfo makes all internal datetimes naive, but there's no guard against callers passing timezone-aware datetimes (e.g. datetime.now(timezone.utc)). The comparison in _validate_time will blow up with a confusing TypeError. Easy fix — strip tzinfo at the top of _validate_time if it's set, or use the existing ensure_utc utility already in data/utils.py.
| xr_array = _sync_async(self.fetch, time, variable) | ||
| return xr_array | ||
|
|
||
| async def fetch( # type: ignore[override] |
There was a problem hiding this comment.
DynamicalGFS and DynamicalGEFS below have identical __call__ and fetch overrides — the only difference is the collection string passed to super().__init__. An _AnalysisBase(_DynamicalBase) in between would eliminate ~130 lines of duplication and make future analysis sources a one-liner.
| "eccodes>=2.39.2", | ||
| "eccodeslib>=2.39.2", | ||
| "ecmwf-opendata>=0.3.29", | ||
| "icechunk>=2.0.0; python_version>='3.12'", |
There was a problem hiding this comment.
PR description says icechunk>=1.0.0 but this pins >=2.0.0 — worth aligning so users reading the description know what they're actually getting.
There was a problem hiding this comment.
And I think since it brings in icechunk we should update the changelog? @NickGeneva
| self._lat = lat[self._lat_idx] | ||
| self._lon = lon_mod[self._lon_idx] | ||
|
|
||
| def _reorder(self, arr: np.ndarray) -> np.ndarray: |
There was a problem hiding this comment.
The _setup_grid / _reorder pair is doing something other data sources will likely need too. I'd suggest pulling this out into reusable utilities in data/utils.py rather than keeping it private to _DynamicalBase. Something like:
def normalize_latitude(da, ascending=False, lat_dim=None):
lat_dim = lat_dim or _find_dim(da, ("latitude", "lat"))
return da.sortby(lat_dim, ascending=ascending)
def normalize_longitude(da, lon_range=(0.0, 360.0), lon_dim=None):
lon_dim = lon_dim or _find_dim(da, ("longitude", "lon"))
start = lon_range[0]
wrapped = (da[lon_dim].values - start) % 360.0 + start
return da.assign_coords({lon_dim: wrapped}).sortby(lon_dim)
def normalize_lat_lon(da, lat_ascending=False, lon_range=(0.0, 360.0), lat_dim=None, lon_dim=None):
da = normalize_latitude(da, ascending=lat_ascending, lat_dim=lat_dim)
da = normalize_longitude(da, lon_range=lon_range, lon_dim=lon_dim)
return daA few benefits of this shape:
- Works on both
xr.DataArrayandxr.Dataset - Auto-detects
latitude/latandlongitude/lonso it isn't tied to dynamical.org's naming - Parameterized — callers that need
(-180, 180)or ascending lat can pass that in normalize_latitudeandnormalize_longitudeare independently useful
Then _setup_grid becomes a one-liner (ref = normalize_lat_lon(ref)) and the fetch loop just calls normalize_lat_lon(da) after .sel() instead of _reorder.
There was a problem hiding this comment.
And @NickGeneva I want to suggest us overall adding some data utils that would be reusuable? similar for da, obs utils, also for regrid, pressure interpol, etc?
3 participants
Earth2Studio Pull Request
Updates from PR: #947
Description
Adds gridded data sources backed by the public dynamical.org STAC catalog, reading directly from the anonymously accessible Icechunk repositories advertised in each collection's STAC metadata (no authentication):
Both take a STAC collection id (e.g.
"noaa-gfs-analysis","noaa-gfs-forecast","ecmwf-aifs-single-forecast"), resolve it againststac.dynamical.org/catalog.json, open the store lazily with xarray (chunks=None, the dynamical.org-recommended access pattern), and normalize the grid to the Earth2Studio convention (latitude 90→-90, longitude 0–360) via cheap numpy indexing on the fetched arrays.Variables (
DynamicalLexicon, with native pass-through for any collection variable not in the lexicon):u10m/v10mwind_u/v_10mmslpressure_reduced_to_mean_sea_levelu100m/v100mwind_u/v_100mtcctotal_cloud_cover_atmospheret2mtemperature_2mtcwvprecipitable_water_atmosphered2m/r2mdew_point/relative_humidity_2mz{500,850,925}geopotential_height_{lvl}hpasppressure_surfacet{500,850,925}temperature_{lvl}hpaUnit conversions are driven by each collection's STAC-reported
unitat fetch time (not baked into the lexicon, since units vary per collection): °C→K, geopotential height m→m² s⁻² (forz<level>), cloud cover %→fraction.Grid guardrails: only regular lat/lon collections are supported; projected datasets (e.g. HRRR, MRMS) raise a clear
ValueError. Requesting a level/variable a collection doesn't serve raises aKeyErrorlisting what's available.Deliberate deviations from the
earth2studio-create-datasourceskillThis source connects to a lazy Icechunk/zarr store, not the discrete-remote-file model the skill's async machinery targets. Two prescribed pieces are intentionally not adopted because forcing them would fight the storage model:
fetch/_create_tasks/gather_with_concurrency/async_retry). For an Icechunk store,xr.open_zarr(..., chunks=None)+.sel()reads only the requested chunks on indexing; this is the vendor-recommended pattern. The skill's "avoid xarray for loading / full file downloads" guidance targets the file-fetch model.cacheis a documented no-op. Icechunk fetches chunks on demand and manages its own caching; a local whole-file cache doesn't map. The param is retained for API parity.Also intentional:
available()is an instance method (validity is per-collection, checked against the STAC temporal extent); STAC JSON is fetched withurllib+ a custom User-Agent (dynamical.org 403s the default agent — metadata only, not bulk data).Checklist
Tests: mock (no-network) tests for both sources cover grid normalization, unit conversions, native pass-through, unknown-collection and projected-grid errors, and
available(). Real-network fetch tests are included asslow/xfail.Dependencies
icechunk>=1.0.0— added to thedataoptional-dependency extra (Apache-2.0). Guarded with the standardOptionalDependencyFailure/check_optional_dependenciespattern.