STFNO

Model Overview

STFNO is a two-dimensional NIMROD field-prediction model based on the Sparsified Time-dependent Fourier Neural Operator. It performs one-step prediction of coupled multiphysics fields from fusion-plasma simulations.

This repository is an independent reproduction of the STFNO NIMROD H-mode experiment, implemented through the OneScience workflow from the paper description and official configuration.

Paper: Sparsified time-dependent Fourier neural operators for fusion simulations

Model Description

STFNO extends the Fourier Neural Operator with the sparse variable-dependency structure described in the paper. It jointly models velocity, density, temperature, and magnetic-field components from the NIMROD2D hyperdiffusivity S=64 data. The model takes one time step of a multichannel two-dimensional field as input and predicts the multichannel field at the next time step.

Use Cases

Use Case	Description
Fusion-plasma simulation surrogate	Rapidly predict the next step of two-dimensional multiphysics fields generated by NIMROD
Coupled multiphysics PDE learning	Validate training and inference for a sparse-dependency neural operator on multivariable, time-dependent PDEs

Usage

1. OneCode

Use the online OneCode environment for an intelligent, one-click AI for Science (AI4S) programming experience:

Launch OneCode for one-click AI4S programming

2. Manual Setup

Hardware Requirements

A GPU or DCU is recommended.
A CPU can be used for import checks and small-scale pipeline validation, but full training and inference will be slow.
DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version for the target cluster, is recommended.

Download the Model Package from Hugging Face

pip install -U huggingface_hub
hf download OneScience-Group/STFNO-NIMROD-HMode --local-dir ./STFNO-NIMROD-HMode
cd STFNO-NIMROD-HMode

Set Up the Runtime Environment

DCU Environment

# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
# Installation with uv is also supported
pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai

GPU Environment

# Activate Conda first
conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
conda activate onescience311
# Installation with uv is also supported
pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai

Training Data

The OneScience community provides the STFNO_NIMROD training data as a Hugging Face dataset. Download it with:

hf download OneScience-Group/STFNO_NIMROD --repo-type dataset --local-dir ./data

Then set paths.data_root in config/config.yaml to the directory containing the downloaded data.

This reproduction uses two-dimensional S=64 field data. Ten channels are selected as the model input and output. The frame array has shape [259, 64, 64, 10] and is divided chronologically into 128 training windows and 129 test windows.

Training

The paper evaluates STFNO on NIMROD and GTC data. This reproduction runs the complete NIMROD2D S=64 one-step prediction workflow; it does not reproduce every variable group or ablation study from the paper.

python scripts/train.py

The default training configuration saves the checkpoint to:

./weight/best_model.pth

The checkpoint is saved after epoch 500 and is not selected using test-set metrics.

Model Weights

The weight/ directory contains a checkpoint trained on the NIMROD2D data and ready for direct inference.

Inference

python scripts/inference.py

Measured metrics for the included checkpoint:

relative_l2 = 0.0647946754
mae = 0.0267913828
max_relative_l2 = 0.0796131641

Evaluation and Visualization

python scripts/result.py

Official OneScience Resources

Platform	OneScience Repository	Skills Repository
Gitee	https://gitee.com/onescience-ai/onescience	https://gitee.com/onescience-ai/oneskills
GitHub	https://github.com/onescience-ai/OneScience	https://github.com/onescience-ai/oneskills

Citations and License

Original STFNO paper: Sparsified time-dependent Fourier neural operators for fusion simulations.
Use of this project must comply with the licenses and terms that apply to the original paper, dataset, and related resources.

Downloads last month: -; Downloads are not tracked for this model. How to track

Inference Providers NEW

This model isn't deployed by any Inference Provider. 🙋 Ask for provider support