updated padding behavior to have a fallback

Author: scaomath

torch_cfd/boundaries.py

@@ -345,15 +345,15 @@ def pad_and_impose_bc(
         self,
         u: GridVariable,
         offset_to_pad_to: Optional[Tuple[float, ...]] = None,
-        mode: Optional[str] = "extend",
+        mode: Optional[str] = "",
     ) -> GridVariable:
         """Returns GridVariable with correct boundary values.
 
         Some grid points of GridVariable might coincide with boundary, thus this function is only used with the trimmed GridVariable.
         Args:
             - u: a `GridVariable` object that specifies only scalar values on the internal nodes.
             - offset_to_pad_to: a Tuple of desired offset to pad to. Note that if the function is given just an interior array in dirichlet case, it can pad to both 0 offset and 1 offset.
-            - mode: type of padding to use in non-periodic case. Mirror mirrors the flow across the boundary. Extend extends the last well-defined value past the boundary.
+            - mode: type of padding to use in non-periodic case. Mirror mirrors the flow across the boundary. Extend extends the last well-defined value past the boundary. None means no ghost cell padding.
 
         Returns:
         A GridVariable that has correct boundary values.
@@ -364,16 +364,27 @@ def pad_and_impose_bc(
         for dim in range(-u.grid.ndim, 0):
             _ = self._is_aligned(u, dim)
             if self.types[dim][0] != BCType.PERIODIC:
-                # if the offset is either 0 or 1, u is aligned with the boundary and is defined on cell edges on one side of the boundary, if trim_boundary is called before this function.
-                # u needs to be padded on both sides
-                # if the offset is 0.5, one ghost cell is needed on each side.
-                # it will be taken care by grids.pad function automatically.
-                u = grids.pad(u, (1, 1), dim, self)
-            elif self.types[dim][0] == BCType.PERIODIC:
-                return GridVariable(u.data, u.offset, u.grid, self)
-        return u
+                if mode:
+                    # if the offset is either 0 or 1, u is aligned with the boundary and is defined on cell edges on one side of the boundary, if trim_boundary is called before this function.
+                    # u needs to be padded on both sides
+                    # if the offset is 0.5, one ghost cell is needed on each side.
+                    # it will be taken care by grids.pad function automatically.
+                    u = grids.pad(u, (1, 1), dim, self, mode=mode)
+                elif self.types[dim][0] == BCType.DIRICHLET and not mode:
+                    if math.isclose(offset_to_pad_to[dim], 1.0):
+                        u = grids.pad(u, 1, dim, self)
+                    elif math.isclose(offset_to_pad_to[dim], 0.0):
+                        u = grids.pad(u, -1, dim, self)
+                elif self.types[dim][0] == BCType.NEUMANN and not mode:
+                    if not math.isclose(offset_to_pad_to[dim], 0.5):
+                        raise ValueError("Neumann bc is not defined on edges.")
+                else:
+                    raise NotImplementedError(
+                        f"Padding for {self.types[dim][0]} boundary conditions is not implemented."
+                    )
+        return GridVariable(u.data, u.offset, u.grid, self)
 
-    def impose_bc(self, u: GridVariable) -> GridVariable:
+    def impose_bc(self, u: GridVariable, mode: str="") -> GridVariable:
         """Returns GridVariable with correct boundary condition.
 
         Some grid points of GridVariable might coincide with boundary. This ensures
@@ -382,11 +393,11 @@ def impose_bc(self, u: GridVariable) -> GridVariable:
         u: a `GridVariable` object.
 
         Returns:
-        A GridVariable that has correct boundary values with ghost cells added on the other side of DoFs living at cell center if the bc is Dirichlet or Neumann.
+        A GridVariable that has correct boundary values. If ghost_cell == True, then ghost cells are added on the other side of DoFs living at cell center if the bc is Dirichlet or Neumann.
         """
         offset = u.offset
         u = self.trim_boundary(u)
-        u = self.pad_and_impose_bc(u, offset)
+        u = self.pad_and_impose_bc(u, offset, mode)
         return u
 
 

torch_cfd/grids.py

@@ -36,6 +36,21 @@
 T = TypeVar("T")  # for GridVariable vector
 
 
+class BCType:
+    PERIODIC = "periodic"
+    DIRICHLET = "dirichlet"
+    NEUMANN = "neumann"
+    ROBIN = "robin"
+    NONE = None
+
+
+class Padding:
+    MIRROR = "reflect"
+    EXTEND = "replicate"
+    SYMMETRIC = "symmetric"
+    NONE = ""
+
+
 @dataclasses.dataclass(init=False, frozen=True)
 class Grid:
     """
