Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

General update to regridding #535

Merged
merged 33 commits into from
Oct 10, 2023
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
33 commits
Select commit Hold shift + click to select a range
d0cd459
Remove tool specific methods from documentation, until decided one wa…
jasonb5 Aug 18, 2023
5a57b7c
Fixes adding missing vertical bounds
jasonb5 Aug 18, 2023
70f2af4
Updates docstrings
jasonb5 Aug 18, 2023
83155c2
Fixes horizontal notebook
jasonb5 Aug 24, 2023
cc38765
Adds additional grid test for multiple dimensions
jasonb5 Sep 1, 2023
9899f91
Updates notebooks with example of opening grid from separate file
jasonb5 Sep 1, 2023
c590924
Fixes mutable list of dimensions
jasonb5 Sep 1, 2023
9520bb7
Reuse _validate_grid_has_single_axis_dim function
jasonb5 Sep 1, 2023
f051ede
Fixes matching error
jasonb5 Sep 1, 2023
954efd2
Fixes using raw string
jasonb5 Sep 1, 2023
8002066
Deprecates horizontal_xesmf and horizontal_regrid2
jasonb5 Sep 1, 2023
7bbdba8
Handles datsets with multiple variables
jasonb5 Sep 1, 2023
57b8073
Fixes typings
jasonb5 Sep 1, 2023
5c11971
Removes storing xESMF regridder instance
jasonb5 Sep 1, 2023
6a99906
Fixes deprecation notice
jasonb5 Sep 11, 2023
3d4dce6
Fixes black formatting
jasonb5 Sep 12, 2023
4f42e49
Adds regridding FAQ entries.
jasonb5 Sep 12, 2023
ab072b2
Merge branch 'main' into fix_regrid_docs
jasonb5 Sep 12, 2023
a96b525
Apply suggestions from code review
jasonb5 Sep 14, 2023
2121557
Updates planned to exploring unstructed grid support
jasonb5 Sep 14, 2023
8828bc8
Fixes docstring formatting
jasonb5 Sep 14, 2023
c205e66
Update docs/faqs.rst
jasonb5 Sep 25, 2023
884fc8b
Adds missing docstring
jasonb5 Sep 25, 2023
4698d0d
address review comment for regridding notebooks
chengzhuzhang Sep 26, 2023
6a56cf2
fixing typos
chengzhuzhang Sep 26, 2023
a0ad524
Fixes perserving mask on input grid
jasonb5 Sep 27, 2023
d40a0bc
Merge branch 'main' into fix_regrid_docs
jasonb5 Sep 27, 2023
8ba0cac
Fixes inverted mask when using regrid2
jasonb5 Sep 27, 2023
c6d63b5
Updates horizontal notebook
jasonb5 Sep 28, 2023
7932861
Fixes sum that wasn't treating nan's correctly when masking input data
jasonb5 Sep 28, 2023
3f68df2
Updates horizontal notebook with fixed regrid2 masking
jasonb5 Sep 28, 2023
b76fb39
Apply suggestions from code review
jasonb5 Oct 3, 2023
29bdcda
Adds test for preserving mask on input
jasonb5 Oct 10, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 1 addition & 3 deletions docs/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -123,8 +123,6 @@ Methods
Dataset.temporal.climatology
Dataset.temporal.departures
Dataset.regridder.horizontal
Dataset.regridder.horizontal_xesmf
Dataset.regridder.horizontal_regrid2
Dataset.regridder.vertical

.. _dsmeth_1:
Expand Down Expand Up @@ -184,5 +182,5 @@ It is especially useful for those who are transitioning over from CDAT to xarray
- ``Dataset.temporal.departures("VAR_KEY", freq=<"season"|"month"|"day">)``, subset results for individual seasons, months, or days
- ``cdutil.SEASONALCYCLE.departures()``, ``cdutil.ANNUALCYCLE.departures()``, ``cdutil.<DJF|MAM|JJA|SON>.departures()``, ``cdutil.<JAN|FEB|...|DEC>.departures()``
* - Regrid horizontally?
- ``Dataset.regridder.horizontal_regrid2()``, ``Dataset.regridder.horizontal_xesmf()``
- ``Dataset.regridder.horizontal(tool="regrid2")``
- ``cdms2.regrid2()``
299 changes: 142 additions & 157 deletions docs/examples/regridding-horizontal.ipynb

Large diffs are not rendered by default.

326 changes: 166 additions & 160 deletions docs/examples/regridding-vertical.ipynb

