Repository: xarray-contrib/xwrf (language: python)

Issue:
First draft "destagger" function
<!-- Thanks for submitting a PR, your contribution is really appreciated! -->
<!-- Unless your change is trivial, please create an issue to discuss the change before creating a PR -->
<!-- Below are a few things we ask you kindly to self-check before getting a review. Remove checks that are not relevant.
-->

## Change Summary

Here's an attempt at a "destaggering" function. This is based on the function in WRF-python (https://github.com/NCAR/wrf-python/blob/22fb45c54f5193b849fdff0279445532c1a6c89f/src/wrf/destag.py). 

I've tested in on "east_west_stag" and "north_south_stag" coordinates. The function takes an xarray data-array and guesses the name of the staggered coordinate (it ends in "_stag"). If there is more than one (I don't think there are in WRF?), a NotImplenetedError is raised. 

I'm also not sure if this should ultimately look like this at all, but I wanted to go ahead and throw this code out there.


## Related issue number
This is related to issue #35 

<!-- Are there any issues opened that will be resolved by merging this change? -->
<!--
Please note any issues this fixes using [closing keywords]( https://help.github.com/articles/closing-issues-using-keywords/ ):
-->

## Checklist

I don't have any unit tests to check this -- I'm open to ideas on how to make unit tests (do they need to be on "real" data?)
Maybe that's a separate issue. 

- [ ] Unit tests for the changes exist
- [ ] Tests pass on CI
- [ ] Documentation reflects the changes where applicable

<!--
Please add any other relevant info below:
-->

I'm new to collaborating on open-source projects, and writing code for wide usage, so any feedback is welcome!


[FEATURE]: destagger function enabled for xarray
### Description

Hello,
I want to raise a hopefully straightforward issue to "dip my toes" into contributing to this project

One of the common tasks (alluded to in #10) dealing with wrf data is "destaggering" wrf variables on edge-cells to cell-center points.  This gets done in WRF-Python. Wind fields (U and V) are common examples 


### Implementation

In WRF-Python, this is done by simply taking the mean of the cells to the left/right or top/bottom of the center point. (https://github.com/NCAR/wrf-python/blob/22fb45c54f5193b849fdff0279445532c1a6c89f/src/wrf/destag.py).

In xarray this is easy enough to do using the ".loc" indexer and taking the mean of the the 1:n and 0:n-1 slices along the staggered coordinate. 

### Tests

This could be tested against the wrf-python feature. The dimensions of the output coordinate should be one less than the dimension of the staggered coordinate. 

### Questions

I think it makes sense to have a completely generic function to do this for one variable at a time (similar to the wrf-python version). I would picture such a function (let's say "_destagger_var"), being called by another class/function like "destagger_all_vars" to apply to all of the staggered vars... 

Maybe this is something that "xgcm" does and doesn't need to be handled separately? I'm not familiar with it.

I have a PR ready for one but I wanted to raise the issue first.

Thanks!
First draft "destagger" function
<!-- Thanks for submitting a PR, your contribution is really appreciated! -->
<!-- Unless your change is trivial, please create an issue to discuss the change before creating a PR -->
<!-- Below are a few things we ask you kindly to self-check before getting a review. Remove checks that are not relevant.
-->

## Change Summary

Here's an attempt at a "destaggering" function. This is based on the function in WRF-python (https://github.com/NCAR/wrf-python/blob/22fb45c54f5193b849fdff0279445532c1a6c89f/src/wrf/destag.py). 

I've tested in on "east_west_stag" and "north_south_stag" coordinates. The function takes an xarray data-array and guesses the name of the staggered coordinate (it ends in "_stag"). If there is more than one (I don't think there are in WRF?), a NotImplenetedError is raised. 

I'm also not sure if this should ultimately look like this at all, but I wanted to go ahead and throw this code out there.


## Related issue number
This is related to issue #35 

<!-- Are there any issues opened that will be resolved by merging this change? -->
<!--
Please note any issues this fixes using [closing keywords]( https://help.github.com/articles/closing-issues-using-keywords/ ):
-->

## Checklist

I don't have any unit tests to check this -- I'm open to ideas on how to make unit tests (do they need to be on "real" data?)
Maybe that's a separate issue. 

- [ ] Unit tests for the changes exist
- [ ] Tests pass on CI
- [ ] Documentation reflects the changes where applicable

<!--
Please add any other relevant info below:
-->

I'm new to collaborating on open-source projects, and writing code for wide usage, so any feedback is welcome!




Discussion / hints:

Fix patch:
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
index d55a4d9c..3cc27dec 100644
--- a/.github/workflows/ci.yaml
+++ b/.github/workflows/ci.yaml
@@ -23,7 +23,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ['3.7', '3.8', '3.9']
+        python-version: ['3.8', '3.9', '3.10']
     steps:
       - uses: actions/checkout@v2
       - uses: conda-incubator/setup-miniconda@v2
diff --git a/.github/workflows/upstream-dev-ci.yaml b/.github/workflows/upstream-dev-ci.yaml
index 7a1382c4..9041357a 100644
--- a/.github/workflows/upstream-dev-ci.yaml
+++ b/.github/workflows/upstream-dev-ci.yaml
@@ -22,7 +22,7 @@ jobs:
     strategy:
       fail-fast: false
       matrix:
-        python-version: ['3.9']
+        python-version: ['3.10']
     steps:
       - uses: actions/checkout@v2
       - uses: conda-incubator/setup-miniconda@v2
diff --git a/setup.py b/setup.py
index 8aacb1e3..3b1578c6 100644
--- a/setup.py
+++ b/setup.py
@@ -18,9 +18,9 @@
     'Intended Audience :: Science/Research',
     'Programming Language :: Python',
     'Programming Language :: Python :: 3',
-    'Programming Language :: Python :: 3.7',
     'Programming Language :: Python :: 3.8',
     'Programming Language :: Python :: 3.9',
+    'Programming Language :: Python :: 3.10',
     'Topic :: Scientific/Engineering',
 ]
 
@@ -29,7 +29,7 @@
     description='A lightweight interface for working with the Weather Research and Forecasting (WRF) model output in Xarray.',
     long_description=long_description,
     long_description_content_type='text/markdown',
-    python_requires='>=3.7',
+    python_requires='>=3.8',
     maintainer='xWRF Developers',
     classifiers=CLASSIFIERS,
     url='https://xwrf.readthedocs.io',
diff --git a/xwrf/accessors.py b/xwrf/accessors.py
index 17482fee..5473cc21 100644
--- a/xwrf/accessors.py
+++ b/xwrf/accessors.py
@@ -2,6 +2,7 @@
 
 import xarray as xr
 
+from .destagger import _destag_variable, _rename_staggered_coordinate
 from .postprocess import (
     _assign_coord_to_dim_of_different_name,
     _calc_base_diagnostics,
@@ -27,6 +28,66 @@ def __init__(self, xarray_obj: xr.Dataset | xr.DataArray) -> WRFAccessor:
 class WRFDataArrayAccessor(WRFAccessor):
     """Adds a number of WRF specific methods to xarray.DataArray objects."""
 
+    def destagger(
+        self,
+        stagger_dim: str | None = None,
+        unstaggered_dim_name: str | None = None,
+        exclude_staggered_auxiliary_coords: bool = False,
+    ) -> xr.DataArray:
+        """
+        Destagger a single WRF xarray.DataArray
+
+        Parameters
+        ----------
+        stagger_dim : str, optional
+            Name of dimension to unstagger. Defaults to guessing based on name (ends in "_stag")
+        unstaggered_dim_name : str, optional
+            String to which to rename the dimension after destaggering. Example would be
+            "west_east" for "west_east_stag". By default the dimenions will be renamed the text in
+            front of "_stag" from the "stagger_dim" field.
+        exclude_staggered_auxiliary_coords : bool, optional
+            If True, auxiliary coordinates (such as latitude and longitude) that originally include
+            staggered dimensions are excluded from being destaggered along with this DataArray's
+            data. If False, auxiliary coordinates are destaggered in the same fashion as the data
+            (average of cell boundary values), which may introduce error in comparison to the
+            analogous non-staggered coordinates from the original WRF data. Defaults to False.
+
+        Returns
+        -------
+        xarray.DataArray
+            The destaggered DataArray with renamed dimension and adjusted coordinates.
+
+        Notes
+        -----
+        While destaggered coordinate variables are made available by default
+        (exclude_staggered_auxiliary_coords=False), these are approximations made from the
+        staggered coordinate values and may not be sufficiently accurate for all grid projections
+        and/or use cases. For full accuracy, auxiliary coordinates should be re-computed from
+        dimension coordinates or obtained from the original dataset.
+        """
+        new_variable = _destag_variable(
+            self.xarray_obj.variable, stagger_dim=stagger_dim, unstag_dim_name=unstaggered_dim_name
+        )
+
+        # Need to recalculate staggered coordinates, as they don't already exist independently
+        # in a DataArray context
+        new_coords = {}
+        for coord_name, coord_data in self.xarray_obj.coords.items():
+            if set(coord_data.dims).difference(set(new_variable.dims)):
+                # Has a dimension not in the destaggered output (and so still staggered)
+                new_name = _rename_staggered_coordinate(
+                    coord_name, stagger_dim=stagger_dim, unstag_dim_name=unstaggered_dim_name
+                )
+                if not exclude_staggered_auxiliary_coords or new_name in new_variable.dims:
+                    # Skip if excluding and this isn't a dimension coordinate of output
+                    new_coords[new_name] = _destag_variable(
+                        coord_data, stagger_dim=stagger_dim, unstag_dim_name=unstaggered_dim_name
+                    )
+            else:
+                new_coords[coord_name] = coord_data.variable
+
+        return xr.DataArray(new_variable, coords=new_coords)
+
 
 @xr.register_dataset_accessor('xwrf')
 class WRFDatasetAccessor(WRFAccessor):
@@ -83,3 +144,49 @@ def postprocess(
         ds = ds.pipe(_include_projection_coordinates)
 
         return ds.pipe(_rename_dims)
+
+    def destagger(self, staggered_to_unstaggered_dims: dict[str, str] | None = None) -> xr.Dataset:
+        """
+        Destagger all data variables in a WRF xarray.Dataset
+
+        Parameters
+        ----------
+        staggered_to_unstaggered_dims : dict, optional
+            Mapping of target staggered dimensions to corresponding unstaggered dimensions
+
+        Returns
+        -------
+        xarray.Dataset
+            The destaggered dataset.
+
+        Notes
+        -----
+        Does not alter coordinates, only data variables. Staggered coordinates will remain on the
+        dataset, but will not be associated with any data variables.
+        """
+        staggered_dims = (
+            {dim for dim in self.xarray_obj.dims if dim.endswith('_stag')}
+            if staggered_to_unstaggered_dims is None
+            else set(staggered_to_unstaggered_dims)
+        )
+        new_data_vars = {}
+        for var_name, var_data in self.xarray_obj.data_vars.items():
+            if this_staggered_dims := set(var_data.dims).intersection(staggered_dims):
+                # Found a staggered dim
+                # TODO: should we raise an error if somehow end up with more than just one
+                # staggered dim, or just pick one from the set like below?
+                this_staggered_dim = this_staggered_dims.pop()
+                new_data_vars[var_name] = _destag_variable(
+                    var_data.variable,
+                    stagger_dim=this_staggered_dim,
+                    unstag_dim_name=(
+                        None
+                        if staggered_to_unstaggered_dims is None
+                        else staggered_to_unstaggered_dims[this_staggered_dim]
+                    ),
+                )
+            else:
+                # No staggered dims
+                new_data_vars[var_name] = var_data.variable
+
+        return xr.Dataset(new_data_vars, self.xarray_obj.coords, self.xarray_obj.attrs)
diff --git a/xwrf/destagger.py b/xwrf/destagger.py
new file mode 100644
index 00000000..9752db87
--- /dev/null
+++ b/xwrf/destagger.py
@@ -0,0 +1,90 @@
+import xarray as xr
+
+
+def _drop_attrs(attrs_dict, attrs_to_drop):
+    try:
+        return {k: v for k, v in attrs_dict.items() if k not in attrs_to_drop}
+    except AttributeError:
+        return None
+
+
+def _destag_variable(datavar, stagger_dim=None, unstag_dim_name=None):
+    """
+    Destaggering function for a single wrf xarray.Variable
+
+    Based off of the wrf-python destagger function
+    https://github.com/NCAR/wrf-python/blob/22fb45c54f5193b849fdff0279445532c1a6c89f/src/wrf/destag.py
+    Copyright 2016 University Corporation for Atmospheric Research, reused with modification
+    according to the terms of the Apache License, Version 2.0
+
+    Parameters
+    ----------
+    datavar : xarray.Variable
+        Data variable to be destaggered
+    stagger_dim : str, optional
+        Name of dimension to unstagger. Defaults to guessing based on name (ends in "_stag")
+    unstag_dim_name : str, option
+        String to which to rename the dimension after destaggering. Example would be
+        "west_east" for "west_east_stag". By default the dimenions will be renamed the text in
+        front of "_stag" from the "stagger_dim" field.
+
+    Returns
+    -------
+    xarray.Variable
+        The destaggered variable with renamed dimension
+    """
+    # get the coordinate to unstagger
+    # option 1) user has provided the dimension
+    if stagger_dim and stagger_dim not in datavar.dims:
+        # check that the user-passed in stag dim is actually in there
+        raise ValueError(f'{stagger_dim} not in {datavar.dims}')
+
+    # option 2) guess the staggered dimension
+    elif stagger_dim is None:
+        # guess the name of the coordinate
+        stagger_dim = [x for x in datavar.dims if x.endswith('_stag')]
+
+        if len(stagger_dim) > 1:
+            raise NotImplementedError(
+                'Expected a single destagger dimensions. Found multiple destagger dimensions: '
+                f'{stagger_dim}'
+            )
+
+        # we need a string, not a list
+        stagger_dim = stagger_dim[0]
+
+    # get the size of the staggereed coordinate
+    stagger_dim_size = datavar.sizes[stagger_dim]
+
+    # I think the "dict(a="...")"  format is preferrable... but you cant stick an fx arg string
+    # into that...
+    left_or_bottom_cells = datavar.isel({stagger_dim: slice(0, stagger_dim_size - 1)})
+    right_or_top_cells = datavar.isel({stagger_dim: slice(1, stagger_dim_size)})
+    center_mean = (left_or_bottom_cells + right_or_top_cells) * 0.5
+
+    # now change the variable name of the unstaggered coordinate
+    # we can pass this in if we want to, for whatever reason
+    if unstag_dim_name is None:
+        unstag_dim_name = stagger_dim.split('_stag')[
+            0
+        ]  # get the part of the name before the "_stag"
+
+    # return a data variable with renamed dimensions
+    return xr.Variable(
+        dims=tuple(str(unstag_dim_name) if dim == stagger_dim else dim for dim in center_mean.dims),
+        data=center_mean.data,
+        attrs=_drop_attrs(center_mean.attrs, ('stagger',)),
+        encoding=center_mean.encoding,
+        fastpath=True,
+    )
+
+
+def _rename_staggered_coordinate(name, stagger_dim=None, unstag_dim_name=None):
+    if name == stagger_dim and unstag_dim_name is not None:
+        return unstag_dim_name
+    elif name[-2:].lower() in ('_u', '_v'):
+        return name[:-2]
+    elif name[-5:].lower() == '_stag':
+        return name[:-5]
+    else:
+        return name


Test patch:
diff --git a/tests/test_accessors.py b/tests/test_accessors.py
index db7cb1de..a64db416 100644
--- a/tests/test_accessors.py
+++ b/tests/test_accessors.py
@@ -1,11 +1,17 @@
 import pandas as pd
 import pytest
+import xarray as xr
 
 import xwrf
 
 from . import importorskip
 
 
+@pytest.fixture(scope='session')
+def test_grid(request):
+    return xwrf.tutorial.open_dataset(request.param)
+
+
 @importorskip('cf_xarray')
 @pytest.mark.parametrize(
     'name, cf_grid_mapping_name', [('lambert_conformal', 'lambert_conformal_conic')]
@@ -59,3 +65,50 @@ def test_postprocess(name, cf_grid_mapping_name):
     assert 'PB' not in ds.data_vars
     assert 'PH' not in ds.data_vars
     assert 'PHB' not in ds.data_vars
+
+
+@pytest.mark.parametrize('test_grid', ['lambert_conformal', 'mercator'], indirect=True)
+def test_dataarray_destagger(test_grid):
+    data = test_grid['U']
+    destaggered = data.xwrf.destagger()
+
+    # Check shape reduction and dim name adjustment
+    assert destaggered.sizes['west_east'] == data.sizes['west_east_stag'] - 1
+
+    # Check coordinate reduction
+    xr.testing.assert_allclose(destaggered['XLAT'], test_grid['XLAT'])
+    xr.testing.assert_allclose(destaggered['XLONG'], test_grid['XLONG'])
+
+
+@pytest.mark.parametrize('test_grid', ['lambert_conformal', 'mercator'], indirect=True)
+def test_dataarray_destagger_with_exclude(test_grid):
+    data = test_grid['V']
+    destaggered = data.xwrf.destagger(exclude_staggered_auxiliary_coords=True)
+
+    # Check shape reduction and dim name adjustment
+    assert destaggered.sizes['south_north'] == data.sizes['south_north_stag'] - 1
+
+    # Verify that no XLAT/XLONG are present in output, but XTIME (which is auxiliary, but not
+    # staggered) remains
+    assert not any(coord_name.startswith('XL') for coord_name in destaggered.coords)
+    assert 'XTIME' in destaggered.coords
+
+
+@pytest.mark.parametrize('test_grid', ['lambert_conformal', 'mercator'], indirect=True)
+def test_dataset_destagger(test_grid):
+    destaggered = (
+        test_grid.isel(Time=slice(0, 2))
+        .xwrf.postprocess(calculate_diagnostic_variables=False)
+        .xwrf.destagger()
+    )
+
+    # Check elimination of staggered dims and "stagger" attr
+    for varname in destaggered.data_vars:
+        assert not {'x_stag', 'y_stag', 'z_stag'}.intersection(set(destaggered[varname].dims))
+        assert (
+            'stagger' not in destaggered[varname].attrs
+            or destaggered[varname].attrs['stagger'] == ''
+        )
+
+    # Check that attrs are preserved
+    assert destaggered.attrs == test_grid.attrs
diff --git a/tests/test_destagger.py b/tests/test_destagger.py
new file mode 100644
index 00000000..695068f8
--- /dev/null
+++ b/tests/test_destagger.py
@@ -0,0 +1,79 @@
+from collections import Counter
+
+import numpy as np
+import pytest
+import xarray as xr
+
+from xwrf.destagger import _destag_variable, _drop_attrs, _rename_staggered_coordinate
+
+
+@pytest.mark.parametrize(
+    'input_attrs, output_attrs',
+    [
+        ({'a': 0, 'b': 1, 'c': 2}, {'a': 0, 'c': 2}),
+        ({'b': 3}, {}),
+        ({'a': 4, 'c': 5}, {'a': 4, 'c': 5}),
+        (Counter(('a', 'a', 'b', 'c')), {'a': 2, 'c': 1}),
+    ],
+)
+def test_drop_attrs_successful(input_attrs, output_attrs):
+    result = _drop_attrs(input_attrs, ('b',))
+    assert result == output_attrs
+    assert isinstance(result, dict)
+
+
+def test_drop_attrs_unsuccessful():
+    assert _drop_attrs('not a Mapping', ('a',)) is None
+
+
+@pytest.mark.parametrize(
+    'input_name, stagger_dim, unstag_dim_name, output_name',
+    [
+        (
+            'bottom_top',
+            'bottom_top_stag',
+            None,
+            'bottom_top',
+        ),
+        ('bottom_top_stag', 'bottom_top_stag', 'z', 'z'),
+        ('bottom_top_stag', 'bottom_top_stag', None, 'bottom_top'),
+        ('XLAT_U', 'west_east', None, 'XLAT'),
+        ('XLONG_V', 'south_north', None, 'XLONG'),
+    ],
+)
+def test_rename_staggered_coordinate(input_name, stagger_dim, unstag_dim_name, output_name):
+    assert _rename_staggered_coordinate(input_name, stagger_dim, unstag_dim_name) == output_name
+
+
+def test_destag_variable_missing_dim():
+    with pytest.raises(ValueError):
+        _destag_variable(xr.Variable(('x', 'y'), np.zeros((2, 2))), 'z_stag')
+
+
+def test_destag_variable_multiple_dims():
+    with pytest.raises(NotImplementedError):
+        _destag_variable(xr.Variable(('x_stag', 'y_stag'), np.zeros((2, 2))))
+
+
+@pytest.mark.parametrize(
+    'unstag_dim_name, expected_output_dim_name',
+    [
+        ('z', 'z'),
+        (None, 'bottom_top'),
+    ],
+)
+def test_destag_variable_1d(unstag_dim_name, expected_output_dim_name):
+    staggered = xr.Variable(('bottom_top_stag',), np.arange(5), attrs={'stagger': 'Z'})
+    output = _destag_variable(staggered, unstag_dim_name=unstag_dim_name)
+    # Check values
+    np.testing.assert_array_almost_equal(output.values, 0.5 + np.arange(4))
+    # Check dim name
+    assert output.dims[0] == expected_output_dim_name
+    # Check attrs
+    assert not output.attrs
+
+
+def test_destag_variable_2d():
+    staggered = xr.Variable(('x', 'y_stag'), np.arange(9).reshape(3, 3))
+    expected = xr.Variable(('x', 'y'), [[0.5, 1.5], [3.5, 4.5], [6.5, 7.5]])
+    xr.testing.assert_equal(_destag_variable(staggered), expected)