polymathic-ai/TFNO-acoustic_scattering_maze

Benchmarking Models on the Well

The Well is a 15TB dataset collection of physics simulations. This model is part of the models that have been benchmarked on the Well.

The models have been trained for a fixed time of 12 hours or up to 500 epochs, whichever happens first. The training was performed on a NVIDIA H100 96GB GPU. In the time dimension, the context length was set to 4. The batch size was set to maximize the memory usage. We experiment with 5 different learning rates for each model on each dataset. We use the model performing best on the validation set to report test set results.

The reported results are here to provide a simple baseline. They should not be considered as state-of-the-art. We hope that the community will build upon these results to develop better architectures for PDE surrogate modeling.

Tensorized Fourier Neural Operator

Implementation of the Tensorized Fourier Neural Operator provided by neuraloperator v0.3.0.

Model Details

For benchmarking on the Well, we used the following parameters.

Parameters	Values
Modes	16
Blocks	4
Hidden Size	128

Trained Model Versions

Below is the list of checkpoints available for the training of TFNO on different datasets of the Well.

Dataset	Learning Rate	Epoch	VRMSE
acoustic_scattering_maze	1E-3	27	0.5034
active_matter	1E-3	243	0.3342
convective_envelope_rsg	1E-3	13	0.0195
gray_scott_reaction_diffusion	5E-3	45	0.1784
helmholtz_staircase	5E-4	131	0.00031
MHD_64	1E-3	155	0.3347
planetswe	5E-4	49	0.1061
post_neutron_star_merger	5E-4	99	0.4064
rayleigh_benard	1E-4	31	0.8568
rayleigh_taylor_instability	1E-4	175	0.2251
shear_flow	1E-3	24	0.3626
supernova_explosion_64	1E-4	35	0.3645
turbulence_gravity_cooling	5E-4	10	0.2789
turbulent_radiative_layer_2D	1E-3	500	0.4938
viscoelastic_instability	5E-3	199	0.7021

Loading the model from Hugging Face

To load the TFNO model trained on the acoustic_scattering_maze of the Well, use the following commands.

from the_well.benchmark.models import TFNO

model = TFNO.from_pretrained("polymathic-ai/TFNO-acoustic_scattering_maze")