Large diffs are not rendered by default.

58 changes: 57 additions & 1 deletion docs/faqs.rst
Original file line number Diff line number Diff line change
Expand Up @@ -224,4 +224,60 @@ caution. You should understand the potential implications of these workarounds.*

For more information on these options, visit the `xarray.open_mfdataset`_ documentation.

.. _`xarray.open_mfdataset`: https://xarray.pydata.org/en/stable/generated/xarray.open_mfdataset.html#xarray-open-mfdataset
.. _`xarray.open_mfdataset`: https://xarray.pydata.org/en/stable/generated/xarray.open_mfdataset.html#xarray-open-mfdataset

Regridding
----------
``xcdat`` extends and provides a uniform interface to `xESMF`_ and `xgcm`_. In addition, ``xcdat`` provides a port of the ``CDAT`` `regrid2`_ package.

Structured rectilinear and curvilinear grids are supported.

.. _`xESMF`: https://xesmf.readthedocs.io/en/stable/
jasonb5 marked this conversation as resolved.
Show resolved Hide resolved
.. _`xgcm`: https://xgcm.readthedocs.io/en/latest/
.. _`regrid2`: https://cdms.readthedocs.io/en/latest/regrid2.html

How can I retrieve the grid from a dataset?
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The :py:func:`xcdat.regridder.accessor.RegridderAccessor.grid` property is provided to extract the grid information from a dataset.

.. code-block:: python

ds = xcdat.open_dataset(...)
grid = ds.regridder.grid

How do I perform horizontal regridding?
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The :py:func:`xcdat.regridder.accessor.RegridderAccessor.horizontal` method provides access to the `xESMF`_ and `Regrid2`_ packages.

The arguments for each regridder can be found:

* :py:func:`xcdat.regridder.xesmf.XESMFRegridder`
* :py:func:`xcdat.regridder.regrid2.Regrid2Regridder`

An example of `horizontal`_ regridding can be found in the `gallery`_.

.. _`Regrid2`: generated/xcdat.regridder.regrid2.Regrid2Regridder.html
.. _`horizontal`: examples/regridding-horizontal.html
.. _`gallery`: gallery.html

How do I perform vertical regridding?
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
The :py:func:`xcdat.regridder.accessor.RegridderAccessor.vertical` method provides access to the `xgcm`_ package.

The arguments for each regridder can be found:

* :py:func:`xcdat.regridder.xgcm.XGCMRegridder`

An example of `vertical`_ regridding can be found in the `gallery`_.

.. _`vertical`: examples/regridding-vertical.html

Can ``xcdat`` automatically derive Parametric Vertical Coordinates in a dataset?
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Automatically deriving `Parametric Vertical Coordinates`_ is a planned feature for ``xcdat``.

.. _`Parametric Vertical Coordinates`: http://cfconventions.org/cf-conventions/cf-conventions.html#parametric-vertical-coordinate

Can I regrid data on unstructured grids?
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Regridding data on unstructured grids is a feature we are exploring for ``xcdat``.
22 changes: 22 additions & 0 deletions tests/fixtures.py
Original file line number Diff line number Diff line change
Expand Up @@ -405,6 +405,28 @@ def generate_lev_dataset(position="center") -> xr.Dataset:
return ds


def generate_multiple_variable_dataset(
copies: int, separate_dims: bool = False, **kwargs
) -> xr.Dataset:
ds_base = generate_dataset(**kwargs)

datasets = [ds_base]

for idx in range(copies):
ds_copy = ds_base.copy(deep=True)

var_names = list(["ts"])

if separate_dims:
var_names += list(ds_base.dims.keys()) # type: ignore[arg-type]

ds_copy = ds_copy.rename({x: f"{x}{idx+1}" for x in var_names})

datasets.append(ds_copy)

return xr.merge(datasets, compat="override")


def generate_dataset(
decode_times: bool, cf_compliant: bool, has_bounds: bool
) -> xr.Dataset:
Expand Down
83 changes: 73 additions & 10 deletions tests/test_regrid.py
Original file line number Diff line number Diff line change
Expand Up @@ -502,17 +502,14 @@ def test_regrid_input_mask(self):

expected_output = np.array(
[
[1.0, 1.0, 1.0, 1.0],
[1e20, 1e20, 1e20, 1e20],
[1e20, 1e20, 1e20, 1e20],
[1.0, 1.0, 1.0, 1.0],
[0.0, 0.0, 0.0, 0.0],
[0.70710677, 0.70710677, 0.70710677, 0.70710677],
[0.70710677, 0.70710677, 0.70710677, 0.70710677],
[0.0, 0.0, 0.0, 0.0],
],
dtype=np.float32,
)

