def test_ecmwf():
ds = xarray_store.open_dataset(TEST_DATA1)
res = cfcoords.translate_coords(ds, coord_model=datamodels.ECMWF)
assert set(
res.dims) == {"latitude", "level", "longitude", "number", "time"}
assert set(res.coords) == {
"latitude",
"level",
"longitude",
"number",
"step",
"time",
"valid_time",
}
ds = xarray_store.open_dataset(TEST_DATA2)
res = cfcoords.translate_coords(ds, coord_model=datamodels.ECMWF)
assert set(res.dims) == {"x", "y"}
assert set(res.coords) == {
"heightAboveGround",
"latitude",
"longitude",
"step",
"time",
"valid_time",
}
def test_ecmwf():
ds = xarray_store.open_dataset(TEST_DATA1)
res = cfcoords.translate_coords(ds, coord_model=datamodels.ECMWF)
assert set(
res.dims) == {'latitude', 'level', 'longitude', 'number', 'time'}
assert set(res.coords) == {
'latitude',
'level',
'longitude',
'number',
'step',
'time',
'valid_time',
}
ds = xarray_store.open_dataset(TEST_DATA2)
res = cfcoords.translate_coords(ds, coord_model=datamodels.ECMWF)
assert set(res.dims) == {'x', 'y'}
assert set(res.coords) == {
'heightAboveGround',
'latitude',
'longitude',
'step',
'time',
'valid_time',
}
def test_cds():
ds = xarray_store.open_dataset(TEST_DATA1)
res = cfcoords.translate_coords(ds, coord_model=datamodels.CDS)
assert set(res.dims) == {
"forecast_reference_time", "lat", "lon", "plev", "realization"
}
assert set(res.coords) == {
"forecast_reference_time",
"lat",
"leadtime",
"lon",
"plev",
"realization",
"time",
}
ds = xarray_store.open_dataset(TEST_DATA2)
res = cfcoords.translate_coords(ds, coord_model=datamodels.CDS)
assert set(res.dims) == {"x", "y"}
assert set(res.coords) == {
"forecast_reference_time",
"heightAboveGround",
"lat",
"leadtime",
"lon",
"time",
}
def test_cds():
ds = xarray_store.open_dataset(TEST_DATA1)
res = cfcoords.translate_coords(ds, coord_model=datamodels.CDS)
assert set(res.dims) == {'time', 'lat', 'lon', 'plev', 'realization'}
assert set(res.coords) == {
'forecast_reference_time',
'lat',
'leadtime',
'lon',
'plev',
'realization',
'time',
}
ds = xarray_store.open_dataset(TEST_DATA2)
res = cfcoords.translate_coords(ds, coord_model=datamodels.CDS)
assert set(res.dims) == {'x', 'y'}
assert set(res.coords) == {
'forecast_reference_time',
'heightAboveGround',
'lat',
'leadtime',
'lon',
'time',
}
def test_open_dataset_fail(grib_name):
grib_path = os.path.join(SAMPLE_DATA_FOLDER, grib_name + ".grib")
with pytest.raises(ValueError):
xarray_store.open_dataset(grib_path,
cache=False,
backend_kwargs={"errors": "raise"})
def test_open_dataset_corrupted():
res = xarray_store.open_dataset(TEST_CORRUPTED)
assert res.attrs['GRIB_edition'] == 1
assert len(res.data_vars) == 1
with pytest.raises(eccodes.EcCodesError):
xarray_store.open_dataset(TEST_CORRUPTED, errors='strict')
def test_open_dataset_corrupted():
res = xarray_store.open_dataset(TEST_CORRUPTED)
assert res.attrs['GRIB_edition'] == 1
assert len(res.data_vars) == 1
with pytest.raises(Exception):
xarray_store.open_dataset(TEST_CORRUPTED, backend_kwargs={'grib_errors': 'raise'})
def test_open_dataset_corrupted() -> None:
res = xarray_store.open_dataset(TEST_CORRUPTED)
assert res.attrs["GRIB_edition"] == 1
assert len(res.data_vars) == 1
with pytest.raises(Exception):
xarray_store.open_dataset(TEST_CORRUPTED,
backend_kwargs={"grib_errors": "raise"})
def test_to_grib_ecmwf(grib_name, tmpdir):
grib_path = os.path.join(SAMPLE_DATA_FOLDER, grib_name + '.grib')
out_path = str(tmpdir.join(grib_name + '.grib'))
res = xarray_store.open_dataset(grib_path, cache=False)
xarray_store.to_grib(res, out_path)
reread = xarray_store.open_dataset(out_path, cache=False)
assert res.equals(reread)
def test_canonical_dataset_to_grib(grib_name, tmpdir):
grib_path = os.path.join(SAMPLE_DATA_FOLDER, grib_name + '.grib')
out_path = str(tmpdir.join(grib_name + '.grib'))
res = xarray_store.open_dataset(grib_path)
xarray_to_grib.canonical_dataset_to_grib(res, out_path)
reread = xarray_store.open_dataset(out_path)
assert res.equals(reread)
def test_canonical_dataset_to_grib(grib_name, tmpdir):
grib_path = os.path.join(SAMPLE_DATA_FOLDER, grib_name + ".grib")
out_path = str(tmpdir.join(grib_name + ".grib"))
res = xarray_store.open_dataset(grib_path)
with pytest.warns(FutureWarning):
xarray_to_grib.canonical_dataset_to_grib(res, out_path)
reread = xarray_store.open_dataset(out_path)
assert res.equals(reread)
def test_all(cache):
da = xarray_store.open_dataset(TEST_DATA, cache=cache).data_vars['t']
va = da.values[:]
assert va.shape == (10, 4, 2, 61, 120)
assert da.mean() == va.mean()
def test_getitem_int(cache: bool) -> None:
da = xarray_store.open_dataset(TEST_DATA, cache=cache).data_vars["t"]
va = da.values[:]
assert da.isel(isobaricInhPa=1).values.shape == va[:, :, 1].shape
assert da.isel(isobaricInhPa=1).mean() == va[:, :, 1].mean()
assert da.sel(isobaricInhPa=500).mean() == va[:, :, 1].mean()
def test_getitem_int(cache):
da = xarray_store.open_dataset(TEST_DATA, cache=cache).data_vars['t']
va = da.values[:]
assert da.isel(air_pressure=1).values.shape == va[:, :, 1].shape
assert da.isel(air_pressure=1).mean() == va[:, :, 1].mean()
assert da.sel(air_pressure=500).mean() == va[:, :, 1].mean()
def test_open_dataset():
res = xarray_store.open_dataset(TEST_DATA)
assert res.attrs['GRIB_edition'] == 1
var = res['t']
assert var.attrs['GRIB_gridType'] == 'regular_ll'
assert var.attrs['units'] == 'K'
assert var.dims == \
('number', 'time', 'isobaricInhPa', 'latitude', 'longitude')
assert var.mean() > 0.
with pytest.warns(FutureWarning):
xarray_store.open_dataset(
TEST_DATA, filter_by_keys={'typeOfLevel': 'isobaricInhPa'})
def test_open_dataset_encode_cf_vertical():
backend_kwargs = {'encode_cf': ('vertical',)}
res = xarray_store.open_dataset(TEST_DATA, backend_kwargs=backend_kwargs)
var = res['t']
assert var.dims == ('number', 'dataDate', 'dataTime', 'isobaricInhPa', 'values')
assert var.mean() > 0.0
def test_getitem_list(cache):
da = xarray_store.open_dataset(TEST_DATA, cache=cache).data_vars['t']
va = da.values[:]
assert da.isel(number=[2, 3, 4, 5]).mean() == va[[2, 3, 4, 5]].mean()
assert da.isel(number=[4, 3, 2, 5]).mean() == va[[4, 3, 2, 5]].mean()
assert da.sel(number=[2, 3, 4, 5]).mean() == va[[2, 3, 4, 5]].mean()
assert da.sel(number=[4, 3, 2, 5]).mean() == va[[4, 3, 2, 5]].mean()
def test_getitem_latlon(cache):
da = xarray_store.open_dataset(TEST_DATA, cache=cache).data_vars['t']
va = da.values[:]
assert da.isel(latitude=slice(0, 3),
longitude=slice(0, 33)).mean() == va[..., :3, :33].mean()
assert da.sel(latitude=slice(90, 0),
longitude=slice(0, 90)).mean() == va[..., :31, :31].mean()
def test_open_dataset():
res = xarray_store.open_dataset(TEST_DATA)
assert res.attrs['GRIB_edition'] == 1
var = res['t']
assert var.attrs['GRIB_gridType'] == 'regular_ll'
assert var.attrs['units'] == 'K'
assert var.dims == \
('number', 'time', 'isobaricInhPa', 'latitude', 'longitude')
assert var.mean() > 0.
with pytest.raises(ValueError):
xarray_store.open_dataset(TEST_DATA, engine='netcdf4')
res = xarray_store.open_dataset(TEST_IGNORE,
backend_kwargs={'errors': 'warn'})
assert 'isobaricInhPa' in res.dims
res = xarray_store.open_dataset(TEST_IGNORE,
backend_kwargs={'errors': 'ignore'})
assert 'isobaricInhPa' in res.dims
with pytest.raises(ValueError):
xarray_store.open_dataset(TEST_IGNORE,
backend_kwargs={'errors': 'raise'})
def test_open_dataset() -> None:
res = xarray_store.open_dataset(TEST_DATA)
assert res.attrs["GRIB_edition"] == 1
var = res["t"]
assert var.attrs["GRIB_gridType"] == "regular_ll"
assert var.attrs["units"] == "K"
assert var.dims == ("number", "time", "isobaricInhPa", "latitude",
"longitude")
assert var.mean() > 0.0
with pytest.raises(ValueError):
xarray_store.open_dataset(TEST_DATA, engine="any-other-engine")
res = xarray_store.open_dataset(TEST_IGNORE,
backend_kwargs={"errors": "warn"})
assert "isobaricInhPa" in res.dims
res = xarray_store.open_dataset(TEST_IGNORE,
backend_kwargs={"errors": "ignore"})
assert "isobaricInhPa" in res.dims
with pytest.raises(ValueError):
xarray_store.open_dataset(TEST_IGNORE,
backend_kwargs={"errors": "raise"})
def test_open_dataset_encode_vertical():
res = xarray_store.open_dataset(TEST_DATA,
flavour_name='eccodes',
encode_vertical=True)
var = res['t']
assert var.dims == ('number', 'dataDate', 'dataTime', 'air_pressure', 'i')
assert var.mean() > 0.
def test_getitem_slice(cache):
da = xarray_store.open_dataset(TEST_DATA, cache=cache).data_vars['t']
va = da.values[:]
assert da.isel(number=slice(2, 6)).mean() == va[2:6].mean()
assert da.isel(number=slice(2, 6, 2)).mean() == va[2:6:2].mean()
# NOTE: label based indexing in xarray is inclusive of both the start and stop bounds.
assert da.sel(number=slice(2, 6)).mean() == va[2:7].mean()
assert da.sel(number=slice(2, 6, 2)).mean() == va[2:7:2].mean()
def test_open_dataset_encode_cf_vertical() -> None:
backend_kwargs = {"encode_cf": ("vertical", )}
res = xarray_store.open_dataset(TEST_DATA, backend_kwargs=backend_kwargs)
var = res["t"]
assert var.dims == ("number", "dataDate", "dataTime", "isobaricInhPa",
"values")
assert var.mean() > 0.0
def test_open_dataset_encode_cf_time() -> None:
backend_kwargs = {"encode_cf": ("time", )}
res = xarray_store.open_dataset(TEST_DATA, backend_kwargs=backend_kwargs)
assert res.attrs["GRIB_edition"] == 1
assert res["t"].attrs["GRIB_gridType"] == "regular_ll"
assert res["t"].attrs["GRIB_units"] == "K"
assert res["t"].dims == ("number", "time", "level", "values")
assert res["t"].mean() > 0.0
def test_open_dataset_encode_cf_time():
backend_kwargs = {'encode_cf': ('time', )}
res = xarray_store.open_dataset(TEST_DATA, backend_kwargs=backend_kwargs)
assert res.attrs['GRIB_edition'] == 1
assert res['t'].attrs['GRIB_gridType'] == 'regular_ll'
assert res['t'].attrs['GRIB_units'] == 'K'
assert res['t'].dims == ('number', 'time', 'level', 'values')
assert res['t'].mean() > 0.
def test_open_dataset_eccodes():
res = xarray_store.open_dataset(TEST_DATA)
assert res.attrs['GRIB_edition'] == 1
var = res['t']
assert var.attrs['GRIB_gridType'] == 'regular_ll'
assert var.attrs['units'] == 'K'
assert var.dims == ('number', 'time', 'isobaricInhPa', 'latitude', 'longitude')
assert var.mean() > 0.0
def test_open_dataset_encode_time():
res = xarray_store.open_dataset(TEST_DATA,
flavour_name='eccodes',
encode_time=True)
assert res.attrs['GRIB_edition'] == 1
assert res['t'].attrs['GRIB_gridType'] == 'regular_ll'
assert res['t'].attrs['units'] == 'K'
assert res['t'].dims == ('number', 'time', 'topLevel', 'i')
assert res['t'].mean() > 0.
def to_dataset(self):
# soft dependency on cfgrib
try:
from cfgrib import xarray_store
except ImportError:
print(
"Package cfgrib/xarray_store not found. Try running 'pip install cfgrib'."
)
raise
dataset = xarray_store.open_dataset(self.url())
return dataset
def test_open_dataset_extra_coords(grib_name: str, ndims: T.Any) -> None:
grib_path = os.path.join(SAMPLE_DATA_FOLDER, grib_name + ".grib")
res = xarray_store.open_dataset(
grib_path,
backend_kwargs={"extra_coords": {
"experimentVersionNumber": "time"
}},
cache=False,
)
assert "experimentVersionNumber" in res.coords
assert len(res["experimentVersionNumber"].dims) == ndims
def test_open_dataset_cds():
res = xarray_store.open_dataset(TEST_DATA, flavour_name='cds')
assert res.attrs['GRIB_edition'] == 1
var = res['t']
assert var.attrs['GRIB_gridType'] == 'regular_ll'
assert var.attrs['units'] == 'K'
assert var.dims == ('realization', 'forecast_reference_time', 'plev',
'lat', 'lon')
assert var.mean() > 0.