Merge pull request #8 from scaomath/0.2.6-dev

- added immersed boundary conditions.
- refactored the pseudo-spectral method codes, fixing some typing errors in initial conditions and forcings.

Author: scaomath

torch_cfd/README.md

@@ -8,6 +8,8 @@
 - [ ] rewrite `shift` and `pad` using `torch.roll`.
 - [x] support for function-valued boundary conditions.
 - [ ] change the multigrid implementation using convolution kernel.
+- [ ] add a simulation script similar to `funcutils.trajectory` but does not rely on flattening pytrees.
+- [ ] add several validation problems and tests (Taylor-Green, pressure gradient flow)
 
 # Changelog
 

torch_cfd/advection.py

@@ -31,9 +31,11 @@
 GridVariable = grids.GridVariable
 GridVariableVector = grids.GridVariableVector
 FluxInterpFn = Callable[[GridVariable, GridVariableVector, float], GridVariable]
+BoundaryConditions = boundaries.BoundaryConditions
 
 
 def default(value, d):
+    """Returns `value` if it is not None, otherwise returns `d` which is the default value."""
     return d if value is None else value
 
 
@@ -232,8 +234,8 @@ class AdvectAligned(nn.Module):
     def __init__(
         self,
         grid: Grid,
-        bcs_c: Tuple[boundaries.BoundaryConditions, ...],
-        bcs_v: Tuple[boundaries.BoundaryConditions, ...],
+        bcs_c: Tuple[BoundaryConditions, ...],
+        bcs_v: Tuple[BoundaryConditions, ...],
         offsets: Tuple[Tuple[float, ...], ...] = ((1.0, 0.5), (0.5, 1.0)),
         **kwargs,
     ):
@@ -318,7 +320,7 @@ def __init__(
         self,
         grid: Grid,
         target_offset: Tuple[float, ...] = (0.5, 0.5),
-        bc: Optional[boundaries.BoundaryConditions] = None,
+        bc: Optional[BoundaryConditions] = None,
         **kwargs,
     ):
         super().__init__()
@@ -360,22 +362,17 @@ def __init__(
         self,
         grid: Grid,
         target_offset: Tuple[float, ...],
-        low_interp: FluxInterpFn = None,
-        high_interp: FluxInterpFn = None,
+        low_interp: Optional[FluxInterpFn] = None,
+        high_interp: Optional[FluxInterpFn] = None,
         limiter: Callable = van_leer_limiter,
     ):
         super().__init__()
         self.grid = grid
