Repository: xgcm/xgcm
Pull request #456 (2022-03-30 18:21:47)

Enable boundary kwargs in ufunc decorator
Fixes a bug I found whilst writing examples for #431.

 - [x] Tests added
 - [x] Passes `pre-commit run --all-files`


Patch:
diff --git a/xgcm/grid.py b/xgcm/grid.py
index 2353c96..a5e2e92 100644
--- a/xgcm/grid.py
+++ b/xgcm/grid.py
@@ -1647,7 +1647,7 @@ class Grid:
             boundary = {ax_name: boundary for ax_name in self.axes.keys()}
         if fill_value is None:
             fill_value = 0.0
-        if isinstance(fill_value, float):
+        if isinstance(fill_value, (int, float)):
             fill_value = {ax_name: fill_value for ax_name in self.axes.keys()}
         # xgcm uses a default fill value of 0, while xarray uses nan.
         fill_value = {k: 0.0 if v is None else v for k, v in fill_value.items()}
diff --git a/xgcm/grid_ufunc.py b/xgcm/grid_ufunc.py
index 105e41a..00c1f0d 100644
--- a/xgcm/grid_ufunc.py
+++ b/xgcm/grid_ufunc.py
@@ -285,6 +285,22 @@ class GridUFunc:
     boundary_width : Dict[str: Tuple[int, int], optional
         The widths of the boundaries at the edge of each array.
         Supplied in a mapping of the form {axis_name: (lower_width, upper_width)}.
+    boundary : {None, 'fill', 'extend', 'extrapolate', dict}, optional
+        A flag indicating how to handle boundaries:
+        * None: Do not apply any boundary conditions. Raise an error if
+          boundary conditions are required for the operation.
+        * 'fill':  Set values outside the array boundary to fill_value
+          (i.e. a Dirichlet boundary condition.)
+        * 'extend': Set values outside the array to the nearest array
+          value. (i.e. a limited form of Neumann boundary condition.)
+        * 'extrapolate': Set values by extrapolating linearly from the two
+          points nearest to the edge
+        Optionally a dict mapping axis name to separate values for each axis
+        can be passed.
+    fill_value : {float, dict}, optional
+        The value to use in boundary conditions with `boundary='fill'`.
+        Optionally a dict mapping axis name to separate values for each axis
+        can be passed. Default is 0.
     dask : {"forbidden", "allowed", "parallelized"}, default: "forbidden"
         How to handle applying to objects containing lazy data in the form of
         dask arrays. Passed directly on to `xarray.apply_ufunc`.
@@ -316,6 +332,8 @@ class GridUFunc:
     ufunc: Callable
     signature: _GridUFuncSignature
     boundary_width: Optional[Mapping[str, Tuple[int, int]]]
+    boundary: Optional[Union[str, Mapping[str, str]]]
+    fill_value: Optional[Union[float, Mapping[str, float]]]
     dask: Literal["forbidden", "parallelized", "allowed"]
     map_overlap: bool
 
@@ -325,10 +343,14 @@ class GridUFunc:
         str_sig = kwargs.pop("signature")
         self.signature = self._get_signature_from_str_or_type_hints(ufunc, str_sig)
         self.boundary_width = kwargs.pop("boundary_width", None)
+        self.boundary = kwargs.pop("boundary", None)
+        self.fill_value = kwargs.pop("fill_value", None)
         self.dask = kwargs.pop("dask", "forbidden")
         self.map_overlap = kwargs.pop("map_overlap", False)
         if kwargs:
-            raise TypeError("Unsupported keyword argument(s) provided")
+            raise TypeError(
+                f"Unsupported keyword argument(s) provided: {list(kwargs.keys())}"
+            )
 
     @staticmethod
     def _get_signature_from_str_or_type_hints(
@@ -379,7 +401,7 @@ class GridUFunc:
         axis: Sequence[str],
         **kwargs,
     ):
-        boundary = kwargs.pop("boundary", None)
+        boundary = kwargs.pop("boundary", self.boundary)
         dask = kwargs.pop("dask", self.dask)
         map_overlap = kwargs.pop("map_overlap", self.map_overlap)
         return apply_as_grid_ufunc(
@@ -415,6 +437,22 @@ def as_grid_ufunc(
     boundary_width : Dict[str: Tuple[int, int], optional
         The widths of the boundaries at the edge of each array.
         Supplied in a mapping of the form {axis_name: (lower_width, upper_width)}.
+    boundary : {None, 'fill', 'extend', 'extrapolate', dict}, optional
+        A flag indicating how to handle boundaries:
+        * None: Do not apply any boundary conditions. Raise an error if
+          boundary conditions are required for the operation.
+        * 'fill':  Set values outside the array boundary to fill_value
+          (i.e. a Dirichlet boundary condition.)
+        * 'extend': Set values outside the array to the nearest array
+          value. (i.e. a limited form of Neumann boundary condition.)
+        * 'extrapolate': Set values by extrapolating linearly from the two
+          points nearest to the edge
+        Optionally a dict mapping axis name to separate values for each axis
+        can be passed.
+    fill_value : {float, dict}, optional
+        The value to use in boundary conditions with `boundary='fill'`.
+        Optionally a dict mapping axis name to separate values for each axis
+        can be passed. Default is 0.
     dask : {"forbidden", "allowed", "parallelized"}, default: "forbidden"
         How to handle applying to objects containing lazy data in the form of
         dask arrays. Passed directly on to `xarray.apply_ufunc`.
@@ -442,11 +480,15 @@ def as_grid_ufunc(
     Grid.apply_as_grid_ufunc
     """
     _allowedkwargs = {
+        "boundary",
+        "fill_value",
         "dask",
         "map_overlap",
     }
     if kwargs.keys() - _allowedkwargs:
-        raise TypeError("Unsupported keyword argument(s) provided")
+        raise TypeError(
+            f"Unsupported keyword argument(s) provided: {list(kwargs.keys())}"
+        )
 
     def _as_grid_ufunc(ufunc):
         return GridUFunc(


Test patch:
diff --git a/xgcm/test/test_grid_ufunc.py b/xgcm/test/test_grid_ufunc.py
index 37a7088..f6a2c3a 100644
--- a/xgcm/test/test_grid_ufunc.py
+++ b/xgcm/test/test_grid_ufunc.py
@@ -965,6 +965,39 @@ class TestDaskOverlap:
             )
 
 
+class TestBoundary:
+    def test_boundary_constant(self):
+        def interp(a):
+            return 0.5 * (a[..., :-1] + a[..., 1:])
+
+        @as_grid_ufunc(
+            signature="(X:center)->(X:left)",
+            boundary_width={"X": (1, 0)},
+            boundary="fill",
+            fill_value=0,
+        )
+        def interp_center_to_left(a):
+            return interp(a)
+
+        grid = create_1d_test_grid("lat")
+        arr = np.arange(9)
+        da = grid._ds.lat_c.copy(data=arr)
+
+        # test that bound kwargs are used
+        result = interp_center_to_left(grid, da, axis=[["lat"]])
+        interped_arr_padded_with_zero = interp(np.concatenate([[0], arr]))
+        expected = grid._ds.lat_g.copy(data=interped_arr_padded_with_zero)
+        assert_equal(result, expected)
+
+        # test that bound kwargs can be overridden at call time
+        result = interp_center_to_left(
+            grid, da, axis=[["lat"]], boundary="fill", fill_value=1
+        )
+        interped_arr_padded_with_one = interp(np.concatenate([[1], arr]))
+        expected = grid._ds.lat_g.copy(data=interped_arr_padded_with_one)
+        assert_equal(result, expected)
+
+
 # TODO tests for handling dask in gri.diff etc. should eventually live in test_grid.py
 class TestMapOverlapGridops:
     def test_chunked_core_dims_unchanging_chunksize_center_to_right(self):