Improved corner interpolation of fields with boundary conditions (#547)

Improved corner interpolation of fields with boundary conditions

pde/fields/base.py

@@ -1372,7 +1372,7 @@ def interpolate(
         """
         if bc is not None:
             # impose boundary conditions and then interpolate using ghost cells
-            self.set_ghost_cells(bc)
+            self.set_ghost_cells(bc, set_corners=True)
             interpolator = self.make_interpolator(fill=fill, with_ghost_cells=True)
 
         else:
@@ -1502,19 +1502,23 @@ def get_boundary_values(
         return (self._data_full[i_wall] + self._data_full[i_ghost]) / 2  # type: ignore
 
     @fill_in_docstring
-    def set_ghost_cells(self, bc: BoundariesData, *, args=None) -> None:
+    def set_ghost_cells(
+        self, bc: BoundariesData, *, set_corners: bool = False, args=None
+    ) -> None:
         """set the boundary values on virtual points for all boundaries
 
         Args:
             bc (str or list or tuple or dict):
                 The boundary conditions applied to the field.
                 {ARG_BOUNDARIES}
+            set_corners (bool):
+                Determines whether the corner cells are set using interpolation
             args:
                 Additional arguments that might be supported by special boundary
                 conditions.
         """
         bcs = self.grid.get_boundary_conditions(bc, rank=self.rank)
-        bcs.set_ghost_cells(self._data_full, args=args)
+        bcs.set_ghost_cells(self._data_full, args=args, set_corners=set_corners)
 
     @property
     @abstractmethod

pde/grids/boundaries/axes.py

@@ -8,6 +8,8 @@
 
 from __future__ import annotations
 
+import itertools
+import logging
 from collections.abc import Iterator, Sequence
 from typing import Union
 
@@ -277,19 +279,49 @@ def get_mathematical_representation(self, field_name: str = "C") -> str:
 
         return "\n".join(result)
 
-    def set_ghost_cells(self, data_full: np.ndarray, *, args=None) -> None:
+    def set_ghost_cells(
+        self, data_full: np.ndarray, *, set_corners: bool = False, args=None
+    ) -> None:
         """set the ghost cells for all boundaries
 
         Args:
             data_full (:class:`~numpy.ndarray`):
                 The full field data including ghost points
+            set_corners (bool):
+                Determines whether the corner cells are set using interpolation
             args:
                 Additional arguments that might be supported by special boundary
                 conditions.
         """
         for b in self:
             b.set_ghost_cells(data_full, args=args)
 
+        if set_corners and self.grid.num_axes >= 2:
+            d = data_full  # abbreviation
+            nxt = [1, -2]  # maps 0 to 1 and -1 to -2 to obtain neighboring cells
+            if self.grid.num_axes == 2:
+                # iterate all corners
+                for i, j in itertools.product([0, -1], [0, -1]):
+                    d[..., i, j] = (d[..., nxt[i], j] + d[..., i, nxt[j]]) / 2
+
+            elif self.grid.num_axes >= 3:
+                # iterate all edges
+                for i, j in itertools.product([0, -1], [0, -1]):
+                    d[..., :, i, j] = (+d[..., :, nxt[i], j] + d[..., :, i, nxt[j]]) / 2
+                    d[..., i, :, j] = (+d[..., nxt[i], :, j] + d[..., i, :, nxt[j]]) / 2
+                    d[..., i, j, :] = (+d[..., nxt[i], j, :] + d[..., i, nxt[j], :]) / 2
+                # iterate all corners
+                for i, j, k in itertools.product(*[[0, -1]] * 3):
+                    d[..., i, j, k] = (
+                        d[..., nxt[i], j, k]
+                        + d[..., i, nxt[j], k]
+                        + d[..., i, j, nxt[k]]
+                    ) / 3
+            else:
+                logging.getLogger(self.__class__.__name__).warning(
+                    f"Can't interpolate corners for grid with {self.grid.num_axes} axes"
+                )
+
     def make_ghost_cell_setter(self) -> GhostCellSetter:
         """return function that sets the ghost cells on a full array"""
         ghost_cell_setters = tuple(b.make_ghost_cell_setter() for b in self)

tests/fields/test_scalar_fields.py

@@ -558,3 +558,27 @@ def test_field_split(decomp, rng):
         np.testing.assert_allclose(field.data, field_data)
     else:
         assert field_data is None
+
+
+def test_field_corner_interpolation_2d():
+    """test corner interpolation for a 2d field"""
+    f = ScalarField(UnitGrid([1, 1]), 0)
+    bc_x = [{"value": -1}, {"value": 2}]
+    bc_y = [{"value": -2}, {"value": 1}]
+    f.set_ghost_cells(bc=[bc_x, bc_y], set_corners=True)
+    expect = np.array([[-1.5, -2, 0], [-1, 0, 2], [0, 1, 1.5]])
+    np.testing.assert_allclose(f._data_full, 2 * expect.T)
+
+
+def test_field_corner_interpolation_3d():
+    """test corner interpolation for a 3d field"""
+    f = ScalarField(UnitGrid([1, 1, 1]), 0)
+    f.set_ghost_cells(bc=[[{"value": -3}, {"value": 3}]] * 3, set_corners=True)
+    expect = np.array(
+        [
+            [[-6, -6, -2], [-6, -6, 0], [-2, 0, 2]],
+            [[-6, -6, 0], [-6, 0, 6], [0, 6, 6]],
+            [[-2, 0, 2], [0, 6, 6], [2, 6, 6]],
+        ]
+    )
+    np.testing.assert_allclose(f._data_full, expect)