This is a Fourier Neural Operator (FNO) trained on the
turbulent_radiative_layer_2D dataset from
The Well
(PolymathicAI, NeurIPS 2024).
The model learns to predict the next timestep of a
turbulent radiative layer simulation given 4 consecutive
input frames. It predicts all 4 physical fields: density,
pressure, velocity x, and velocity y.
The model shows good prediction quality for velocity
fields. Density and pressure fields exhibit spectral
bias artifacts (vertical stripes) consistent with
small-batch FNO training. This is a baseline result
— improvements possible with larger batch size and
longer training.
1import torch
2from the_well.benchmark.models import FNO
3from the_well.data import WellDataset
4
5# Load dataset
6testset = WellDataset(
7 well_base_path='hf://datasets/polymathic-ai/',
8 well_dataset_name='turbulent_radiative_layer_2D',
9 well_split_name='test',
10 n_steps_input=4,
11)
12
13# Load model
14checkpoint = torch.load('best.pt', map_location='cuda')
15
16model = FNO(
17 dim_in=16,
18 dim_out=4,
19 n_spatial_dims=2,
20 spatial_resolution=(128, 384),
21 modes1=16,
22 modes2=16,
23 hidden_channels=128,
24).cuda()
25
26model.load_state_dict(checkpoint['model_state_dict'])
27model.eval()
If you use this model, please cite The Well dataset:
@inproceedings{the_well_2024,
title={The Well: a Large-Scale Collection of
Diverse Physics Simulations for Machine Learning},
author={Ohana et al.},
booktitle={NeurIPS 2024}
}
Trained by
Sriram Reddy
as a learning project. First ML model trained from scratch.
