General update to regridding (#535)

* Remove tool specific methods from documentation, until decided one way or the other * Fixes adding missing vertical bounds * Updates docstrings * Fixes horizontal notebook * Adds additional grid test for multiple dimensions * Updates notebooks with example of opening grid from separate file * Fixes mutable list of dimensions * Reuse _validate_grid_has_single_axis_dim function * Fixes matching error * Fixes using raw string * Deprecates horizontal_xesmf and horizontal_regrid2 * Handles datsets with multiple variables * Fixes typings * Removes storing xESMF regridder instance * Fixes deprecation notice * Fixes black formatting * Adds regridding FAQ entries. * Apply suggestions from code review Co-authored-by: Tom Vo <tomvothecoder@gmail.com> * Updates planned to exploring unstructed grid support * Fixes docstring formatting * Update docs/faqs.rst Co-authored-by: Stephen Po-Chedley <pochedls@uw.edu> * Adds missing docstring * address review comment for regridding notebooks * fixing typos * Fixes perserving mask on input grid * Fixes inverted mask when using regrid2 * Updates horizontal notebook * Fixes sum that wasn't treating nan's correctly when masking input data * Updates horizontal notebook with fixed regrid2 masking * Apply suggestions from code review Co-authored-by: Tom Vo <tomvothecoder@gmail.com> * Adds test for preserving mask on input --------- Co-authored-by: Tom Vo <tomvothecoder@gmail.com> Co-authored-by: Stephen Po-Chedley <pochedls@uw.edu> Co-authored-by: chengzhuzhang <zhang40@llnl.gov>
xCDAT · Oct 10, 2023 · bb1fff6 · bb1fff6
1 parent c33f0f9
commit bb1fff6
Show file tree

Hide file tree

Showing 10 changed files with 690 additions and 592 deletions.
diff --git a/docs/api.rst b/docs/api.rst
@@ -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:
@@ -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()``
diff --git a/docs/examples/regridding-horizontal.ipynb b/docs/examples/regridding-horizontal.ipynb
diff --git a/docs/examples/regridding-vertical.ipynb b/docs/examples/regridding-vertical.ipynb
diff --git a/docs/faqs.rst b/docs/faqs.rst
@@ -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/
+.. _`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``.
diff --git a/tests/fixtures.py b/tests/fixtures.py
@@ -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:

diff --git a/tests/test_regrid.py b/tests/test_regrid.py
@@ -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):
@@ -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))
 
@@ -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))
 
@@ -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
@@ -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"
 
@@ -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"