Add support for met_em files (#46)

Co-authored-by: jthielen <jon@thielen.science>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

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):

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'

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

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:

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: '%'

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
 
 

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',
 }