update a lid-driven cavity example

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

@@ -933,9 +933,9 @@ def data(self, values: Tuple[torch.Tensor, ...]):
             v.data = value
 
     @property
-    def array(self) -> Tuple[GridVariable, ...]:
+    def array(self) -> GridVariableVector:
         """Returns a tuple of GridVariables without boundary conditions."""
-        return tuple(v.array for v in self)
+        return GridVariableVector(tuple(v.array for v in self))
 
     @array.setter
     def array(self, values: Tuple[GridVariable, ...]):