diff --git a/tests/test_config.py b/tests/test_config.py index 15ef338a..1e5973ef 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -5,7 +5,7 @@ @pytest.fixture(scope='session') def config(): - return xwrf.config.config + return xwrf.config def test_defaults_exist(config): diff --git a/tests/test_postprocess.py b/tests/test_postprocess.py index 7527203c..8919cd22 100644 --- a/tests/test_postprocess.py +++ b/tests/test_postprocess.py @@ -13,6 +13,11 @@ def dummy_attrs_only_dataset(request): return xwrf.tutorial.open_dataset(request.param) +@pytest.fixture(scope='session', params=['met_em_sample']) +def met_em_dataset(request): + return xwrf.tutorial.open_dataset(request.param) + + @pytest.mark.parametrize('variable', ('Q2', 'PSFC')) def test_cf_attrs_added(dummy_dataset, variable): dataset = xwrf.postprocess._modify_attrs_to_cf(dummy_dataset) @@ -22,6 +27,12 @@ def test_cf_attrs_added(dummy_dataset, variable): @pytest.mark.parametrize('variable', ('THIS_IS_AN_IDEAL_RUN', 'SAVE_TOPO_FROM_REAL')) -def test_remove_units_from_bool_arrays(dummy_attrs_only_dataset, variable): - dataset = xwrf.postprocess._remove_units_from_bool_arrays(dummy_attrs_only_dataset) +def test_remove_invalid_units(dummy_attrs_only_dataset, variable): + dataset = xwrf.postprocess._make_units_pint_friendly(dummy_attrs_only_dataset) assert 'units' not in dataset[variable].attrs + + +@pytest.mark.parametrize('variable', ('OL4', 'GREENFRAC')) +def test_met_em_parsing(met_em_dataset, variable): + dataset = xwrf.postprocess._make_units_pint_friendly(met_em_dataset) + assert dataset[variable].attrs['units'] == '1' diff --git a/xwrf/__init__.py b/xwrf/__init__.py index 8c1c887b..b1a05282 100644 --- a/xwrf/__init__.py +++ b/xwrf/__init__.py @@ -4,8 +4,9 @@ from pkg_resources import DistributionNotFound, get_distribution -from . import config, postprocess, tutorial +from . import postprocess, tutorial from .accessors import WRFDataArrayAccessor, WRFDatasetAccessor +from .config import config try: __version__ = get_distribution(__name__).version diff --git a/xwrf/accessors.py b/xwrf/accessors.py index 8fd110c4..a9e46c03 100644 --- a/xwrf/accessors.py +++ b/xwrf/accessors.py @@ -5,8 +5,8 @@ from .postprocess import ( _collapse_time_dim, _decode_times, + _make_units_pint_friendly, _modify_attrs_to_cf, - _remove_units_from_bool_arrays, ) @@ -41,7 +41,7 @@ def postprocess(self, decode_times=True) -> xr.Dataset: """ ds = ( self.xarray_obj.pipe(_modify_attrs_to_cf) - .pipe(_remove_units_from_bool_arrays) + .pipe(_make_units_pint_friendly) .pipe(_collapse_time_dim) ) if decode_times: diff --git a/xwrf/config.yaml b/xwrf/config.yaml index 18646088..5ecec1d1 100644 --- a/xwrf/config.yaml +++ b/xwrf/config.yaml @@ -22,9 +22,20 @@ time_coords: - Time - time -boolean_units_attrs: - - '-' - - flag +unit_harmonization_map: + kelvin: + - Kelvin + '1': + - fraction + invalid: + - '-' + - flag + - '0/1 Flag' + - whoknows + - category + - none + meters: + - 'meters MSL' cf_attribute_map: ZNW: @@ -82,3 +93,18 @@ cf_attribute_map: standard_name: integral_of_surface_upward_heat_flux_in_air_wrt_time ACLHF: standard_name: integral_of_surface_upward_latent_heat_flux_in_air_wrf_time + VAR_SSO: + units: m2 + description: Variance of Subgrid Scale Orography MSL + HGT_V: + standard_name: surface_altitude + HGT_U: + standard_name: surface_altitude + HGT_M: + standard_name: surface_altitude + PRES: + units: Pa + ST: + units: kelvin + RH: + units: '%' diff --git a/xwrf/postprocess.py b/xwrf/postprocess.py index 2efa68c0..67808c9b 100644 --- a/xwrf/postprocess.py +++ b/xwrf/postprocess.py @@ -10,22 +10,34 @@ def _decode_times(ds: xr.Dataset) -> xr.Dataset: """ Decode the time variable to datetime64. """ - ds = ds.assign_coords( - { - 'Time': pd.to_datetime( - ds.Times.data.astype('str'), errors='raise', format='%Y-%m-%d_%H:%M:%S' - ) - } - ) + try: + _time = pd.to_datetime( + ds.Times.data.astype('str'), errors='raise', format='%Y-%m-%d_%H:%M:%S' + ) + except ValueError: + _time = pd.to_datetime( + ds.Times.data.astype('str'), errors='raise', format='%Y-%m-%dT%H:%M:%S.%f' + ) + ds = ds.assign_coords({'Time': _time}) ds.Time.attrs = {'long_name': 'Time', 'standard_name': 'time'} return ds -def _remove_units_from_bool_arrays(ds: xr.Dataset) -> xr.Dataset: - boolean_units_attrs = config.get('boolean_units_attrs') +def _make_units_pint_friendly(ds: xr.Dataset) -> xr.Dataset: + """ + Harmonizes awkward WRF units into pint-friendly ones + """ + # We have to invert the mapping from "new_unit -> wrf_units" to "wrf_unit -> new_unit" + wrf_units_map = { + v: k for (k, val_list) in config.get('unit_harmonization_map').items() for v in val_list + } for variable in ds.data_vars: - if ds[variable].attrs.get('units') in boolean_units_attrs: - ds[variable].attrs.pop('units', None) + if ds[variable].attrs.get('units') in wrf_units_map: + harmonized_unit = wrf_units_map[ds[variable].attrs['units']] + if harmonized_unit == 'invalid': + ds[variable].attrs.pop('units', None) + else: + ds[variable].attrs['units'] = harmonized_unit return ds diff --git a/xwrf/tutorial.py b/xwrf/tutorial.py index 45c5ab3c..0ac9269f 100644 --- a/xwrf/tutorial.py +++ b/xwrf/tutorial.py @@ -34,6 +34,7 @@ def _construct_cache_dir(path): 'lambert_conformal': 'data/lambert_conformal_sample.nc', 'mercator': 'data/mercator_sample.nc', 'tiny': 'data/tiny.nc', + 'met_em_sample': 'data/met_em.d01.2005-08-28_12:00:00.nc', }