rename sfno folder to fno to avoid ModuleNotFoundError

Author: scaomath

README.md

@@ -12,7 +12,7 @@ The main changes are documented in the `README.md` under the [`torch_cfd` direct
   - All ops take into consideration the batch dimension of tensors `(b, *, n, m)` regardless of `*` dimension, for example, `(b, T, C, n, n)`, which is similar to PyTorch behavior, not a single trajectory like Google's original Jax-CFD package.
 
 ### Part II: Spectral-Refiner: Neural Operator-Assisted Navier-Stokes Equations simulator.
-  - The **Spatiotemporal Fourier Neural Operator** (SFNO) is a spacetime tensor-to-tensor learner (or trajectory-to-trajectory), available in the [`sfno` directory](./sfno/). We draw inspiration from the [3D FNO in Nvidia's Neural Operator repo](https://github.com/neuraloperator/neuraloperator), as well as Temam's book on functional analysis for NSE. Major architectural changes can be found in [the documentation of the `SFNO` class](./sfno/sfno.py#L547). 
+  - The **Spatiotemporal Fourier Neural Operator** (SFNO) is a spacetime tensor-to-tensor learner (or trajectory-to-trajectory), available in the [`fno` directory](./fno). Different components of FNO have been re-implemented keeping the conciseness of the original implementation while allowing modern expansions. We draw inspiration from the [3D FNO in Nvidia's Neural Operator repo](https://github.com/neuraloperator/neuraloperator), [Transformers-based neural operators](https://github.com/thuml/Neural-Solver-Library), as well as Temam's book on functional analysis for NSE. Major architectural changes can be found in [the documentation of the `SFNO` class](./fno/sfno.py#L485). 
   - Data generation for the meta-example of the isotropic turbulence in [McWilliams1984]. After the warmup phase, the energy spectra match the inverse cascade of Kolmogorov flow in a periodic box.
   - Pipelines for the *a posteriori* error estimation to fine-tune the SFNO to reach the scientific computing level of accuracy ($\le 10^{-6}$) in Bochner norm using FLOPs on par with a single evaluation, and only a fraction of FLOPs of a single `.backward()`.
   - [Examples](#examples) can be found below.
@@ -33,7 +33,7 @@ pip install -r requirements.txt
 
 ## Data
 The data are available at [https://huggingface.co/datasets/scaomath/navier-stokes-dataset](https://huggingface.co/datasets/scaomath/navier-stokes-dataset).
-Data generation instructions are available in the [SFNO folder](./sfno/).
+Data generation instructions are available in the [SFNO folder](./fno).
 
 
 ## Examples
@@ -75,4 +75,4 @@ 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 also want to thank the free A6000 credits at the SSE ML cluster from the University of Missouri.
 
-For individual paper's acknowledgement please see [here](./sfno/README.md).
+For individual paper's acknowledgement please see [here](./fno/README.md).

examples/ex2_FNO3d_train_normalized.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 2,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -32,12 +32,12 @@
     "\n",
     "import torch\n",
     "\n",
-    "from sfno.pipeline import *\n",
-    "from sfno.datasets import *\n",
-    "from sfno.losses import SobolevLoss\n",
+    "from fno.pipeline import *\n",
+    "from fno.datasets import *\n",
+    "from fno.losses import SobolevLoss\n",
     "\n",
-    "from sfno.fno3d import FNO3d\n",
-    "from sfno.visualizations import plot_contour_trajectory\n",
+    "from fno.fno3d import FNO3d\n",
+    "from fno.visualizations import plot_contour_trajectory\n",
     "from torch.utils.data import DataLoader\n",
     "\n",
     "from fno.utils import get_seed\n",

examples/ex2_SFNO_5ep_spectra.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -11,11 +11,11 @@
     "\n",
     "import torch\n",
     "\n",
-    "from sfno.utils import get_num_params,  get_seed\n",
-    "from sfno.pipeline import *\n",
-    "from sfno.losses import SobolevLoss\n",
-    "from sfno.datasets import BochnerDataset\n",
-    "from sfno.sfno import SFNO\n",
+    "from fno.utils import get_num_params,  get_seed\n",
+    "from fno.pipeline import *\n",
+    "from fno.losses import SobolevLoss\n",
+    "from fno.datasets import BochnerDataset\n",
+    "from fno.sfno import SFNO\n",
     "from torch.utils.data import DataLoader\n",
     "\n",
     "get_seed(42, printout=False)\n",

examples/ex2_SFNO_finetune_McWilliams2d.ipynb

@@ -144,7 +144,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {
@@ -163,14 +163,14 @@
     "vort_out_gt = vort_out_gt.unsqueeze(0).to(device)\n",
     "\n",
     "model.eval()\n",
-    "model.add_latent_hook(\"r\")\n",
+    "model.add_latent_hook(\"reduction\")\n",
     "with torch.no_grad():\n",
     "    preds = model(vort_inp_gt)\n",
     "    preds_no = preds.detach()\n",
     "print(f\"{loss_func(preds_no, vort_out_gt).item():.7e}\")\n",
     "\n",
     "\n",
-    "v_latent = model.latent_tensors[\"r\"]"
+    "v_latent = model.latent_tensors[\"reduction\"]"
    ]
   },
   {

examples/ex2_SFNO_train.ipynb

@@ -2,20 +2,20 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 5,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import os\n",
     "import numpy as np\n",
     "\n",
     "import torch\n",
-    "from sfno.utils import get_num_params, get_seed\n",
-    "from sfno.pipeline import *\n",
-    "from sfno.datasets import *\n",
-    "from sfno.losses import *\n",
-    "from sfno.sfno import SFNO\n",
-    "from sfno.visualizations import plot_contour_trajectory\n",
+    "from fno.utils import get_num_params, get_seed\n",
+    "from fno.pipeline import *\n",
+    "from fno.datasets import *\n",
+    "from fno.losses import *\n",
+    "from fno.sfno import SFNO\n",
+    "from fno.visualizations import plot_contour_trajectory\n",
     "from torch.utils.data import DataLoader\n",
     "get_seed(1127825)\n",
     "\n",

examples/ex2_SFNO_train_fnodata.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
@@ -11,12 +11,12 @@
     "import numpy as np\n",
     "\n",
     "import torch\n",
-    "from sfno.utils import get_num_params, get_seed\n",
-    "from sfno.pipeline import *\n",
-    "from sfno.datasets import *\n",
-    "from sfno.losses import *\n",
-    "from sfno.sfno import SFNO\n",
-    "from sfno.visualizations import plot_contour_trajectory\n",
+    "from fno.utils import get_num_params, get_seed\n",
+    "from fno.pipeline import *\n",
+    "from fno.datasets import *\n",
+    "from fno.losses import *\n",
+    "from fno.sfno import SFNO\n",
+    "from fno.visualizations import plot_contour_trajectory\n",
     "from torch.utils.data import DataLoader\n",
     "\n",
     "get_seed(1127825)\n",

sfno/README.md renamed to fno/README.md

@@ -1,4 +1,5 @@
-# Spatiotemporal Fourier Neural Operator
+# Spatiotemporal Fourier Neural Operator (SFNO)
+This is a new concise implementation of the Fourier Neural Operator see [`base.py`](./base.py#L172) for a template class.
 
 ## Learning maps between Bochner spaces
 SFNO now can learn a `trajectory-to-trajectory` map that inputs arbitrary-length trajectory, and outputs arbitrary-lengthed trajectory (if length is not specified, then the output length is the same with the input).
@@ -37,6 +38,13 @@ Generate the isotropic turbulence in [1] with the inverse cascade frequency sign
 
 ## Training and evaluation scripts
 
+
+### Testing the arbitrary input and output discretization sizes (including time)
+Run the part below `__name__ == "__main__"` in [`sfno.py`](sfno.py)
+```bash
+>>> python fno/sfno.py
+```
+
 ### FNO NSE dataset
 Train SFNO for the FNO dataset:
 ```bash