feat: add G-FNO (Group Equivariant Fourier Neural Operator) model

examples/G-FNO/README.md

@@ -0,0 +1,130 @@
+# G-FNO
+
+## 1. Background
+
+G-FNO (Group Equivariant Fourier Neural Operators) is a family of operator-learning models for PDE surrogate modeling. This PaddleCFD integration keeps the Paddle version of the main 2D/3D FNO, GCNN, GFNO, Ghybrid, and radialNO variants.
+
+![G-FNO network](assets/network_visual.png)
+
+## 2. Code Layout
+
+- Core model code: `ppcfd/models/g_fno`
+- Training and data generation scripts: `examples/G-FNO`
+
+## 3. Installation
+
+At the PaddleCFD repository root:
+
+```bash
+python -m pip install -r requirements.txt
+python -m pip install -e .
+```
+
+No extra Python packages beyond PaddleCFD root requirements are needed for the Paddle G-FNO runtime.
+
+## 4. Import Models From Installed PaddleCFD
+
+```python
+from ppcfd.models.g_fno import FNO2d, GFNO2d
+
+fno = FNO2d(
+    num_channels=1,
+    modes1=12,
+    modes2=12,
+    width=20,
+    initial_step=10,
+    grid_type="symmetric",
+)
+
+gfno = GFNO2d(
+    num_channels=1,
+    modes=12,
+    width=10,
+    initial_step=10,
+    reflection=False,
+    grid_type="symmetric",
+)
+```
+
+## 5. Data Preparation
+
+### 5.1 Navier-Stokes with Symmetric Forcing
+
+From `examples/G-FNO/data_generation/navier_stokes`:
+
+```bash
+python "ns_2d_rt.py" --nu=1e-4 --T=30 --N=1200 --save_path="./data" --ntest=100 --period=4 --device=auto
+```
+
+### 5.2 Other datasets
+
+- NS / PDEArena / PDEBench shallow-water datasets still follow the original upstream data sources referenced by the original paper.
+- Place datasets under a local `data/` directory and pass absolute or repository-relative paths to `experiments.py`.
+
+## 6. Training From Examples
+
+### Supported `--model_type`
+
+The Paddle version currently supports the model types below.
+
+| Status            | Model types                                                                                                                                                                                                                                                   |
+| ----------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| Supported         | `FNO2d`, `FNO2d_aug`, `FNO3d`, `FNO3d_aug`, `GCNN2d_p4`, `GCNN2d_p4m`, `GCNN3d_p4`, `GCNN3d_p4m`, `GFNO2d_p4`, `GFNO2d_p4m`, `GFNO3d_p4`, `GFNO3d_p4m`, `Ghybrid2d_p4`, `Ghybrid2d_p4m`, `radialNO2d_p4`, `radialNO2d_p4m`, `radialNO3d_p4`, `radialNO3d_p4m` |
+| Removed in Paddle | `GFNO2d_p4_steer`, `GFNO2d_p4m_steer`, `Unet_Rot2d`, `Unet_Rot_M2d`, `Unet_Rot_3D`                                                                                                                                                                            |
+
+The removed Paddle-only model types depended on `e2cnn` or `escnn`, which require Torch at runtime. The four main experiment commands documented below all use `GFNO2d_p4`, which remains supported.
+
+From `examples/G-FNO`:
+
+### NS
+
+```bash
+python "experiments.py" --seed=1 --data_path="./data/ns_V1e-4_N10000_T30.mat" \
+    --results_path="./results/ns_V1e-4_N10000_T30.mat/GFNO2d_p4" --strategy=teacher_forcing \
+    --T=20 --ntrain=1000 --nvalid=100 --ntest=100 --model_type=GFNO2d_p4 --modes=12 --width=10 \
+    --batch_size=20 --epochs=100 --suffix=seed1 --txt_suffix="ns_V1e-4_N10000_T30.mat_GFNO2d_p4_seed1" \
+    --learning_rate=1e-3 --early_stopping=100 --verbose --super \
+    --super_path="./data/ns_data_V1e-4_N20_T50_R256test.mat" --device=auto
+```
+
+### NS-Sym
+
+```bash
+python "experiments.py" --seed=1 --data_path="./data/ns_V0.0001_N1200_T30_cos4.mat" \
+    --results_path="./results/ns_V0.0001_N1200_T30_cos4.mat/GFNO2d_p4" --strategy=teacher_forcing \
+    --T=10 --ntrain=1000 --nvalid=100 --ntest=100 --model_type=GFNO2d_p4 --modes=12 --width=10 \
+    --batch_size=20 --epochs=100 --suffix=seed1 --txt_suffix="ns_V0.0001_N1200_T30_cos4.mat_GFNO2d_p4_seed1" \
+    --learning_rate=1e-3 --early_stopping=100 --verbose --super \
+    --super_path="./data/ns_V0.0001_N1200_T30_cos4_super.mat" --device=auto
+```
+
+### SWE (PDEArena)
+
+```bash
+python "experiments.py" --seed=1 --data_path="./data/ShallowWater2D" \
+    --results_path="./results/ShallowWater2D/GFNO2d_p4" --strategy=teacher_forcing \
+    --T=9 --ntrain=5600 --nvalid=1120 --ntest=1120 --model_type=GFNO2d_p4 --modes=32 --width=10 \
+    --batch_size=20 --epochs=100 --suffix=seed1 --txt_suffix="ShallowWater2D_GFNO2d_p4_seed1" \
+    --learning_rate=1e-3 --early_stopping=100 --verbose --time_pad --device=auto
+```
+
+### SWE-Sym (PDEBench)
+
+```bash
+python "experiments.py" --seed=1 --data_path="./data/2D_rdb_NA_NA.h5" \
+    --results_path="./results/2D_rdb_NA_NA.h5/GFNO2d_p4" --strategy=teacher_forcing \
+    --T=24 --ntrain=800 --nvalid=100 --ntest=100 --model_type=GFNO2d_p4 --modes=12 --width=10 \
+    --batch_size=20 --epochs=100 --suffix=seed1 --txt_suffix="2D_rdb_NA_NA.h5_GFNO2d_p4_seed1" \
+    --learning_rate=1e-3 --early_stopping=100 --verbose --super --device=auto
+```
+
+## 7. Citation
+
+```latex
+@inproceedings{helwig2023group,
+author = {Jacob Helwig and Xuan Zhang and Cong Fu and Jerry Kurtin and Stephan Wojtowytsch and Shuiwang Ji},
+title = {Group Equivariant {Fourier} Neural Operators for Partial Differential Equations},
+booktitle = {Proceedings of the 40th International Conference on Machine Learning},
+year = {2023},
+}
+```

