added test_advection for analytical solutions (#4)

Author: scaomath

examples/Kolmogrov2d_rk4_fvm_forced_turbulence.ipynb

@@ -90,7 +90,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -115,7 +115,7 @@
     "    drag=0.1,\n",
     "    smooth=True,\n",
     "    forcing_fn=forcing_fn,\n",
-    "    solver=RK4CrankNicolsonStepper(),\n",
+    "    step_fn=RK4CrankNicolsonStepper(),\n",
     ").to(device)\n",
     "\n",
     "result = get_trajectory_imex(\n",

examples/Kolmogrov2d_rk4_spectral_forced_turbulence.ipynb

@@ -122,7 +122,7 @@ def main(args):
         drag=0.1,
         smooth=True,
         forcing_fn=forcing_fn,
-        solver=RK4CrankNicolsonStepper(),
+        step_fn=RK4CrankNicolsonStepper(),
     ).to(device)
 
     num_batches = total_samples // batch_size

fno/data_gen/data_gen_Kolmogorov2d.py

@@ -119,7 +119,7 @@ def main(args):
         drag=0,
         smooth=True,
         forcing_fn=None,
-        solver=RK4CrankNicolsonStepper(),
+        step_fn=RK4CrankNicolsonStepper(),
     ).to(device)
 
     num_batches = total_samples // batch_size

fno/data_gen/data_gen_McWilliams2d.py

@@ -176,8 +176,7 @@ def main(args):
         grid=grid,
         smooth=True,
         forcing_fn=forcing_fn,
-        solver=step_fn,
-        order=2,
+        step_fn=step_fn,
     ).to(device)
 
     if os.path.exists(data_filepath) and not force_rerun:

fno/data_gen/data_gen_fno.py

@@ -34,7 +34,7 @@ def __init__(
         n=128,
         alpha=2,
         tau=3,
-        device="cuda",
+        device:torch.device=device,
         dtype=torch.float,
         normalize=False,
         smoothing=False,