Skip to content

Commit

Permalink
Rename user-visible arg to encode_cf to be coherent with xarray `…
Browse files Browse the repository at this point in the history 
…decode_cf`. Ref #23
  • Loading branch information
@alexamici
alexamici committed Nov 3, 2018
1 parent e228736 commit bf9e894
Show file tree
Hide file tree
Showing 3 changed files with 37 additions and 35 deletions.
34 changes: 18 additions & 16 deletions cfgrib/dataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -223,14 +223,14 @@ def dtype(self):

def build_geography_coordinates(
index, # type: messages.FileIndex
cfencode, # type: T.Sequence[str]
encode_cf, # type: T.Sequence[str]
log=LOG, # type: logging.Logger
):
# type: (...) -> T.Tuple[T.Tuple[str, ...], T.Tuple[int, ...], T.Dict[str, Variable]]
first = index.first()
geo_coord_vars = collections.OrderedDict() # type: T.Dict[str, Variable]
grid_type = index.getone('gridType')
if 'geography' in cfencode and grid_type in GRID_TYPES_COORD_VAR:
if 'geography' in encode_cf and grid_type in GRID_TYPES_COORD_VAR:
geo_dims = ('latitude', 'longitude') # type: T.Tuple[str, ...]
geo_shape = (index.getone('Nj'), index.getone('Ni')) # type: T.Tuple[int, ...]
geo_coord_vars['latitude'] = Variable(
Expand All @@ -241,7 +241,7 @@ def build_geography_coordinates(
dimensions=('longitude',), data=np.array(first['distinctLongitudes']),
attributes=cfmessage.COORD_ATTRS['longitude'],
)
elif 'geography' in cfencode and grid_type in GRID_TYPES_2D_AUX_COORD_VAR:
elif 'geography' in encode_cf and grid_type in GRID_TYPES_2D_AUX_COORD_VAR:
geo_dims = ('y', 'x')
geo_shape = (index.getone('Ny'), index.getone('Nx'))
geo_coord_vars['latitude'] = Variable(
Expand Down Expand Up @@ -272,31 +272,31 @@ def build_geography_coordinates(
return geo_dims, geo_shape, geo_coord_vars


def do_cfencode_first(data_var_attrs, coords_map, cfencode):
if 'parameter' in cfencode:
def encode_cf_first(data_var_attrs, coords_map, encode_cf):
if 'parameter' in encode_cf:
if 'GRIB_cfName' in data_var_attrs:
data_var_attrs['standard_name'] = data_var_attrs['GRIB_cfName']
if 'GRIB_name' in data_var_attrs:
data_var_attrs['long_name'] = data_var_attrs['GRIB_name']
if 'GRIB_units' in data_var_attrs:
data_var_attrs['units'] = data_var_attrs['GRIB_units']
if 'time' in cfencode:
if 'time' in encode_cf:
coords_map.extend(REF_TIME_COORDINATE_MAP)
else:
coords_map.extend(DATA_TIME_COORDINATE_MAP)
if 'vertical' in cfencode and data_var_attrs.get('GRIB_typeOfLevel') in PLEV_TYPE_OF_LEVELS:
if 'vertical' in encode_cf and data_var_attrs.get('GRIB_typeOfLevel') in PLEV_TYPE_OF_LEVELS:
coords_map.extend(PLEV_COORDINATE_MAP)
else:
coords_map.extend(VERTICAL_COORDINATE_MAP)


def build_data_var_components(index, cfencode=(), filter_by_keys={}, log=LOG):
def build_variable_components(index, encode_cf=(), filter_by_keys={}, log=LOG):
data_var_attrs_keys = DATA_ATTRIBUTES_KEYS[:]
data_var_attrs_keys.extend(GRID_TYPE_MAP.get(index.getone('gridType'), []))
data_var_attrs = enforce_unique_attributes(index, data_var_attrs_keys, filter_by_keys)
coords_map = HEADER_COORDINATES_MAP[:]

do_cfencode_first(data_var_attrs, coords_map, cfencode)
encode_cf_first(data_var_attrs, coords_map, encode_cf)

coord_vars = collections.OrderedDict()
for coord_key, increasing in coords_map:
Expand All @@ -315,7 +315,7 @@ def build_data_var_components(index, cfencode=(), filter_by_keys={}, log=LOG):
header_dimensions = tuple(d for d, c in coord_vars.items() if c.data.size > 1)
header_shape = tuple(coord_vars[d].data.size for d in header_dimensions)

geo_dims, geo_shape, geo_coord_vars = build_geography_coordinates(index, cfencode)
geo_dims, geo_shape, geo_coord_vars = build_geography_coordinates(index, encode_cf)
dimensions = header_dimensions + geo_dims
shape = header_shape + geo_shape
coord_vars.update(geo_coord_vars)
Expand All @@ -333,7 +333,7 @@ def build_data_var_components(index, cfencode=(), filter_by_keys={}, log=LOG):
geo_ndim=len(geo_dims),
)

if 'time' in coord_vars and 'time' in cfencode:
if 'time' in coord_vars and 'time' in encode_cf:
# add the 'valid_time' secondary coordinate
step_data = coord_vars['step'].data if 'data' in coord_vars else np.array(0.)
dims, time_data = cfmessage.build_valid_time(
Expand All @@ -342,7 +342,7 @@ def build_data_var_components(index, cfencode=(), filter_by_keys={}, log=LOG):
attrs = cfmessage.COORD_ATTRS['valid_time']
coord_vars['valid_time'] = Variable(dimensions=dims, data=time_data, attributes=attrs)

if 'level' in coord_vars and 'vertical' in cfencode and 'GRIB_typeOfLevel' in data_var_attrs:
if 'level' in coord_vars and 'vertical' in encode_cf and 'GRIB_typeOfLevel' in data_var_attrs:
type_of_level = data_var_attrs['GRIB_typeOfLevel']
coord_vars = collections.OrderedDict(
(type_of_level if k == 'level' else k, v) for k, v in coord_vars.items()
Expand All @@ -369,7 +369,7 @@ def dict_merge(master, update):

def build_dataset_components(
stream, indexpath='{path}.{short_hash}.idx', filter_by_keys={},
cfencode=('parameter', 'time', 'geography', 'vertical'), log=LOG,
encode_cf=('parameter', 'time', 'geography', 'vertical'), log=LOG,
):
filter_by_keys = dict(filter_by_keys)
index = stream.index(ALL_KEYS, indexpath=indexpath).subindex(filter_by_keys)
Expand All @@ -378,8 +378,10 @@ def build_dataset_components(
variables = collections.OrderedDict()
for param_id, short_name, var_name in zip(param_ids, index['shortName'], index['cfVarName']):
var_index = index.subindex(paramId=param_id)
dims, data_var, coord_vars = build_data_var_components(var_index, cfencode, filter_by_keys)
if 'parameter' in cfencode and var_name not in ('undef', 'unknown'):
dims, data_var, coord_vars = build_variable_components(
var_index, encode_cf, filter_by_keys,
)
if 'parameter' in encode_cf and var_name not in ('undef', 'unknown'):
short_name = var_name
vars = collections.OrderedDict([(short_name, data_var)])
vars.update(coord_vars)
Expand All @@ -394,7 +396,7 @@ def build_dataset_components(
encoding = {
'source': stream.path,
'filter_by_keys': filter_by_keys,
'cfencode': cfencode,
'encode_cf': encode_cf,
}
open_text = ', '.join('%s=%r' % it for it in encoding.items())
attributes['history'] = 'GRIB to CDM+CF via ' \
Expand Down
22 changes: 11 additions & 11 deletions tests/test_30_dataset.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -27,19 +27,19 @@ def test_dict_merge():

def test_build_data_var_components_no_encode():
index = messages.FileStream(path=TEST_DATA).index(dataset.ALL_KEYS).subindex(paramId=130)
dims, data_var, coord_vars = dataset.build_data_var_components(index=index)
dims, data_var, coord_vars = dataset.build_variable_components(index=index)
assert dims == {'number': 10, 'dataDate': 2, 'dataTime': 2, 'level': 2, 'i': 7320}
assert data_var.data.shape == (10, 2, 2, 2, 7320)

# equivalent to not np.isnan without importing numpy
assert data_var.data[:, :, :, :, :].mean() > 0.


def test_build_data_var_components_cfencode_geography():
def test_build_data_var_components_encode_cf_geography():
stream = messages.FileStream(path=TEST_DATA, message_class=cfmessage.CfMessage)
index = stream.index(dataset.ALL_KEYS).subindex(paramId=130)
dims, data_var, coord_vars = dataset.build_data_var_components(
index=index, cfencode='geography',
dims, data_var, coord_vars = dataset.build_variable_components(
index=index, encode_cf='geography',
)
assert dims == {
'number': 10, 'dataDate': 2, 'dataTime': 2,
Expand All @@ -62,16 +62,16 @@ def test_Dataset():

def test_Dataset_no_encode():
res = dataset.Dataset.from_path(
TEST_DATA, cfencode=()
TEST_DATA, encode_cf=()
)
assert 'history' in res.attributes
assert res.attributes['GRIB_edition'] == 1
assert tuple(res.dimensions.keys()) == ('number', 'dataDate', 'dataTime', 'level', 'i')
assert len(res.variables) == 9


def test_Dataset_cfencode_time():
res = dataset.Dataset.from_path(TEST_DATA, cfencode=('time',))
def test_Dataset_encode_cf_time():
res = dataset.Dataset.from_path(TEST_DATA, encode_cf=('time',))
assert 'history' in res.attributes
assert res.attributes['GRIB_edition'] == 1
assert tuple(res.dimensions.keys()) == ('number', 'time', 'level', 'i')
Expand All @@ -81,8 +81,8 @@ def test_Dataset_cfencode_time():
assert res.variables['t'].data[:, :, :, :].mean() > 0.


def test_Dataset_cfencode_geography():
res = dataset.Dataset.from_path(TEST_DATA, cfencode=('geography',))
def test_Dataset_encode_cf_geography():
res = dataset.Dataset.from_path(TEST_DATA, encode_cf=('geography',))
assert 'history' in res.attributes
assert res.attributes['GRIB_edition'] == 1
assert tuple(res.dimensions.keys()) == \
Expand All @@ -93,8 +93,8 @@ def test_Dataset_cfencode_geography():
assert res.variables['t'].data[:, :, :, :, :, :].mean() > 0.


def test_Dataset_cfencode_vertical():
res = dataset.Dataset.from_path(TEST_DATA, cfencode=('vertical',))
def test_Dataset_encode_cf_vertical():
res = dataset.Dataset.from_path(TEST_DATA, encode_cf=('vertical',))
assert 'history' in res.attributes
assert res.attributes['GRIB_edition'] == 1
assert tuple(res.dimensions.keys()) == ('number', 'dataDate', 'dataTime', 'isobaricInhPa', 'i')
Expand Down
16 changes: 8 additions & 8 deletions tests/test_40_xarray_store.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -16,13 +16,13 @@


def test_CfGribDataStore():
datastore = cfgrib_.CfGribDataStore(TEST_DATA, cfencode=())
datastore = cfgrib_.CfGribDataStore(TEST_DATA, encode_cf=())
expected = {'number': 10, 'dataDate': 2, 'dataTime': 2, 'level': 2, 'i': 7320}
assert datastore.get_dimensions() == expected


def test_xarray_open_dataset():
datastore = cfgrib_.CfGribDataStore(TEST_DATA, cfencode=())
datastore = cfgrib_.CfGribDataStore(TEST_DATA, encode_cf=())
res = xr.open_dataset(datastore)

assert res.attrs['GRIB_edition'] == 1
Expand Down Expand Up @@ -57,8 +57,8 @@ def test_open_dataset_corrupted():
xarray_store.open_dataset(TEST_CORRUPTED, backend_kwargs={'errors': 'strict'})


def test_open_dataset_cfencode_time():
backend_kwargs = {'cfencode': ('time',)}
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
Expand All @@ -69,8 +69,8 @@ def test_open_dataset_cfencode_time():
assert res['t'].mean() > 0.


def test_open_dataset_cfencode_vertical():
backend_kwargs = {'cfencode': ('vertical',)}
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']
Expand All @@ -79,8 +79,8 @@ def test_open_dataset_cfencode_vertical():
assert var.mean() > 0.


def test_open_dataset_cfencode_geography():
backend_kwargs = {'cfencode': ('geography',)}
def test_open_dataset_encode_cf_geography():
backend_kwargs = {'encode_cf': ('geography',)}
res = xarray_store.open_dataset(TEST_DATA, backend_kwargs=backend_kwargs)

assert res.attrs['GRIB_edition'] == 1
Expand Down

0 comments on commit bf9e894

Please sign in to comment.