@@ -112,7 +127,7 @@ def __init__(
     def __repr__(self) -> str:
         lines = [f"Grid({self.ndim}D):"]
         lines.append(f"  shape: {self.shape}")
-    
+
         for i in range(self.ndim):
             lower, upper = self.domain[i]
             step = self.step[i]
@@ -241,20 +256,6 @@ def eval_on_mesh(
         return GridVariable(fn(*self.mesh(offset)), offset, self)
 
 
-class BCType:
-    PERIODIC = "periodic"
-    DIRICHLET = "dirichlet"
-    NEUMANN = "neumann"
-    ROBIN = "robin"
-    NONE = None
-
-
-class Padding:
-    MIRROR = "reflect"
-    EXTEND = "replicate"
-    SYMMETRIC = "symmetric"
-
-
 @dataclasses.dataclass(init=False, frozen=True)
 class BoundaryConditions:
     """Base class for boundary conditions on a PDE variable.
@@ -322,12 +323,13 @@ def trim_boundary(
     def impose_bc(
         self,
         u: GridVariable,
+        mode: Optional[str] = ""
     ) -> GridVariable:
         """Impose boundary conditions on the grid variable."""
         raise NotImplementedError(
             "impose_bc() not implemented in BoundaryConditions base class."
         )
-    
+
     def pad_and_impose_bc(
         self,
         u: GridVariable,
@@ -778,15 +780,15 @@ def interior(self) -> GridVariable:
         interior_grid = self._interior_grid()
         return GridVariable(interior_array, self.offset, interior_grid)
 
-    def impose_bc(self) -> GridVariable:
+    def impose_bc(self, mode: str="") -> GridVariable:
         """Returns the GridVariable with edge BC enforced, if applicable.
 
         For GridVariables having nonperiodic BC and offset 0 or 1, there are values
         in the array data that are dependent on the boundary condition.
         impose_bc() changes these boundary values to match the prescribed BC.
         """
         assert self.bc is not None, "Boundary conditions must be set to impose BC."
-        return self.bc.impose_bc(self)
+        return self.bc.impose_bc(self, mode)
 
     def enforce_edge_bc(self, *args) -> GridVariable:
         """Returns the GridVariable with edge BC enforced, if applicable.
@@ -1030,9 +1032,11 @@ def pad(
     assert not (
         u.bc is None and bc is None and bc_types is None
     ), "u.bc, bc, and bc_types cannot be None at the same time"
-    assert mode in [Padding.MIRROR, Padding.EXTEND, Padding.SYMMETRIC], (
-    f"Padding mode must be one of ['{Padding.MIRROR}', '{Padding.EXTEND}', '{Padding.SYMMETRIC}'], got '{mode}'"
-)
+    assert mode in [
+        Padding.MIRROR,
+        Padding.EXTEND,
+        Padding.SYMMETRIC,
+    ], f"Padding mode must be one of ['{Padding.MIRROR}', '{Padding.EXTEND}', '{Padding.SYMMETRIC}'], got '{mode}'"
     bc = bc if bc is not None else u.bc
     bc_types = bc.types[dim] if bc_types is None else bc_types
     values = bc.bc_values if values is None else values
@@ -1112,9 +1116,7 @@ def pad(
                     return GridVariable(data, tuple(new_offset), u.grid, bc)
                 elif mode == Padding.MIRROR:
                     bc_padding = [(0, 0)] * u.grid.ndim
-                    bc_padding[dim] = tuple(
-                        1 if pad > 0 else 0 for pad in padding
-                    )
+                    bc_padding[dim] = tuple(1 if pad > 0 else 0 for pad in padding)
                     # subtract the padded cell
                     full_padding_past_bc = [(0, 0)] * u.grid.ndim
                     full_padding_past_bc[dim] = tuple(
@@ -1125,13 +1127,17 @@ def pad(
                         u.data, bc_padding, mode="constant", constant_values=(0, 0)
                     )
                     padding_values = list(values)
-                    padding_values[dim] = tuple([pad / 2 for pad in padding_values[dim]])
+                    padding_values[dim] = tuple(
+                        [pad / 2 for pad in padding_values[dim]]
+                    )
                     data = 2 * expand_dims_pad(
                         u.data,
                         full_padding,
                         mode="constant",
                         constant_values=padding_values,
-                    ) - expand_dims_pad(expanded_data, full_padding_past_bc, mode="reflect")
+                    ) - expand_dims_pad(
+                        expanded_data, full_padding_past_bc, mode="reflect"
+                    )
                 return GridVariable(data, tuple(new_offset), u.grid, bc)
             else:
                 raise ValueError(
@@ -1203,7 +1209,9 @@ def expand_dims_pad(
     inputs: torch.Tensor,
     pad: Sequence[Tuple[int, int]],
     mode: str = "constant",
-    constant_values: Union[float, Tuple[float, float], Sequence[Tuple[float, float]]] = 0,
+    constant_values: Union[
+        float, Tuple[float, float], Sequence[Tuple[float, float]]
+    ] = 0,
     **kwargs,
 ) -> torch.Tensor:
     """

torch_cfd/tests/test_advection.py

@@ -206,6 +206,10 @@ def test_burgers_analytical_dirichlet_convergence(
         offset,
     ):
         def _step_func(v, dt, method):
+            """
+            dt/2 is used because for Burgers equation
+            the flux is u_t + (0.5*u^2)_x = 0
+            """
             dv_dt = method(c=v[0], v=v, dt=dt) / 2
             return (bc.impose_bc(v[0].data + dt * dv_dt),)
 
@@ -225,10 +229,6 @@ def _velocity_implicit(grid, offset, u, t):
         advect = advect_van_leer(grid, offset)
 
         for _ in range(num_steps):
-            """
-            dt/2 is used because for Burgers equation
-            the flux is u_t + (0.5*u^2)_x = 0
-            """
             v = _step_func(v, dt, method=advect)
 
         expected = bc.impose_bc(

torch_cfd/tests/test_boundaries.py

@@ -24,6 +24,7 @@
 from torch_cfd import boundaries, grids, test_utils
 
 BCType = boundaries.BCType
+Padding = boundaries.Padding
 
 tensor = partial(torch.tensor, dtype=torch.float32)
 
@@ -704,7 +705,7 @@ def test_pad_and_impose_bc_1d(
         grid = grids.Grid((grid_size,))
         array = grids.GridVariable(input_data, input_offset, grid)
         bc = boundaries.ConstantBoundaryConditions(bc_types[0], bc_types[1])
-        actual = bc.pad_and_impose_bc(array, expected_offset)
+        actual = bc.pad_and_impose_bc(array, expected_offset, mode=Padding.EXTEND)
         expected = grids.GridVariable(expected_data, expected_offset, grid)
         self.assertArrayEqual(actual, expected)
 
@@ -765,7 +766,7 @@ def test_impose_bc_1d(
         grid = grids.Grid((grid_size,))
         bc = boundaries.ConstantBoundaryConditions(bc_types[0], bc_types[1])
         array = grids.GridVariable(input_data, input_offset, grid, bc)
-        actual = bc.impose_bc(array)
+        actual = bc.impose_bc(array, mode=Padding.EXTEND)
         expected = grids.GridVariable(expected_data, expected_offset, grid, bc)
         self.assertArrayEqual(actual, expected)
 
@@ -818,7 +819,7 @@ def test_impose_bc_2d_constant_boundary(
         grid = grids.Grid(input_data.shape)
         bc = boundaries.dirichlet_boundary_conditions(grid.ndim, values)
         variable = grids.GridVariable(input_data, offset, grid, bc)
-        variable = variable.impose_bc()
+        variable = variable.impose_bc(mode=Padding.EXTEND)
         expected = grids.GridVariable(expected_data, expected_offset, grid)
         self.assertArrayEqual(variable, expected)