-        self.low_interp = (
-            Upwind(grid, target_offset=target_offset)
-            if low_interp is None
-            else low_interp
-        )
-        self.high_interp = (
-            LaxWendroff(grid, target_offset=target_offset)
-            if high_interp is None
-            else high_interp
+        self.low_interp = default(low_interp, Upwind(grid, target_offset=target_offset))
+        self.high_interp = default(
+            high_interp, LaxWendroff(grid, target_offset=target_offset)
         )
+
         self.limiter = limiter
         self.target_offset = target_offset
 
@@ -394,28 +391,28 @@ def forward(
         Returns:
             Interpolated scalar field c to a target offset using Van Leer flux limiting, which uses a combination of high and low order methods to produce monotonic interpolation method.
         """
-        for axis, axis_offset in enumerate(self.target_offset):
+        for dim, offset in enumerate(self.target_offset):
             interpolation_offset = tuple(
                 [
-                    c_offset if i != axis else axis_offset
+                    c_offset if i != dim else offset
                     for i, c_offset in enumerate(c.offset)
                 ]
             )
             if interpolation_offset != c.offset:
-                if interpolation_offset[axis] - c.offset[axis] != 0.5:
+                if interpolation_offset[dim] - c.offset[dim] != 0.5:
                     raise NotImplementedError(
                         "Only forward interpolation to control volume faces is supported."
                     )
                 c_low = self.low_interp(c, v, dt)
                 c_high = self.high_interp(c, v, dt)
-                c_left = c.shift(-1, axis).data
-                c_right = c.shift(1, axis).data
-                c_next_right = c.shift(2, axis).data
+                c_left = c.shift(-1, dim).data
+                c_right = c.shift(1, dim).data
+                c_next_right = c.shift(2, dim).data
                 pos_r = safe_div(c - c_left, c_right - c)
                 neg_r = safe_div(c_next_right - c_right, c_right - c)
                 pos_phi = self.limiter(pos_r).data
                 neg_phi = self.limiter(neg_r).data
-                u = v[axis]
+                u = v[dim]
                 phi = torch.where(u > 0, pos_phi, neg_phi)
                 interpolated = c_low - (c_low - c_high) * phi
                 c = GridVariable(interpolated.data, interpolation_offset, c.grid, c.bc)
@@ -455,8 +452,8 @@ def __init__(
         self,
         grid: Grid,
         offset: Tuple[float, ...],
-        bc_c: boundaries.BoundaryConditions,
-        bc_v: Tuple[boundaries.BoundaryConditions, ...],
+        bc_c: Optional[BoundaryConditions] = None,
+        bc_v: Optional[Tuple[BoundaryConditions, ...]] = None,
         limiter: Optional[Callable] = None,
     ) -> None:
         super().__init__()
@@ -472,8 +469,12 @@ def __init__(
                 for _ in range(grid.ndim)
             ),
         )
+        self.bc_v = bc_v
         self.advect_aligned = AdvectAligned(
-            grid=grid, bcs_c=(bc_c, bc_c), bcs_v=bc_v, offsets=self.target_offsets
+            grid=grid,
+            bcs_c=(bc_c,) * grid.ndim,
+            bcs_v=bc_v,
+            offsets=self.target_offsets,
         )
         self._flux_interp = nn.ModuleList()  # placeholder
         self._velocity_interp = nn.ModuleList()  # placeholder
@@ -538,13 +539,8 @@ def __init__(
         self,
         grid: Grid,
         offset=(0.5, 0.5),
-        bc_c: boundaries.BoundaryConditions = boundaries.periodic_boundary_conditions(
-            ndim=2
-        ),
-        bc_v: Tuple[boundaries.BoundaryConditions, ...] = (
-            boundaries.periodic_boundary_conditions(ndim=2),
-            boundaries.periodic_boundary_conditions(ndim=2),
-        ),
+        bc_c: Optional[BoundaryConditions] = None,
+        bc_v: Optional[Tuple[BoundaryConditions, ...]] = None,
         **kwargs,
     ):
         super().__init__(grid, offset, bc_c, bc_v)
@@ -578,13 +574,8 @@ def __init__(
         self,
         grid: Grid,
         offset: Tuple[float, ...] = (0.5, 0.5),
-        bc_c: boundaries.BoundaryConditions = boundaries.periodic_boundary_conditions(
-            ndim=2
-        ),
-        bc_v: Tuple[boundaries.BoundaryConditions, ...] = (
-            boundaries.periodic_boundary_conditions(ndim=2),
-            boundaries.periodic_boundary_conditions(ndim=2),
-        ),
+        bc_c: Optional[BoundaryConditions] = None,
+        bc_v: Optional[Tuple[BoundaryConditions, ...]] = None,
         **kwargs,
     ):
         super().__init__(grid, offset, bc_c, bc_v)
@@ -617,13 +608,8 @@ def __init__(
         self,
         grid: Grid,
         offset: Tuple[float, ...] = (0.5, 0.5),
-        bc_c: boundaries.BoundaryConditions = boundaries.periodic_boundary_conditions(
-            ndim=2
-        ),
-        bc_v: Tuple[boundaries.BoundaryConditions, ...] = (
-            boundaries.periodic_boundary_conditions(ndim=2),
-            boundaries.periodic_boundary_conditions(ndim=2),
-        ),
+        bc_c: Optional[BoundaryConditions] = None,
+        bc_v: Optional[Tuple[BoundaryConditions, ...]] = None,
         limiter: Callable = van_leer_limiter,
         **kwargs,
     ):
@@ -639,7 +625,7 @@ def __init__(
 
         self._velocity_interp = nn.ModuleList(
             LinearInterpolation(grid, target_offset=offset, bc=bc)
-            for offset, bc in zip(self.target_offsets, bc_v)
+            for offset, bc in zip(self.target_offsets, self.bc_v)
         )
 
 
@@ -672,16 +658,21 @@ def __init__(
         self,
         grid: Grid,
         offsets: Tuple[Tuple[float, ...], ...] = ((1.0, 0.5), (0.5, 1.0)),
-        bcs: Tuple[boundaries.BoundaryConditions, ...] = (
-            boundaries.periodic_boundary_conditions(ndim=2),
-            boundaries.periodic_boundary_conditions(ndim=2),
-        ),
+        bcs: Optional[Tuple[BoundaryConditions, ...]] = None,
         advect: type[nn.Module] = AdvectionVanLeer,
         limiter: Callable = van_leer_limiter,
         **kwargs,
     ):
         super().__init__()
-
+        self.grid = grid
+        self.offsets = offsets
+        bcs = default(
+            bcs,
+            tuple(
+                boundaries.periodic_boundary_conditions(ndim=grid.ndim)
+                for _ in range(grid.ndim)
+            ),
+        )
         self.advect = nn.ModuleList(
             advect(
                 grid=grid,