examples/G-FNO/data_generation/navier_stokes/ns_2d_rt.py

@@ -0,0 +1,153 @@
+import os
+from pathlib import Path
+
+import paddle
+
+"""
+This is a modified version of ns_2d.py from https://github.com/zongyi-li/fourier_neural_operator
+"""
+import argparse
+import math
+from timeit import default_timer
+
+import scipy.io
+from random_fields import GaussianRF
+from tqdm import tqdm
+
+from ppcfd.models.g_fno.paddle_utils import set_runtime_device
+
+
+def navier_stokes_2d(w0, f, domain_size, visc, T, delta_t=0.0001, record_steps=1):
+    N = w0.size()[-1]
+    k_max = math.floor(N / 2.0)
+    steps = math.ceil(T / delta_t)
+    w_h = paddle.fft.rfft2(w0)
+    f_h = paddle.fft.rfft2(f)
+    if len(f_h.size()) < len(w_h.size()):
+        f_h = paddle.unsqueeze(f_h, 0)
+    record_time = math.floor(steps / record_steps)
+    k_y = paddle.cat(
+        (
+            paddle.arange(start=0, end=k_max, step=1, device=w0.device),
+            paddle.arange(start=-k_max, end=0, step=1, device=w0.device),
+        ),
+        0,
+    ).repeat(N, 1)
+    k_x = k_y.transpose(0, 1)
+    k_x = k_x[..., : k_max + 1]
+    k_y = k_y[..., : k_max + 1]
+    lap = 4 * math.pi**2 * (k_x**2 + k_y**2) / domain_size**2
+    lap[0, 0] = 1.0
+    dealias = paddle.unsqueeze(
+        paddle.logical_and(
+            paddle.abs(k_y) <= 2.0 / 3.0 * k_max, paddle.abs(k_x) <= 2.0 / 3.0 * k_max
+        ).float(),
+        0,
+    )
+    sol = paddle.zeros(*w0.size(), record_steps, device=w0.device)
+    sol_t = paddle.zeros(record_steps, device=w0.device)
+    c = 0
+    t = 0.0
+    for j in range(steps):
+        psi_h = w_h / lap
+        q = 2.0 * math.pi / domain_size * k_y * 1.0j * psi_h
+        q = paddle.fft.irfft2(q, s=(N, N))
+        v = -2.0 * math.pi / domain_size * k_x * 1.0j * psi_h
+        v = paddle.fft.irfft2(v, s=(N, N))
+        w_x = 2.0 * math.pi / domain_size * k_x * 1.0j * w_h
+        w_x = paddle.fft.irfft2(w_x, s=(N, N))
+        w_y = 2.0 * math.pi / domain_size * k_y * 1.0j * w_h
+        w_y = paddle.fft.irfft2(w_y, s=(N, N))
+        F_h = paddle.fft.rfft2(q * w_x + v * w_y)
+        F_h = dealias * F_h
+        w_h = (
+            -delta_t * F_h + delta_t * f_h + (1.0 - 0.5 * delta_t * visc * lap) * w_h
+        ) / (1.0 + 0.5 * delta_t * visc * lap)
+        t += delta_t
+        if (j + 1) % record_time == 0:
+            w = paddle.fft.irfft2(w_h, s=(N, N))
+            sol[..., c] = w
+            sol_t[c] = t
+            c += 1
+    return sol, sol_t
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument("--nu", type=float, required=True)
+parser.add_argument("--s", type=int, default=256)
+parser.add_argument("--T", type=int, required=True, help="Time horizon")
+parser.add_argument("--N", type=int, required=True)
+parser.add_argument("--save_path", type=str, required=True)
+parser.add_argument("--bsize", type=int, default=20)
+parser.add_argument("--suffix", type=str, default=None)
+parser.add_argument(
+    "--ntest", type=int, required=True, help="Number of superresolution examples"
+)
+parser.add_argument("--period", type=int, required=True, help="Period if sym is true")
+parser.add_argument(
+    "--sym", action="store_true", default=True, help="Use a symmetric forcing term"
+)
+parser.add_argument("--domain_size", type=float, default=1)
+parser.add_argument(
+    "--device",
+    type=str,
+    default="auto",
+    help="runtime device, e.g. auto, cpu, gpu, xpu, npu, gcu",
+)
+args = parser.parse_args()
+device = set_runtime_device(args.device)
+s = args.s
+N = args.N
+GRF = GaussianRF(2, s, args.domain_size, alpha=2.5, tau=7, device=device)
+t = paddle.linspace(0, args.domain_size, s + 1, device=device)
+t = t[0:-1]
+X, Y = paddle.meshgrid(t, t, indexing="ij")
+if args.sym:
+    f = 0.1 * (
+        paddle.cos(args.period * math.pi * X) + paddle.cos(args.period * math.pi * Y)
+    )
+else:
+    f = 0.1 * (paddle.sin(2 * math.pi * (X + Y)) + paddle.cos(2 * math.pi * (X + Y)))
+record_steps = args.T * 4
+a = paddle.zeros(N, s, s)
+u = paddle.zeros(N, s, s, record_steps)
+bsize = args.bsize
+c = 0
+t0 = default_timer()
+for j in tqdm(range(N // bsize)):
+    w0 = GRF.sample(shape=bsize)
+    sol, sol_t = navier_stokes_2d(
+        w0, f, args.domain_size, args.nu, args.T, 0.0001, record_steps
+    )
+    a[c : c + bsize, ...] = w0
+    u[c : c + bsize, ...] = sol
+    c += bsize
+    t1 = default_timer()
+    print(j, c, t1 - t0)
+a_super = a[-args.ntest :]
+u_super = u[-args.ntest :]
+space_sub = s // 64
+time_sub = 4
+a = a[..., ::space_sub, ::space_sub]
+u = u[..., ::space_sub, ::space_sub, ::time_sub]
+if args.sym:
+    data_name = f"ns_V{args.nu}_N{args.N}_T{args.T}_cos{args.period}{'_' + args.suffix if args.suffix is not None else ''}.mat"
+else:
+    data_name = f"ns_V{args.nu}_N{args.N}_T{args.T}_sin{'_' + args.suffix if args.suffix is not None else ''}.mat"
+super_name = data_name[:-4] + "_super.mat"
+if not os.path.exists(args.save_path):
+    os.makedirs(args.save_path)
+save_dir = os.path.join(args.save_path, data_name)
+super_dir = os.path.join(args.save_path, super_name)
+scipy.io.savemat(
+    save_dir,
+    mdict={"a": a.cpu().numpy(), "u": u.cpu().numpy(), "t": sol_t.cpu().numpy()},
+)
+scipy.io.savemat(
+    super_dir,
+    mdict={
+        "a": a_super.cpu().numpy(),
+        "u": u_super.cpu().numpy(),
+        "t": sol_t.cpu().numpy(),
+    },
+)

examples/G-FNO/data_generation/navier_stokes/random_fields.py

@@ -0,0 +1,97 @@
+import paddle
+
+"""
+Source: https://github.com/zongyi-li/fourier_neural_operator
+"""
+import math
+from timeit import default_timer
+
+
+class GaussianRF(object):
+    def __init__(
+        self,
+        dim,
+        size,
+        domain_size,
+        alpha=2,
+        tau=3,
+        sigma=None,
+        boundary="periodic",
+        device=None,
+    ):
+        self.dim = dim
+        self.device = device
+        if sigma is None:
+            sigma = tau ** (0.5 * (2 * alpha - self.dim))
+        k_max = size // 2
+        if dim == 1:
+            k = paddle.cat(
+                (
+                    paddle.arange(start=0, end=k_max, step=1, device=device),
+                    paddle.arange(start=-k_max, end=0, step=1, device=device),
+                ),
+                0,
+            )
+            self.sqrt_eig = (
+                size
+                * math.sqrt(2.0)
+                * sigma
+                * (4 * math.pi**2 / domain_size**2 * k**2 + tau**2)
+                ** (-alpha / 2.0)
+            )
+            self.sqrt_eig[0] = 0.0
+        elif dim == 2:
+            wavenumers = paddle.cat(
+                (
+                    paddle.arange(start=0, end=k_max, step=1, device=device),
+                    paddle.arange(start=-k_max, end=0, step=1, device=device),
+                ),
+                0,
+            ).repeat(size, 1)
+            k_x = wavenumers.transpose(0, 1)
+            k_y = wavenumers
+            self.sqrt_eig = (
+                size**2
+                * math.sqrt(2.0)
+                * sigma
+                * (
+                    4 * math.pi**2 / domain_size**2 * (k_x**2 + k_y**2)
+                    + tau**2
+                )
+                ** (-alpha / 2.0)
+            )
+            self.sqrt_eig[0, 0] = 0.0
+        elif dim == 3:
+            wavenumers = paddle.cat(
+                (
+                    paddle.arange(start=0, end=k_max, step=1, device=device),
+                    paddle.arange(start=-k_max, end=0, step=1, device=device),
+                ),
+                0,
+            ).repeat(size, size, 1)
+            k_x = wavenumers.transpose(1, 2)
+            k_y = wavenumers
+            k_z = wavenumers.transpose(0, 2)
+            self.sqrt_eig = (
+                size**3
+                * math.sqrt(2.0)
+                * sigma
+                * (
+                    4
+                    * math.pi**2
+                    / domain_size**2
+                    * (k_x**2 + k_y**2 + k_z**2)
+                    + tau**2
+                )
+                ** (-alpha / 2.0)
+            )
+            self.sqrt_eig[0, 0, 0] = 0.0
+        self.size = []
+        for j in range(self.dim):
+            self.size.append(size)
+        self.size = tuple(self.size)
+
+    def sample(self, N):
+        coeff = paddle.randn(N, *self.size, dtype=paddle.complex64, device=self.device)
+        coeff = self.sqrt_eig * coeff
+        return paddle.fft.ifftn(coeff, dim=list(range(-1, -self.dim - 1, -1))).real()