README.md

---
license: apache-2.0
language:
- en
library_name: pytorch
datasets:
- OneScience-Group/STFNO_NIMROD
tags:
- OneScience
- computational-fluid-dynamics
- neural-operator
- plasma-simulation
- fusion
- fourier-neural-operator
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">STFNO</span>
  </strong>
</p>

# 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](https://pubs.aip.org/aip/pop/article/31/12/123902/3323878/Sparsified-time-dependent-Fourier-neural-operators)

# 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](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

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

```bash
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**

```bash
# 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**

```bash
# 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:

```bash
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.

```bash
python scripts/train.py
```

The default training configuration saves the checkpoint to:

```text
./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

```bash
python scripts/inference.py
```

Measured metrics for the included checkpoint:

```text
relative_l2 = 0.0647946754
mae = 0.0267913828
max_relative_l2 = 0.0796131641
```

### Evaluation and Visualization

```bash
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](https://pubs.aip.org/aip/pop/article/31/12/123902/3323878/Sparsified-time-dependent-Fourier-neural-operators).
- Use of this project must comply with the licenses and terms that apply to the original paper, dataset, and related resources.