StFT

Spatio-temporal Fourier Transformer for Long-term Dynamics Prediction

Overview

StFT is a multi-scale spatiotemporal forecasting model for long-horizon dynamics prediction.
This repository provides training code for plasma MHD data and core model components.

Repository Structure

• train.py — training entrypoint
• StFT_3D.py — StFT model definition
• data_utils.py — dataset/loss/grid utilities
• model_utils.py — Transformer layers and positional embeddings

---

Installation

git clone https://github.com/BerkeleyLab/StFT.git
cd StFT

Install dependencies:

pip install -r requirements.txt

GPU note: On CUDA systems, install a CUDA-compatible PyTorch build first using the official PyTorch instructions.

Training StFT

To train StFT on the plasma MHD dataset:

python train.py

By default, the results will be saved to the ~/ray_results at home directory.
To customize the saved directory, you can change the save_path variable in the train.py file.

Example Results

Mean relative L2 error across autoregressive rollout timesteps.

Qualitative long-horizon comparison across autoregressive baselines.

References

@article{stft2026,
  title={St{FT}: Spatio-temporal Fourier Transformer for Long-term Dynamics Prediction},
  author={Long, Da and Zhe, Shandian and Williams, Samuel and Oliker, Leonid and Bai, Zhe},
  journal={Transactions on Machine Learning Research},
  issn={2835-8856},
  year={2026},
  url={https://openreview.net/forum?id=o9Cb0ri2oW},
}

See the LICENSE file for copyright and licensing information.

---

_ Copyright Notice _

Spatio-temporal Fourier Transformer for Long-term Dynamics Prediction (StFT) Copyright (c) 2025, The Regents of the University of California, through Lawrence Berkeley National Laboratory (subject to receipt of any required approvals from the U.S. Dept. of Energy). All rights reserved.

If you have questions about your rights to use or distribute this software, please contact Berkeley Lab's Intellectual Property Office at IPO@lbl.gov.

NOTICE. This Software was developed under funding from the U.S. Department of Energy and the U.S. Government consequently retains certain rights. As such, the U.S. Government has been granted for itself and others acting on its behalf a paid-up, nonexclusive, irrevocable, worldwide license in the Software to reproduce, distribute copies to the public, prepare derivative works, and perform publicly and display publicly, and to permit others to do so.

Questions? Contact Zhe Bai (zhebai@lbl.gov)

---