fixed the behavior of impose_bc for 1.5 offset

Author: scaomath

torch_cfd/advection.py

@@ -225,7 +225,7 @@ def __init__(
         grid: Grid,
         bcs_c: Tuple[boundaries.BoundaryConditions, ...],
         bcs_v: Tuple[boundaries.BoundaryConditions, ...],
-        offsets: Tuple[Tuple[float, ...], ...] = ((1.5, 0.5), (0.5, 1.5)),
+        offsets: Tuple[Tuple[float, ...], ...] = ((1.0, 0.5), (0.5, 1.0)),
         **kwargs,
     ):
         super().__init__()
@@ -281,10 +281,9 @@ def forward(self, cs: GridVariableVector, v: GridVariableVector) -> GridVariable
         # flux's bc will become None
         flux = GridVariableVector(tuple(c * u for c, u in zip(cs, v)))
 
-        # Apply boundary conditions to flux if not periodic
-        flux = GridVariableVector(
-                tuple(bc.impose_bc(f) for f, bc in zip(flux, self.flux_bcs))
-            )
+        # wrap flux with boundary conditions to flux if not periodic
+        flux = GridVariableVector(tuple(GridVariable(f.data, offset, f.grid, bc) for f, offset, bc in zip(flux, self.offsets, self.flux_bcs)))
+        
 
         # Return negative divergence of flux
         # after taking divergence the bc becomes None
@@ -642,13 +641,14 @@ def __init__(
             boundaries.periodic_boundary_conditions(ndim=2),
             boundaries.periodic_boundary_conditions(ndim=2),
         ),
+        advect: type[nn.Module] = AdvectionVanLeer,
         limiter: Callable = van_leer_limiter,
         **kwargs,
     ):
         super().__init__()
 
         self.advect = nn.ModuleList(
-            AdvectionVanLeer(
+            advect(
                 grid=grid,
                 offset=offset,
                 bc_c=bc,

torch_cfd/fvm.py

@@ -186,6 +186,11 @@ def forward(
 
         Returns:
             Updated velocity field after one time step
+
+        Port note: 
+        - In Jax-CFD, dvdt is wrapped with the same bc with v, 
+          which does not work for inhomogeneous boundary condition.
+          see explicit_terms_with_same_bcs in jax_cfd/base/equation.py
         """
         alpha = self.params["a"]
         beta = self.params["b"]
@@ -206,6 +211,7 @@ def forward(
                 if alpha[i - 1][j] != 0:
                     u_star = u_star + dt * alpha[i - 1][j] * k[j]
 
+            u_star = wrap_field_same_bcs(u_star, u0)
             u[i], _ = equation.pressure_projection(u_star)
             k[i] = equation.explicit_terms(u[i], dt)
 
@@ -214,8 +220,10 @@ def forward(
             if beta[j] != 0:
                 u_star = u_star + dt * beta[j] * k[j]
 
+        u_star = wrap_field_same_bcs(u_star, u0)
         u_final, p = equation.pressure_projection(u_star)
 
+        u_final = wrap_field_same_bcs(u_final, u0)
         return u_final, p
 
 
@@ -309,9 +317,8 @@ def _explicit_terms(self, v, dt, **kwargs):
         dv_dt += self._diffusion(v)
         if forcing is not None:
             dv_dt += GridVariableVector(forcing(grid, v)) / density
-        dv_dt = wrap_field_same_bcs(dv_dt, v)
         if self.drag > 0.0:
-            dv_dt += -self.drag * v
+            dv_dt += -self.drag * v.array
         return dv_dt
 
     def explicit_terms(self, *args, **kwargs):

torch_cfd/grids.py

@@ -320,11 +320,7 @@ def trim_boundary(
             "trim_boundary() not implemented in BoundaryConditions base class."
         )
 
-    def impose_bc(
-        self,
-        u: GridVariable,
-        mode: Optional[str] = ""
-    ) -> GridVariable:
+    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."
@@ -537,7 +533,7 @@ def __post_init__(self):
 
     def __repr__(self) -> str:
         lines = [f"GridVariable:"]
-        lines.append(f"data tensor: \n{self.data.detach().numpy()}")
+        lines.append(f"data tensor: \n{self.data.cpu().detach().numpy()}")
         lines.append(f"data shape: {tuple(s for s in self.data.shape)}")
         lines.append(f"offset: {self.offset}")
         lines.append(f"grid shape: {self.grid.shape}")
@@ -780,7 +776,7 @@ def interior(self) -> GridVariable:
         interior_grid = self._interior_grid()
         return GridVariable(interior_array, self.offset, interior_grid)
 
-    def impose_bc(self, mode: str="") -> 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
@@ -1037,7 +1033,7 @@ def pad(
         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 = bc if bc is not None else u.bc  # use bc in priority
     bc_types = bc.types[dim] if bc_types is None else bc_types
     values = bc.bc_values if values is None else values
     if isinstance(width, int):
@@ -1086,11 +1082,16 @@ def pad(
         if math.isclose(u.offset[dim] % 1, 0.5):  # cell center
             # make the linearly interpolated value equal to the boundary by setting
             # the padded values to the negative symmetric values
+            # on the left side if u.offset is either 0.5 or 1.5, width = -1
+            # then (u - bc_val)/u.offset = (u_padded - bc_val)/(u.offset + width)
+            # here the implemenation assumes that the left pad and the right pad widths are the same when u.offset == 0.5 and only left needs to be padded when u.offset == 1.5
+
             if any(p > 1 for p in padding):
                 mode = Padding.SYMMETRIC
-            data = 2 * expand_dims_pad(
+            _alpha = 1 / u.offset[dim]
+            data = _alpha * expand_dims_pad(
                 u.data, full_padding, mode="constant", constant_values=values
-            ) - expand_dims_pad(u.data, full_padding, mode=mode)
+            ) + (1 - _alpha) * expand_dims_pad(u.data, full_padding, mode=mode)
             return GridVariable(data, tuple(new_offset), u.grid, bc)
         elif math.isclose(u.offset[dim] % 1, 0):  # cell edge
             # First the value on