fixed a mis-aligned args in SFNO

Author: scaomath

README.md

@@ -1,5 +1,7 @@
 # Neural Operator-Assisted Computational Fluid Dynamics in PyTorch
 
+![A decaying turbulence (McWilliams 1984)](examples/McWilliams2d.svg)
+
 ## Summary 
 
 This repository contains mainly two parts:

examples/Kolmogrov2d_rk4_cn_forced_turbulence.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -14,7 +14,7 @@
     "from torch_cfd.initial_conditions import *\n",
     "from torch_cfd.finite_differences import *\n",
     "from torch_cfd.forcings import *\n",
-    "from sfno.data_gen import get_trajectory_rk4\n",
+    "from fno.data_gen import get_trajectory_rk4\n",
     "\n",
     "import xarray\n",
     "import seaborn as sns\n",

examples/McWilliams2d.svg

@@ -128,7 +128,7 @@
    ],
    "source": [
     "torch.cuda.empty_cache()\n",
-    "model = SFNO(modes, modes, modes_t, width, beta).to(device)\n",
+    "model = SFNO(modes, modes, modes_t, width, beta=beta).to(device)\n",
     "print(get_num_params(model))\n",
     "\n",
     "model.load_state_dict(torch.load(path_model))\n",
@@ -211,7 +211,7 @@
    "outputs": [],
    "source": [
     "torch.cuda.empty_cache()\n",
-    "model = SFNO(modes, modes, modes_t, width, beta).to(device)\n",
+    "model = SFNO(modes, modes, modes_t, width, beta=beta).to(device)\n",
     "model.load_state_dict(torch.load(path_model))\n",
     "model.to(dtype)\n",
     "f = torch.zeros((n, n))[None, ...].to(device)"

examples/ex2_SFNO_finetune_McWilliams2d.ipynb

@@ -109,7 +109,7 @@
    ],
    "source": [
     "torch.cuda.empty_cache()\n",
-    "model = SFNO(modes, modes, modes_t, width, beta, output_steps=T_out).to(device)\n",
+    "model = SFNO(modes, modes, modes_t, width, beta=beta, output_steps=T_out).to(device)\n",
     "print(get_num_params(model))\n",
     "\n",
     "model.load_state_dict(torch.load(path_model))\n",

examples/ex2_SFNO_finetune_fnodata.ipynb

@@ -95,7 +95,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -108,7 +108,7 @@
    ],
    "source": [
     "torch.cuda.empty_cache()\n",
-    "model = SFNO(modes, modes, modes_t, width, beta).to(device)\n",
+    "model = SFNO(modes, modes, modes_t, width, beta=beta).to(device)\n",
     "\n",
     "print(get_num_params(model))\n",
     "\n",

examples/ex2_SFNO_train.ipynb

@@ -110,7 +110,7 @@
    ],
    "source": [
     "torch.cuda.empty_cache()\n",
-    "model = SFNO(modes, modes, modes_t, width, beta, \n",
+    "model = SFNO(modes, modes, modes_t, width, beta=beta, \n",
     "             output_steps=T_out).to(device)\n",
     "\n",
     "print(get_num_params(model))\n",

examples/ex2_SFNO_train_fnodata.ipynb

@@ -20,12 +20,13 @@
 from utils import *
 from pipeline import *
 from data_gen import *
+import matplotlib.pyplot as plt
 from datasets import BochnerDataset
 from losses import SobolevLoss
-import matplotlib.pyplot as plt
-from fno.sfno import SFNO
 from torch.utils.data import DataLoader
 
+from fno.sfno import SFNO
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
@@ -44,17 +45,16 @@
 
 
 def main(args):
-    
+
     current_time = datetime.now().strftime("%d_%b_%Y_%Hh%Mm")
     log_name = "".join(os.path.basename(__file__).split(".")[:-1])
 
     log_filename = os.path.join(LOG_PATH, f"{current_time}_{log_name}.log")
     logger = get_logger(log_filename)
     logger.info(f"Saving log at {log_filename}")
 
-
     all_args = {k: v for k, v in vars(args).items() if not callable(v)}
-    logger.info("Arguments: "+" | ".join(f"{k}={v}" for k, v in all_args.items()))
+    logger.info("Arguments: " + " | ".join(f"{k}={v}" for k, v in all_args.items()))
 
     example = args.example
     Ntrain = args.num_samples
@@ -122,14 +122,20 @@ def main(args):
         val_loader = DataLoader(val_dataset, batch_size=batch_size, shuffle=False)
 
         torch.cuda.empty_cache()
-        model = SFNO(modes, modes, modes_t, width, beta,
-                     num_spectral_layers=num_layers, 
-                     output_steps=out_steps,
-                     spatial_padding=spatial_padding,
-                     activation=activation,
-                     pe_trainable=pe_trainable,
-                     spatial_random_feats=spatial_random_feats,
-                     lift_activation=lift_activation)
+        model = SFNO(
+            modes,
+            modes,
+            modes_t,
+            width,
+            beta=beta,
+            num_spectral_layers=num_layers,
+            output_steps=out_steps,
+            spatial_padding=spatial_padding,
+            activation=activation,
+            pe_trainable=pe_trainable,
+            spatial_random_feats=spatial_random_feats,
+            lift_activation=lift_activation,
+        )
         logger.info(f"Number of parameters: {get_num_params(model)}")
         model.to(device)
 
@@ -155,7 +161,9 @@ def main(args):
             with tqdm(train_loader) as pbar:
                 t_ep = datetime.now().strftime("%d-%b-%Y %H:%M:%S")
                 tr_loss_str = f"current train rel L2: 0.0"
-                pbar.set_description(f"{t_ep} - Epoch [{ep+1:3d}/{epochs}]  {tr_loss_str:>35}")
+                pbar.set_description(
+                    f"{t_ep} - Epoch [{ep+1:3d}/{epochs}]  {tr_loss_str:>35}"
+                )
                 for i, data in enumerate(train_loader):
                     l2 = train_batch_ns(
                         model,
@@ -173,7 +181,9 @@ def main(args):
 
                     if i % 4 == 0:
                         tr_loss_str = f"current train rel L2: {l2.item():.4e}"
-                        pbar.set_description(f"{t_ep} - Epoch [{ep+1:3d}/{epochs}]  {tr_loss_str:>35}")
+                        pbar.set_description(
+                            f"{t_ep} - Epoch [{ep+1:3d}/{epochs}]  {tr_loss_str:>35}"
+                        )
                         pbar.update(4)
             val_l2_min = 1e4
             val_l2 = eval_epoch_ns(
@@ -214,13 +224,19 @@ def main(args):
         )
         test_loader = DataLoader(test_dataset, batch_size=1, shuffle=False)
         torch.cuda.empty_cache()
-        model = SFNO(modes, modes, modes_t, width, beta,
-                    num_spectral_layers=num_layers, 
-                    spatial_padding=spatial_padding,
-                    activation=activation,
-                    pe_trainable=pe_trainable,
-                    spatial_random_feats=spatial_random_feats,
-                    lift_activation=lift_activation).to(device)
+        model = SFNO(
+            modes,
+            modes,
+            modes_t,
+            width,
+            beta=beta,
+            num_spectral_layers=num_layers,
+            spatial_padding=spatial_padding,
+            activation=activation,
+            pe_trainable=pe_trainable,
+            spatial_random_feats=spatial_random_feats,
+            lift_activation=lift_activation,
+        ).to(device)
         model.load_state_dict(torch.load(path_model))
         logger.info(f"Loaded model from {path_model}")
         eval_metric = SobolevLoss(n_grid=n_test, norm_order=norm_order, relative=True)
@@ -238,20 +254,21 @@ def main(args):
         if args.demo_plots > 0:
             try:
                 from visualizations import plot_contour_trajectory
+
                 idx = np.random.randint(0, args.num_test_samples)
                 im1 = plot_contour_trajectory(
                     preds[idx],
                     num_snapshots=args.demo_plots,
                     T_start=args.time_warmup,
                     dt=args.dt,
-                    title="SFNO predictions"
+                    title="SFNO predictions",
                 )
                 im2 = plot_contour_trajectory(
                     gt_solns[idx],
                     num_snapshots=args.demo_plots,
                     T_start=args.time_warmup,
                     dt=args.dt,
-                    title="Ground truth generated by IMEX"
+                    title="Ground truth generated by IMEX",
                 )
                 plt.show()
             except Exception as e:
@@ -294,4 +311,4 @@ def main(args):
     parser.add_argument("--demo-plots", type=int, default=0)
 
     args = parser.parse_args()
-    main(args)
+    main(args)