Merge pull request #5 from scaomath/0.2.3-dev

0.2.2->0.2.3

Author: scaomath

examples/Kolmogrov2d_rk4_fvm_forced_turbulence.ipynb

@@ -2,15 +2,15 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import torch\n",
     "import torch.fft as fft\n",
     "\n",
     "from torch_cfd.grids import Grid\n",
-    "from torch_cfd.initial_conditions import vorticity_field\n",
+    "from torch_cfd.initial_conditions import filtered_vorticity_field\n",
     "from torch_cfd.forcings import KolmogorovForcing\n",
     "from torch_cfd.spectral import *\n",
     "from torch_cfd.equations import *\n",
@@ -26,7 +26,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 4,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -53,7 +53,7 @@
     "\n",
     "grid = Grid(shape=(n, n), domain=((0, diam), (0, diam)), device=device)\n",
     "\n",
-    "vort_init = vorticity_field(grid, peak_wavenumber, random_state=random_state, batch_size=batch_size).data\n",
+    "vort_init = filtered_vorticity_field(grid, peak_wavenumber, random_state=random_state, batch_size=batch_size).data\n",
     "vort_hat = fft.rfft2(vort_init).to(device)\n",
     "\n",
     "\n",

examples/Kolmogrov2d_rk4_spectral_forced_turbulence.ipynb

@@ -11,6 +11,7 @@
 
 import torch
 import torch.fft as fft
+import torch.nn.functional as F
 from torch_cfd.finite_differences import curl_2d
 from torch_cfd.forcings import KolmogorovForcing
 

examples/Lid-driven_cavity_rk4_fvm.ipynb

@@ -11,9 +11,10 @@
 
 import torch
 import torch.fft as fft
+import torch.nn.functional as F
 
 from torch_cfd.grids import Grid
-from torch_cfd.initial_conditions import vorticity_field
+from torch_cfd.initial_conditions import filtered_vorticity_field
 from torch_cfd.equations import *
 
 from solvers import get_trajectory_imex
@@ -131,7 +132,7 @@ def main(args):
 
         vort_init = torch.stack(
             [
-                vorticity_field(
+                filtered_vorticity_field(
                     grid, peak_wavenumber, random_state=random_state + idx + k
                 ).data
                 for k in range(batch_size)

fno/data_gen/data_gen_Kolmogorov2d.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,