0.2.1: fixed boundary values imposing bugs from Jax-cfd (#3)

* updated requirements and added a test workflow

* fixed boundary values imposing bugs from Jax-cfd

Author: scaomath

.github/workflows/pytest.yml

@@ -0,0 +1,40 @@
+name: Run Tests 🧪
+
+on:
+  push:
+    tags:
+      - '*'
+  pull_request:
+
+jobs:
+  test:
+    name: Run pytest
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v2
+
+    - name: Set up Python
+      uses: actions/setup-python@v2
+      with:
+        python-version: "3.10"
+
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements.txt
+
+    - name: Extract tag name
+      id: tag
+      run: echo ::set-output name=TAG_NAME::$(echo $GITHUB_REF | cut -d / -f 3)
+    
+    - name: Update version in setup.py
+      run: >-
+        sed -i "s/{{VERSION_PLACEHOLDER}}/${{ steps.tag.outputs.TAG_NAME }}/g" setup.py
+    
+    - name: Install build
+      run: python -m pip install build
+
+    - name: Run pytest
+      run: |
+        pytest --pyargs torch_cfd --verbose

.github/workflows/python-publish.yml

@@ -25,11 +25,10 @@ jobs:
       with:
         python-version: "3.10"
 
-    - name: Install pip
-      run: python -m pip install --upgrade pip
-
     - name: Install dependencies
-      run: pip install -r requirements.txt
+      run: |
+        python -m pip install --upgrade pip
+        pip install -r requirements.txt
 
     - name: Extract tag name
       id: tag

README.md

@@ -80,6 +80,6 @@ If you like to use `torch-cfd` please use the following [paper](https://arxiv.or
 I am grateful for the support from [Long Chen (UC Irvine)](https://github.com/lyc102/ifem) and 
 [Ludmil Zikatanov (Penn State)](https://github.com/HAZmathTeam/hazmath) over the years, and their efforts in open-sourcing scientific computing codes. I also appreciate the support from the National Science Foundation (NSF) to junior researchers. I want to thank the free A6000 credits at the SSE ML cluster from the University of Missouri. 
 
-(Added after `0.2.0`) I also want to acknowledge that University of Missouri's OpenAI Enterprise API key. After from version `0.1.0`, I begin prompt existing codes in VSCode Copilot (using the OpenAI Enterprise API), which arguably significantly improve the efficiency on "porting->debugging->refactoring" cycle, e.g., Copilot helps design unittests and `.vscode/launch.json` for debugging. For details of how Copilot's suggestions on code refactoring, please see [README.md](./torch_cfd/README.md) in `torch_cfd` folder.
+(Added after `0.2.0`) I also want to acknowledge that University of Missouri's OpenAI Enterprise API key. After version `0.1.0`, I began prompt in VSCode Copilot with existing codes (using the OpenAI Enterprise API), which arguably significantly improve the efficiency on "porting->debugging->refactoring" cycle, e.g., Copilot helps design unittests and `.vscode/launch.json` for debugging. For details of how Copilot's suggestions on code refactoring, please see [README.md](./torch_cfd/README.md) in `torch_cfd` folder.
 
 For individual paper's acknowledgment please see [here](./fno/README.md).

requirements.txt

@@ -1,8 +1,15 @@
-numpy>=2.2.0
-torch>=2.5.0
-xarray>=2025.3.1
-tqdm>=4.62.0
-einops>=0.8.0
-dill>=0.4.0
-matplotlib>=3.5.0
-seaborn>=0.13.0
+absl_py==2.2.2
+dill==0.4.0
+einops==0.8.1
+h5py==3.13.0
+matplotlib==3.10.3
+numpy==2.2.6
+plotly==6.0.1
+psutil==7.0.0
+pytest==8.3.5
+scipy==1.15.3
+seaborn==0.13.2
+tensordict==0.7.2
+torch==2.6.0
+tqdm==4.67.1
+xarray==2025.3.1

torch_cfd/README.md

@@ -1,20 +1,37 @@
-## TODO
+# TODO
 
-- [x] add native PyTorch implementation for applying `torch.linalg` and `torch.fft` function directly on `GridArray` and `GridVariable` (added 0.1.0).
+- [x] add native PyTorch implementation for applying `torch.linalg` and `torch.fft` function directly on `GridArray` and `GridVariable` (added 0.1.).
 - [x] add discrete Helmholtz decomposition (pressure projection) in both spatial and spectral domains (added 0.0.1).
 - [x] adjust the functions and routines to act on `(batch, time, *spatial)` tensor, currently only `(*spatial)` is supported (added for key routines in 0.0.1).
 - [x] add native FFT-based vorticity computation, instead of taking finite differences for pseudo-spectral (added in 0.0.4).
 - [ ] add no-slip boundary.
 
-## Changelog
+# Changelog
 
 ### 0.2.0
 
 After version `0.1.0`, I began prompt with existing codes in VSCode Copilot (using the OpenAI Enterprise API kindedly provided by UM), which arguably significantly improve the "porting->debugging->refactoring" cycle. I recorded some several good refactoring suggestions by GPT o4-mini and some by ***Claude Sonnet 3.7*** here. There were definitely over-complicated "poor" refactoring suggestions, which have been stashed after benchmarking. I found that Sonnet 3.7 is exceptionally good at providing templates for me to filling the details, when it is properly prompted with details of the functionality of current codes. Another highlight is that, based on the error or exception raised in the unittests, Sonnet 3.7 directly added configurations in `.vscode/launch.json`, saving me quite some time of copy-paste boilerplates then change by hand.
 
 #### Major change: batch dimension for FVM
 The finite volume solver now accepts the batch dimension, some key updates include
-- Re-implemented flux computations of $(\boldsymbol{u}\cdot\nabla)\boldsymbol{u}$ as `nn.Module`. I originally implemented a tensor only version but did not quite work. Sonnet 3.7 provided a very good refactoring template after being given both the original code and my implementation.
+- Re-implemented flux computations of $(\boldsymbol{u}\cdot\nabla)\boldsymbol{u}$ as `nn.Module`. I originally implemented a tensor only version but did not quite work by pre-assigning `target_offsets`, which was buggy for the second component of the velocity. Sonnet 3.7 provided a very good refactoring template after being given both the original code and my implementation, after which I pretty much just fill in the blanks in [`advection.py`](./advection.py). Later I found out the bug was pretty stupid on my side, from
+  ```python
+  for i in range(2):
+    u = v[i]
+    for j in range(2):
+        v[i] = flux_interpolation(u, v[j]) # offset is updated here
+  ```
+  to 
+  ```python
+  # this is gonna be buggy of course because the offset alignment will go wrong
+  # the target_offsets are looped inside flux_interpolation of AdvectionVanLeer
+  for offset in target_offsets:
+    u = v[i]; u.offset = offset
+    for j in range(2):
+        v[i] = flux_interpolation(u, v[j])
+        v[i].offset = offset 
+  ```
+  The fixed version that loops in `__call__` of [`AdvectionVanLeer` class is here](./advection.py#L451).
 - Implemented a `_constant_pad_tensor` function to improve the behavior of `F.pad`, to help imposing non-homogeneous boundary conditions. It uses naturally ordered `pad` args (like Jax, unlike `F.pad`), while taking the batch dimension into consideration.
 - Changed the behavior of `u.shift` taking into consideration of batch dimension. In general these methods within the `bc` class or `GridVariable` starts from the last dimension instead of the first, e.g., `for dim in range(u.grid.ndim): ...` changes to `for dim in range(-u.grid.ndim, 0): ...`.
 

torch_cfd/boundaries.py

@@ -34,8 +34,8 @@
 
 
 class Padding:
-    MIRROR = "mirror"
-    EXTEND = "extend"
+    MIRROR = "reflect"
+    EXTEND = "replicate"
 
 
 @dataclasses.dataclass(init=False, frozen=True)
@@ -71,30 +71,17 @@ def __init__(
     def shift(
         self,
         u: GridVariable,
-        offset: float,
+        offset: int,
         dim: int,
     ) -> GridVariable:
         """
         A fallback function to make the implementation back-compatible
         see grids.shift
+        bc.shift(u, offset, dim) overrides u.bc
+        grids.shift(u, offset, dim) keeps u.bc
         """
-        return grids.shift(u, offset, dim)
+        return grids.shift(u, offset, dim, self)
 
-    def _count_bc_components(self) -> int:
-        """Counts the number of components in the boundary conditions.
-
-        Returns:
-          The number of components in the boundary conditions.
-        """
-        count = 0
-        ndim = len(self.types)
-        for axis in range(ndim):  # ndim
-            if len(self.types[axis]) != 2:
-                raise ValueError(
-                    f"Boundary conditions for axis {axis} must have two values got {len(self.types[axis])}."
-                )
-            count += len(self.types[axis])
-        return count
 
     def _is_aligned(self, u: GridVariable, dim: int) -> bool:
         """Checks if array u contains all interior domain information.
@@ -240,15 +227,15 @@ def pad_and_impose_bc(
         """
         if offset_to_pad_to is None:
             offset_to_pad_to = u.offset
-            for axis in range(u.grid.ndim):
-                _ = self._is_aligned(u, axis)
-                if self.types[axis][0] == BCType.DIRICHLET and math.isclose(
-                    u.offset[axis], 1.0
-                ):
-                    if math.isclose(offset_to_pad_to[axis], 1.0):
-                        u = grids.pad(u, 1, axis, mode=mode)
-                    elif math.isclose(offset_to_pad_to[axis], 0.0):
-                        u = grids.pad(u, -1, axis, mode=mode)
+        for axis in range(-u.grid.ndim, 0):
+            _ = self._is_aligned(u, axis)
+            if self.types[axis][0] == BCType.DIRICHLET and math.isclose(
+                u.offset[axis], 1.0
+            ):
+                if math.isclose(offset_to_pad_to[axis], 1.0):
+                    u = grids.pad(u, 1, axis, self)
+                elif math.isclose(offset_to_pad_to[axis], 0.0):
+                    u = grids.pad(u, -1, axis, self)
         return GridVariable(u.data, u.offset, u.grid, self)
 
     def impose_bc(self, u: GridVariable) -> GridVariable:
@@ -265,7 +252,8 @@ def impose_bc(self, u: GridVariable) -> GridVariable:
         """
         offset = u.offset
         u = self.trim_boundary(u)
-        return self.pad_and_impose_bc(u, offset)
+        u = self.pad_and_impose_bc(u, offset)
+        return u
 
 
 class HomogeneousBoundaryConditions(ConstantBoundaryConditions):
@@ -383,6 +371,22 @@ def periodic_and_neumann_boundary_conditions(
         )
 
 
+def _count_bc_components(bc: BoundaryConditions) -> int:
+    """Counts the number of components in the boundary conditions.
+
+    Returns:
+        The number of components in the boundary conditions.
+    """
+    count = 0
+    ndim = len(bc.types)
+    for axis in range(ndim):  # ndim
+        if len(bc.types[axis]) != 2:
+            raise ValueError(
+                f"Boundary conditions for axis {axis} must have two values got {len(bc.types[axis])}."
+            )
+        count += len(bc.types[axis])
+    return count
+
 def consistent_boundary_conditions_grid(
     grid, *arrays: GridVariable
 ) -> Tuple[GridVariable, ...]:
@@ -391,7 +395,7 @@ def consistent_boundary_conditions_grid(
     """
     bc_counts = []
     for array in arrays:
-        bc_counts.append(array.bc._count_bc_components())
+        bc_counts.append(_count_bc_components(array.bc))
     bc_count = bc_counts[0]
     if any(bc_counts[i] != bc_count for i in range(1, len(bc_counts))):
         raise Exception("Boundary condition counts are inconsistent")

torch_cfd/fvm.py

@@ -45,14 +45,14 @@ class ProjectionExplicitODE(nn.Module):
     u <- pressure_projection(u)
     """
 
-    def explicit_terms(self, *, u):
+    def explicit_terms(self, *args, **kwargs) -> GridVariableVector:
         """
         Explicit forcing term as du/dt.
-        * allows extra arguments to be passed.
         """
         raise NotImplementedError
 
-    def pressure_projection(self, *, u):
+    def pressure_projection(self, *args, **kwargs) -> Tuple[GridVariableVector, GridVariable]:
+        """Pressure projection step."""
         raise NotImplementedError
 
     def forward(self, u: GridVariableVector, dt: float) -> GridVariableVector:

torch_cfd/grids.py

@@ -856,7 +856,7 @@ def clone(self, *args, **kwargs):
         return super().clone(*args, **kwargs)
 
 
-def shift(u: GridVariable, offset: int, dim: int) -> GridVariable:
+def shift(u: GridVariable, offset: int, dim: int, bc: Optional[BoundaryConditions] = None) -> GridVariable:
     """Shift a GridVariable by `offset`.
 
     Args:
@@ -869,12 +869,12 @@ def shift(u: GridVariable, offset: int, dim: int) -> GridVariable:
         `u.offset + offset`.
     """
     dim = -u.grid.ndim + dim if dim >= 0 else dim
-    padded = pad(u, offset, dim)
+    padded = pad(u, offset, dim, bc)
     trimmed = trim(padded, -offset, dim)
     return trimmed
 
 
-def pad(u: GridVariable, width: int, dim: int) -> GridVariable:
+def pad(u: GridVariable, width: int, dim: int, bc: Optional[BoundaryConditions] = None) -> GridVariable:
     """Pad a GridVariable by `padding`.
 
     Important: the original _pad in jax_cfd makes no sense past 1 ghost cell for nonperiodic boundaries.
@@ -892,11 +892,14 @@ def pad(u: GridVariable, width: int, dim: int) -> GridVariable:
     Note:
         the padding removes the boundary conditions, so u.bc is set to None.
     """
+    assert not (u.bc is None and bc is None), "Both u.bc and bc cannot be None"
+    bc = bc if bc is not None else u.bc
+    
     if width < 0:  # pad lower boundary
-        bc_type = u.bc.types[dim][0]
+        bc_type = bc.types[dim][0]
         padding = (-width, 0)
     else:  # pad upper boundary
-        bc_type = u.bc.types[dim][1]
+        bc_type = bc.types[dim][1]
         padding = (0, width)
 
     full_padding = [(0, 0)] * u.grid.ndim
@@ -913,20 +916,20 @@ def pad(u: GridVariable, width: int, dim: int) -> GridVariable:
     if bc_type == BCType.PERIODIC:
         # self.values are ignored here
         data = expand_dims_pad(u.data, full_padding, mode="circular")
-        return GridVariable(data, tuple(new_offset), u.grid)
+        return GridVariable(data, tuple(new_offset), u.grid, bc)
     elif bc_type == BCType.DIRICHLET:
         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
             data = 2 * expand_dims_pad(
-                u.data, full_padding, mode="constant", constant_values=u.bc._values
-            ) - expand_dims_pad(u.data, full_padding, mode="reflect")
-            return GridVariable(data, tuple(new_offset), u.grid)
+                u.data, full_padding, mode="constant", constant_values=bc._values
+            ) - expand_dims_pad(u.data, full_padding, mode="replicate")
+            return GridVariable(data, tuple(new_offset), u.grid, bc)
         elif math.isclose(u.offset[dim] % 1, 0):  # cell edge
             data = expand_dims_pad(
-                u.data, full_padding, mode="constant", constant_values=u.bc._values
+                u.data, full_padding, mode="constant", constant_values=bc._values
             )
-            return GridVariable(data, tuple(new_offset), u.grid)
+            return GridVariable(data, tuple(new_offset), u.grid, bc)
         else:
             raise ValueError(
                 "expected the new offset to be an edge or cell center, got "
@@ -943,19 +946,19 @@ def pad(u: GridVariable, width: int, dim: int) -> GridVariable:
         else:
             # When the data is cell-centered, computes the backward difference.
             # When the data is on cell edges, boundary is set such that
-            # (u_last_interior - u_boundary)/grid_step = neumann_bc_value.
+            # (u_boundary - u_last_interior)/grid_step = neumann_bc_value (fixed from Jax-cfd).
+            # note: Jax-cfd implementation was wrong, Neumann BC is \nabla u \cdot exterior normal, the order is reversed
             data = expand_dims_pad(
                 u.data, full_padding, mode="replicate"
             ) + u.grid.step[dim] * (
-                expand_dims_pad(u.data, full_padding, mode="constant")
-                - expand_dims_pad(
+                expand_dims_pad(
                     u.data,
                     full_padding,
                     mode="constant",
-                    constant_values=u.bc._values,
-                )
+                    constant_values=bc._values,
+                ) - expand_dims_pad(u.data, full_padding, mode="constant")
             )
-            return GridVariable(data, tuple(new_offset), u.grid)
+            return GridVariable(data, tuple(new_offset), u.grid, bc)
 
     else:
         raise ValueError("invalid boundary type")