README.md

Group Equivariant Fourier Neural Operators for Partial Differential Equations

Note on Hyperparameters

To choose the number of modes and channels relative to FNO, the extra parameters in G-FNO kernels added by the stabilizer dimension should be offset by reducing the number of channels (as opposed to the number of modes) as discussed in Appendix A.2 and Section 8.2 of Cohen and Welling (2016).

Introduction

This is the official implementation of G-FNO:

Jacob Helwig*, Xuan Zhang*, Cong Fu, Jerry Kurtin, Stephan Wojtowytsch and Shuiwang Ji. "Group Equivariant Fourier Neural Operators for Partial Differential Equations". [ICML 2023 Poster]

*Equal contribution

Requirements

To create a GFNO conda environment, run:

source setup.sh

Preparing Data

• The Navier-Stokes data with a non-symmetric forcing term (NS) are available via the FNO GitHub. Note that we use both the dataset ns_data_V1e-4_N20_T50_R256test.mat (20 super-resolution test trajectories) and ns_V1e-4_N10000_T30.mat (10,000 downsampled trajectories).

• We use ns_2d_rt.py to generate the Navier-Stokes data with a symmetric forcing term (NS-Sym). To generate this data (ns_V0.001_N1200_T30_cos4.mat for 1,200 downsampled trajectories and ns_V0.001_N1200_T30_cos4_super.mat for 100 super-resolution test trajectories), run:

python ns_2d_rt.py --nu=1e-4 --T=30 --N=1200 --save_path=./data --ntest=100 --period=4

• Instructions for generating the shallow water equations data from PDEArena (SWE) are available via the PDEArena data generation instructions.

• The shallow water equations data from PDEBench (SWE) is available via the PDEBench GitHub.

Run

We use the shell script run_experiment.sh to run all experiments on all datasets and models. Below are commands for training G-FNO2d-p4 on each of the datasets.

NS:

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

NS-Sym:

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

SWE:

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

SWE-Sym:

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

Citation

@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},
}

Acknowledgments

This work was supported in part by National Science Foundation grant IIS-2006861, and by state allocated funds for the Water Exceptional Item through Texas A&M AgriLife Research facilitated by the Texas Water Resources Institute.