Added earth-relative wind field calculation to base diagnostics

tests/test_postprocess.py

@@ -153,7 +153,11 @@ def test_clean_brackets_from_units(sample_dataset, variable, bracket):
 
 @pytest.mark.parametrize('sample_dataset', ['lambert_conformal'], indirect=True)
 def test_calc_base_diagnostics(sample_dataset):
-    subset = sample_dataset[['T', 'P', 'PB', 'PH', 'PHB']].isel(Time=0).load()
+    subset = (
+        sample_dataset[['T', 'P', 'PB', 'PH', 'PHB', 'U', 'V', 'SINALPHA', 'COSALPHA']]
+        .isel(Time=0)
+        .load()
+    )
     ds_dropped = subset.copy().pipe(xwrf.postprocess._calc_base_diagnostics)
     ds_undropped = subset.copy().pipe(xwrf.postprocess._calc_base_diagnostics, drop=False)
 
@@ -182,6 +186,18 @@ def test_calc_base_diagnostics(sample_dataset):
         assert ds['geopotential_height'].attrs['units'] == 'm'
         assert ds['geopotential_height'].attrs['standard_name'] == 'geopotential_height'
 
+        # Earth-relative eastward winds
+        np.testing.assert_allclose(ds['wind_east'].max().item(), 34.492916)
+        np.testing.assert_allclose(ds['wind_east'].min().item(), -3.2876422)
+        assert ds['wind_east'].attrs['units'] == 'm s-1'
+        assert ds['wind_east'].attrs['standard_name'] == 'eastward_wind'
+
+        # Earth-relative northward winds
+        np.testing.assert_allclose(ds['wind_north'].max().item(), 33.849540)
+        np.testing.assert_allclose(ds['wind_north'].min().item(), -13.434467)
+        assert ds['wind_north'].attrs['units'] == 'm s-1'
+        assert ds['wind_north'].attrs['standard_name'] == 'northward_wind'
+
     # Check dropped or not
     assert 'T' not in ds_dropped
     assert 'P' not in ds_dropped

xwrf/accessors.py

@@ -117,15 +117,16 @@ def postprocess(
         decode_times : bool, optional
             Decode the string-like wrfout times to xarray-friendly Pandas types. Defaults to True.
         calculate_diagnostic_variables : bool, optional
-            Calculate four essential diagnostic variables (potential temperature, air pressure,
+            Calculates essential diagnostic variables (potential temperature, air pressure,
             geopotential, and geopotential height) that are otherwise only present in wrfout files
-            as split components or dependent upon special adjustments. Defaults to True. If the
+            as split components or dependent upon special adjustments. Also calculates earth-relative
+            wind fields, as winds by default are grid-relative. Defaults to True. If the
             underlying fields on which any of these calculated fields depends is missing, that
             calculated variable is skipped. These will be eagerly evalulated, unless your data has
             been chunked with Dask, in which case these fields will also be Dask arrays.
         drop_diagnostic_variable_components : bool, optional
             Determine whether to drop the underlying fields used to calculate the diagnostic
-            variables. Defaults to True.
+            variables. Defaults to True. Never drops grid-relative wind fields.
 
         Returns
         -------

xwrf/postprocess.py

@@ -8,6 +8,7 @@
 import xarray as xr
 
 from .config import config
+from .destagger import _destag_variable
 from .grid import _wrf_grid_from_dataset
 
 
@@ -154,7 +155,11 @@ def _rename_dims(ds: xr.Dataset) -> xr.Dataset:
 
 
 def _calc_base_diagnostics(ds: xr.Dataset, drop: bool = True) -> xr.Dataset:
-    """Calculate the four basic fields that WRF does not have in physically meaningful form.
+    """Calculate the basic fields that WRF does not have in physically meaningful form.
+
+    Includes:
+        * diagnostics 'air_potential_temperature', 'air_pressure', 'geopotential', 'geopotential_height'
+        * earth-relative wind fields ('wind_east', 'wind_north')
 
     Parameters
     ----------
@@ -205,4 +210,22 @@ def _calc_base_diagnostics(ds: xr.Dataset, drop: bool = True) -> xr.Dataset:
         if drop:
             del ds['PH'], ds['PHB']
 
+    # Earth-relative wind fields (computed according to https://forum.mmm.ucar.edu/threads/how-do-i-convert-model-grid-relative-wind-to-earth-relative-wind-so-that-i-can-compare-model-wind-to-observations.179/)
+    if {'U', 'V', 'SINALPHA', 'COSALPHA'}.issubset(ds.data_vars):
+        u_model, v_model = _destag_variable(ds['U'].variable), _destag_variable(ds['V'].variable)
+        ds['wind_east'] = u_model * ds['COSALPHA'] - v_model * ds['SINALPHA']
+        ds['wind_north'] = v_model * ds['COSALPHA'] + u_model * ds['SINALPHA']
+        ds['wind_east'].attrs = dict(
+            description='earth-relative x-wind component',
+            standard_name='eastward_wind',
+            units='m s-1',
+            grid_mapping='wrf_projection',
+        )
+        ds['wind_north'].attrs = dict(
+            description='earth-relative y-wind component',
+            standard_name='northward_wind',
+            units='m s-1',
+            grid_mapping='wrf_projection',
+        )
+
     return ds