Refactors regrid2 main function to work on a DataArray, removes dataset

from land-sea mask generation

Author: jasonb5

tests/test_mask.py

@@ -239,7 +239,7 @@ def test_generate_land_sea_mask_pcmdi(ds):
 
     output = mask.generate_land_sea_mask(ds["ts"], method="pcmdi")
 
-    xr.testing.assert_equal(output.lsmask, expected)
+    xr.testing.assert_equal(output, expected)
 
 
 @mock.patch("xcdat.mask._improve_mask")
@@ -248,7 +248,7 @@ def test_generate_land_sea_mask_pcmdi_multiple_iterations(_improve_mask, ds):
         [
             [1, 1, 1, 1],
             [0, 0, 0, 0],
-            [1, 1, 0, 1],
+            [1, 1, 1, 1],
             [0, 0, 0, 0],
         ],
         dims=("lat", "lon"),
@@ -258,24 +258,24 @@ def test_generate_land_sea_mask_pcmdi_multiple_iterations(_improve_mask, ds):
         [
             [1, 1, 1, 1],
             [0, 0, 0, 0],
-            [1, 1, 1, 1],
+            [1, 1, 0, 1],
             [0, 0, 0, 0],
         ],
         dims=("lat", "lon"),
         coords={"lat": ds.lat.copy(), "lon": ds.lon.copy()},
     )
 
     _improve_mask.side_effect = [
-        xr.Dataset({"sftlf": mask1.copy()}),
-        xr.Dataset({"sftlf": mask2.copy()}),
-        xr.Dataset({"sftlf": mask2.copy()}),
+        mask1.copy(),
+        mask2.copy(),
+        mask2.copy(),
     ]
 
     expected = xr.DataArray(
         [
             [1, 1, 1, 1],
             [0, 0, 0, 0],
-            [1, 1, 1, 1],
+            [1, 1, 0, 1],
             [0, 0, 0, 0],
         ],
         dims=("lat", "lon"),
@@ -284,7 +284,7 @@ def test_generate_land_sea_mask_pcmdi_multiple_iterations(_improve_mask, ds):
 
     output = mask.generate_land_sea_mask(ds["ts"], method="pcmdi")
 
-    xr.testing.assert_equal(output.lsmask, expected)
+    xr.testing.assert_equal(output, expected)
 
 
 def test_get_resource_path(monkeypatch, tmp_path):

tests/test_regrid.py

@@ -655,8 +655,10 @@ def test_regrid_create_nan_mask(self):
 
         output_data = regridder.horizontal("ts", self.coarse_2d_ds)
 
+        assert output_data.ts.dtype == np.float32
+
         # np.nan != np.nan, replace with 1e20
-        output_data = output_data.fillna(1e20)
+        output_data = output_data.fillna(1e20).astype(np.float32)
 
         expected_output = np.array(
             [
@@ -678,8 +680,10 @@ def test_regrid_input_mask(self):
 
         output_data = regridder.horizontal("ts", self.coarse_2d_ds)
 
+        assert output_data.ts.dtype == np.float32
+
         # np.nan != np.nan, replace with 1e20
-        output_data = output_data.fillna(1e20)
+        output_data = output_data.fillna(1e20).astype(np.float32)
 
         expected_output = np.array(
             [

xcdat/mask.py

@@ -11,6 +11,7 @@
 from xcdat.axis import get_dim_coords
 from xcdat.regridder.accessor import obj_to_grid_ds
 from xcdat.regridder.grid import create_grid
+from xcdat.regridder.regrid2 import _horizontal
 
 logger = _setup_custom_logger(__name__)
 
@@ -202,7 +203,6 @@ def _pcmdi_land_sea_mask(
     da: xr.DataArray,
     threshold1: float = 0.2,
     threshold2: float = 0.3,
-    mask_name: str = "lsmask",
 ) -> xr.DataArray:
     """Generate a land-sea mask using the PCMDI method.
 
@@ -233,14 +233,13 @@ def _pcmdi_land_sea_mask(
         "sftlf", obj_to_grid_ds(da), tool="regrid2"
     )
 
-    mask = highres_regrid.copy()
-    mask["sftlf"] = xr.where(highres_regrid.sftlf > 0.5, 1, 0).astype("i")
+    mask = xr.where(highres_regrid.sftlf > 0.5, 1, 0).astype("i")
 
-    lon = mask.sftlf.cf["X"]
-    lon_bnds = mask.bounds.get_bounds("X")
+    lon = mask.cf["X"]
+    lon_bnds = highres_regrid.bounds.get_bounds("X")
     is_circular = _is_circular(lon, lon_bnds)
 
-    surrounds = _generate_surrounds(mask.sftlf, is_circular)
+    surrounds = _generate_surrounds(mask, is_circular)
 
     i = 0
 
@@ -264,8 +263,6 @@ def _pcmdi_land_sea_mask(
 
         i += 1
 
-    mask = mask.rename({"sftlf": mask_name})
-
     return mask
 
 
@@ -332,14 +329,14 @@ def _is_circular(lon: xr.DataArray, lon_bnds: xr.DataArray) -> bool:
 
 
 def _improve_mask(
-    mask: xr.Dataset,
+    mask_da: xr.DataArray,
     source: xr.Dataset,
     data_var: str,
     surrounds: list[np.ndarray],
     is_circular: bool,
     threshold1=0.2,
     threshold2=0.3,
-) -> xr.Dataset:
+) -> xr.DataArray:
     """Improve a land-sea mask.
 
     This function improves a land-sea mask by converting points based on
@@ -367,15 +364,12 @@ def _improve_mask(
     xr.Dataset
         The improved mask.
     """
-    mask_approx = _map2four(
-        mask,
-        data_var,
-    )
+    mask_approx = _map2four(mask_da)
 
-    diff = source[data_var] - mask_approx[data_var]
+    diff = source[data_var] - mask_approx
 
     mask_convert_land = _convert_points(
-        mask[data_var] * 1.0,
+        mask_da * 1.0,
         source[data_var],
         diff,
         threshold1,
@@ -395,12 +389,10 @@ def _improve_mask(
         convert_land=False,
     )
 
-    mask[data_var] = mask_convert_sea.astype("i")
-
-    return mask
+    return mask_convert_sea.astype("i")
 
 
-def _map2four(mask: xr.Dataset, data_var: str) -> xr.Dataset:
+def _map2four(mask_da: xr.DataArray) -> xr.DataArray:
     """Map a mask to four subgrids and back.
 
     This function regrids a mask to four subgrids (odd-odd, odd-even,
@@ -419,9 +411,7 @@ def _map2four(mask: xr.Dataset, data_var: str) -> xr.Dataset:
     xr.Dataset
         The processed mask.
     """
-    mask_temp = mask.copy()
-
-    lat, lon = mask_temp[data_var].cf["Y"], mask_temp[data_var].cf["X"]
+    lat, lon = mask_da.cf["Y"], mask_da.cf["X"]
     lat_odd, lat_even = lat[::2], lat[1::2]
     lon_odd, lon_even = lon[::2], lon[1::2]
 
@@ -430,23 +420,21 @@ def _map2four(mask: xr.Dataset, data_var: str) -> xr.Dataset:
     even_odd = create_grid(y=lat_even, x=lon_odd, add_bounds=True)
     even_even = create_grid(y=lat_even, x=lon_even, add_bounds=True)
 
-    regrid_odd_odd = mask_temp.regridder.horizontal(data_var, odd_odd, tool="regrid2")
-    regrid_odd_even = mask_temp.regridder.horizontal(data_var, odd_even, tool="regrid2")
-    regrid_even_odd = mask_temp.regridder.horizontal(data_var, even_odd, tool="regrid2")
-    regrid_even_even = mask_temp.regridder.horizontal(
-        data_var, even_even, tool="regrid2"
+    regrid_odd_odd, _ = _horizontal(mask_da.copy(), obj_to_grid_ds(mask_da), odd_odd)
+    regrid_odd_even, _ = _horizontal(mask_da.copy(), obj_to_grid_ds(mask_da), odd_even)
+    regrid_even_odd, _ = _horizontal(mask_da.copy(), obj_to_grid_ds(mask_da), even_odd)
+    regrid_even_even, _ = _horizontal(
+        mask_da.copy(), obj_to_grid_ds(mask_da), even_even
     )
 
-    output = np.zeros(mask_temp[data_var].shape, dtype="f")
-
-    output[::2, ::2] = regrid_odd_odd[data_var].data
-    output[::2, 1::2] = regrid_odd_even[data_var].data
-    output[1::2, ::2] = regrid_even_odd[data_var].data
-    output[1::2, 1::2] = regrid_even_even[data_var].data
+    output = xr.zeros_like(mask_da, dtype="f")
 
-    mask_temp[data_var] = (mask_temp[data_var].dims, output)
+    output[::2, ::2] = regrid_odd_odd.data
+    output[::2, 1::2] = regrid_odd_even.data
+    output[1::2, ::2] = regrid_even_odd.data
+    output[1::2, 1::2] = regrid_even_even.data
 
-    return mask_temp
+    return output
 
 
 def _convert_points(