Metadata-Version: 2.4
Name: stft-transformer
Version: 0.1.15
Summary: Spatial-temporal Fourier Transformer
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: numpy
Requires-Dist: torch
Requires-Dist: einops
Requires-Dist: ray[train]
Dynamic: license-file

<h1 align="center">StFT</h1>
<h3 align="center">Spatio-temporal Fourier Transformer for Long-term Dynamics Prediction</h3>

<p align="center">
  <img src="https://raw.githubusercontent.com/BerkeleyLab/StFT/main/figs/stft-architecture.png" alt="Spatio-temporal Fourier Transformer architecture" width="80%"/>
  <br>
</p>

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

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

Install dependencies:

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

```bash
python train.py
```

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

## Example Results

<p align="center"> 
  <img src="https://raw.githubusercontent.com/BerkeleyLab/StFT/main/figs/stft-error-vs-time.png" alt="Autoregressive error over time" width="80%"/><br>
  Mean relative L2 error across autoregressive rollout timesteps. 
</p> 
<br>
<p align="center"> 
  <img src="https://raw.githubusercontent.com/BerkeleyLab/StFT/main/figs/stft-qualitative-comparison.png" alt="Qualitative comparison" width="80%"/><br>
  Qualitative long-horizon error comparison across autoregressive baselines. 
</p>

## References

```bibtex

@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](LICENSE) 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)

---
