Sevenzoro321/trl2D-surrogate

FNO Surrogate Model — Turbulent Radiative Layer 2D

Model Description

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.

Dataset

• Name: turbulent_radiative_layer_2D
• Size: 6.9GB
• Resolution: 384×128
• Fields: density, pressure, velocity x, velocity y
• Physics: Hot and cold gas mixing via Kelvin-Helmholtz instability

Training Details

• Architecture: FNO (Fourier Neural Operator)
• Hidden channels: 128
• Modes: 16×16
• Input timesteps: 4 consecutive frames (16 channels)
• Epochs: 200 (best checkpoint at epoch 77)
• Batch size: 4
• Learning rate: 1e-3 with cosine annealing
• GPU: NVIDIA L4 (24GB VRAM)
• Best validation loss: 7.49

Performance Notes

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.

Usage

import torch
from the_well.benchmark.models import FNO
from the_well.data import WellDataset

# Load dataset
testset = WellDataset(
    well_base_path='hf://datasets/polymathic-ai/',
    well_dataset_name='turbulent_radiative_layer_2D',
    well_split_name='test',
    n_steps_input=4,
)

# Load model
checkpoint = torch.load('best.pt', map_location='cuda')

model = FNO(
    dim_in=16,
    dim_out=4,
    n_spatial_dims=2,
    spatial_resolution=(128, 384),
    modes1=16,
    modes2=16,
    hidden_channels=128,
).cuda()

model.load_state_dict(checkpoint['model_state_dict'])
model.eval()

Visualization

Citation

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

Author

Trained by Sriram Reddy as a learning project. First ML model trained from scratch.