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