Fix tests and update docstring for create_grid()

xCDAT · Mar 20, 2024 · cac4a90 · cac4a90
1 parent bebdc9f
commit cac4a90
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 51 deletions.
diff --git a/tests/test_dataset.py b/tests/test_dataset.py
@@ -76,9 +76,6 @@ def test_skips_adding_bounds(self):
         ds = generate_dataset(decode_times=True, cf_compliant=True, has_bounds=False)
         ds.to_netcdf(self.file_path)
 
-        result = open_dataset(self.file_path, add_bounds=False)
-        assert result.identical(ds)
-
         result = open_dataset(self.file_path, add_bounds=None)
         assert result.identical(ds)
 
@@ -367,9 +364,6 @@ def test_skips_adding_bounds(self):
         ds = generate_dataset(decode_times=True, cf_compliant=True, has_bounds=False)
         ds.to_netcdf(self.file_path1)
 
-        result = open_mfdataset(self.file_path1, add_bounds=False)
-        assert result.identical(ds)
-
         result = open_mfdataset(self.file_path1, add_bounds=None)
         assert result.identical(ds)
 

diff --git a/tests/test_regrid.py b/tests/test_regrid.py
@@ -37,7 +37,8 @@ class TestXGCMRegridder:
     def setup(self):
         self.ds = fixtures.generate_lev_dataset()
 
-        self.output_grid = grid.create_grid(lev=np.linspace(10000, 2000, 2))
+        z = grid.create_axis("lev", np.linspace(10000, 2000, 2), generate_bounds=False)
+        self.output_grid = grid.create_grid(z=z)
 
     def test_multiple_z_axes(self):
         self.ds = self.ds.assign_coords({"ilev": self.ds.lev.copy().rename("ilev")})
@@ -919,8 +920,8 @@ def test_gaussian_grid(self):
 
     def test_global_mean_grid(self):
         source_grid = grid.create_grid(
-            lat=np.array([-80, -40, 0, 40, 80]),
-            lon=np.array([0, 45, 90, 180, 270, 360]),
+            x=np.array([0, 45, 90, 180, 270, 360]),
+            y=np.array([-80, -40, 0, 40, 80]),
         )
 
         mean_grid = grid.create_global_mean_grid(source_grid)
@@ -1001,7 +1002,8 @@ def test_raises_error_for_global_mean_grid_if_an_axis_has_multiple_dimensions(se
 
     def test_zonal_grid(self):
         source_grid = grid.create_grid(
-            lat=np.array([-80, -40, 0, 40, 80]), lon=np.array([-160, -80, 80, 160])
+            x=np.array([-160, -80, 80, 160]),
+            y=np.array([-80, -40, 0, 40, 80]),
         )
 
         zonal_grid = grid.create_zonal_grid(source_grid)
@@ -1126,7 +1128,9 @@ def test_horizontal(self):
         assert output_data.ts.shape == (15, 4, 4)
 
     def test_vertical(self):
-        output_grid = grid.create_grid(lev=np.linspace(10000, 2000, 2))
+        z = grid.create_axis("lev", np.linspace(10000, 2000, 2), generate_bounds=False)
+
+        output_grid = grid.create_grid(z=z)
 
         output_data = self.vertical_ds.regridder.vertical(
             "so", output_grid, tool="xgcm", method="linear"
@@ -1142,7 +1146,8 @@ def test_vertical(self):
         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))
+        z = grid.create_axis("lev", np.linspace(10000, 2000, 2), generate_bounds=False)
+        output_grid = grid.create_grid(z=z)
 
         self.vertical_ds = self.vertical_ds.assign_coords(
             {"ilev": self.vertical_ds.lev.copy().rename("ilev")}
@@ -1244,22 +1249,6 @@ def test_vertical_tool_check(self, _get_input_grid):
         ):
             self.ac.vertical("ts", mock_data, tool="dummy", target_data=None)  # type: ignore
 
-    @pytest.mark.filterwarnings("ignore:.*invalid value.*divide.*:RuntimeWarning")
-    def test_convenience_methods(self):
-        ds = fixtures.generate_dataset(
-            decode_times=True, cf_compliant=False, has_bounds=True
-        )
-
-        out_grid = grid.create_gaussian_grid(32)
-
-        output_xesmf = ds.regridder.horizontal_xesmf("ts", out_grid, method="bilinear")
-
-        assert output_xesmf.ts.shape == (15, 32, 65)
-
-        output_regrid2 = ds.regridder.horizontal_regrid2("ts", out_grid)
-
-        assert output_regrid2.ts.shape == (15, 32, 65)
-
 
 class TestBase:
     def test_preserve_bounds(self):

diff --git a/xcdat/regridder/grid.py b/xcdat/regridder/grid.py
@@ -1,11 +1,10 @@
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import xarray as xr
 
 from xcdat.axis import COORD_DEFAULT_ATTRS, VAR_NAME_MAP, CFAxisKey, get_dim_coords
 from xcdat.bounds import create_bounds
-from xcdat.regridder.base import CoordOptionalBnds
 
 # First 50 zeros for the bessel function
 # Taken from https://github.com/CDAT/cdms/blob/dd41a8dd3b5bac10a4bfdf6e56f6465e11efc51d/regrid2/Src/_regridmodule.c#L3390-L3402
@@ -434,35 +433,20 @@ def create_zonal_grid(grid: xr.Dataset) -> xr.Dataset:
 
 
 def create_grid(
-    x: Optional[
-        Union[
-            xr.DataArray,
-            Tuple[xr.DataArray, Optional[xr.DataArray]],
-        ]
-    ] = None,
-    y: Optional[
-        Union[
-            xr.DataArray,
-            Tuple[xr.DataArray, Optional[xr.DataArray]],
-        ]
-    ] = None,
-    z: Optional[
-        Union[
-            xr.DataArray,
-            Tuple[xr.DataArray, Optional[xr.DataArray]],
-        ]
-    ] = None,
+    x: xr.DataArray | Tuple[xr.DataArray, Optional[xr.DataArray]] | None = None,
+    y: xr.DataArray | Tuple[xr.DataArray, Optional[xr.DataArray]] | None = None,
+    z: xr.DataArray | Tuple[xr.DataArray, Optional[xr.DataArray]] | None = None,
     attrs: Optional[Dict[str, str]] = None,
 ) -> xr.Dataset:
     """Creates a grid dataset using the specified axes.
 
     Parameters
     ----------
-    x : Optional[Union[xr.DataArray, Tuple[xr.DataArray]]]
+    x : xr.DataArray | Tuple[xr.DataArray, Optional[xr.DataArray]] | None
         Data with optional bounds to use for the "X" axis, by default None.
-    y : Optional[Union[xr.DataArray, Tuple[xr.DataArray]]]
+    y : xr.DataArray | Tuple[xr.DataArray, Optional[xr.DataArray]] | None
         Data with optional bounds to use for the "Y" axis, by default None.
-    z : Optional[Union[xr.DataArray, Tuple[xr.DataArray]]]
+    z : xr.DataArray | Tuple[xr.DataArray, Optional[xr.DataArray]] | None
         Data with optional bounds to use for the "Z" axis, by default None.
     attrs : Optional[Dict[str, str]]
         Custom attributes to be added to the generated `xr.Dataset`.