# need to replace nans since nan != nan
output_data["ts"] = output_data.ts.fillna(1e20)

assert np.all(output_data.ts.values == expected_output)

def test_regrid_output_mask(self):
Expand Down Expand Up @@ -1182,8 +1179,56 @@ class TestAccessor:
def setup(self):
self.data = mock.MagicMock()
self.ac = accessor.RegridderAccessor(self.data)
self.horizontal_ds = fixtures.generate_dataset(
decode_times=True, cf_compliant=True, has_bounds=True
)
self.vertical_ds = fixtures.generate_lev_dataset()

def test_preserve_mask_from_input(self):
mask = xr.zeros_like(self.horizontal_ds.isel(time=0), dtype=int)
mask = mask.drop_vars("time")

mask.ts[1:3] = 1.0

self.horizontal_ds["mask"] = mask.ts

grid = accessor._get_input_grid(self.horizontal_ds, "ts", ["X", "Y"])

assert "mask" in grid

xr.testing.assert_allclose(mask.ts, grid.mask)

@requires_xesmf
def test_horizontal(self):
output_grid = grid.create_gaussian_grid(32)

output_data = self.horizontal_ds.regridder.horizontal(
"ts", output_grid, tool="xesmf", method="bilinear"
)

assert output_data.ts.shape == (15, 32, 65)

# use dataset with time dimension
output_data = self.horizontal_ds.regridder.horizontal(
"ts", self.horizontal_ds, tool="xesmf", method="bilinear"
)

assert output_data.ts.shape == (15, 4, 4)

ds_multi = fixtures.generate_multiple_variable_dataset(
1,
separate_dims=True,
decode_times=True,
cf_compliant=True,
has_bounds=True,
)

output_data = ds_multi.regridder.horizontal(
"ts", self.horizontal_ds, tool="xesmf", method="bilinear"
)

assert output_data.ts.shape == (15, 4, 4)

def test_vertical(self):
output_grid = grid.create_grid(lev=np.linspace(10000, 2000, 2))

Expand All @@ -1193,6 +1238,13 @@ def test_vertical(self):

assert output_data.so.shape == (15, 2, 4, 4)

# use dataset with time dimension
output_data = self.vertical_ds.regridder.vertical(
"so", self.vertical_ds, tool="xgcm", method="linear"
)

assert output_data.so.shape == (15, 4, 4, 4)

def test_vertical_multiple_z_axes(self):
output_grid = grid.create_grid(lev=np.linspace(10000, 2000, 2))

Expand Down Expand Up @@ -1235,6 +1287,16 @@ def test_grid(self):
assert "lat_bnds" in grid
assert "lon_bnds" in grid

ds_multi = fixtures.generate_multiple_variable_dataset(
1, separate_dims=True, decode_times=True, cf_compliant=True, has_bounds=True
)

with pytest.raises(
ValueError,
match=r".*lon\d?.*lon\d?.*",
):
ds_multi.regridder.grid

def test_grid_raises_error_when_dataset_has_multiple_dims_for_an_axis(self):
ds_bounds = fixtures.generate_dataset(
decode_times=True, cf_compliant=True, has_bounds=True
Expand All @@ -1246,7 +1308,8 @@ def test_grid_raises_error_when_dataset_has_multiple_dims_for_an_axis(self):
with pytest.raises(ValueError):
ds_bounds.regridder.grid

def test_horizontal_tool_check(self):
@mock.patch("xcdat.regridder.accessor._get_input_grid")
def test_horizontal_tool_check(self, _get_input_grid):
mock_regridder = mock.MagicMock()
mock_regridder.return_value.horizontal.return_value = "output data"

Expand All @@ -1266,8 +1329,8 @@ def test_horizontal_tool_check(self):
):
self.ac.horizontal("ts", mock_data, tool="dummy") # type: ignore

@mock.patch("xcdat.regridder.accessor._get_vertical_input_grid")
def test_vertical_tool_check(self, _get_vertical_input_grid):
@mock.patch("xcdat.regridder.accessor._get_input_grid")
def test_vertical_tool_check(self, _get_input_grid):
mock_regridder = mock.MagicMock()
mock_regridder.return_value.vertical.return_value = "output data"

Expand Down
Loading
Loading