How to use cfgrib - 10 common examples

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_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_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_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_reanalysis_Dataset(test_file):
    dataset, request, key_count = TEST_FILES[test_file]
    path = cdscommon.ensure_data(dataset, request, name='cds-' + test_file + '-{uuid}.grib')

    res = cfgrib.xarray_store.open_dataset(path)
    res.to_netcdf(path[:-5] + '.nc')

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)

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_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.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

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