README.md

---
frameworks:
  - pytorch
language:
  - en
license: apache-2.0
tags:
  - OneScience
  - FactFormer
  - neural-operator
  - factorized-attention
  - computational-fluid-dynamics
  - Kolmogorov-flow
tasks:
  - time-series-prediction
---
<p align="center">
  <strong>
    <span style="font-size: 30px;">FactFormer</span>
  </strong>
</p>

# Model Overview

FactFormer (Factorized Transformer) is a Transformer architecture for PDE surrogate modeling. This model predicts the temporal evolution of the vorticity field in two-dimensional Kolmogorov flow, using the first 10 time steps to predict the subsequent 16 by default.

Paper: [Scalable Transformer for PDE Surrogate Modeling](https://arxiv.org/abs/2305.17560)

# Model Description

FactFormer factorizes global attention on a two-dimensional regular grid into attention operations along the two spatial axes. This preserves long-range spatial interactions while reducing computational and memory costs. The model predicts future states in blocks through latent propagation and rolls out the complete forecast window autoregressively.

# Use Cases

| Use Case | Description |
| :---: | :--- |
| Vorticity time-series prediction | Predict future Kolmogorov-flow states from historical two-dimensional vorticity fields |
| PDE surrogate modeling | Learn mappings from historical to future physical fields on regular grids |
| Factorized-attention research | Evaluate axial attention for high-resolution physical-field modeling |
| Pipeline validation | Validate training and inference using the bundled weights or a small-scale configuration |

# 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 for training and inference.
- A CPU can be used for small-scale pipeline validation, although training with the default configuration will be slow.
- DCU users must install DTK and a PyTorch environment compatible with the target cluster.

### Download the Model Package

```bash
modelscope download --model OneScience/FactFormer --local_dir ./FactFormer
cd FactFormer
```

### Set Up the Runtime Environment

**DCU Environment**

```bash
# Activate DTK and Conda first
conda create -n onescience311 python=3.11 -y
conda activate onescience311
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
pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
```

### Training Data

The model uses two-dimensional Kolmogorov-flow data with an original array shape of `[120, 320, 256, 256]`, representing trajectories, time steps, and the two spatial dimensions, respectively. Because the data file is large, it is read on demand using memory mapping.

Download the data with the following command and verify that the data path in `conf/config.yaml` is configured correctly:

```bash
modelscope download --dataset OneScience/Kolmogorov_flow_2d --local_dir ./data
```

### Training

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

The default weights are saved to `weight/factformer_kolmogorov.pt`.

### Model Weights

This repository provides weights trained on the two-dimensional Kolmogorov-flow dataset in the `weight/` directory.

### Inference, Evaluation, and Visualization

The model package includes weights for pipeline validation. After preparing the data, run:

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

The script loads `weight/factformer_kolmogorov.pt` by default. Predicted tensors and visualizations are saved to the `result/` directory. Training and inference parameters can be modified in `conf/config.yaml`.

# 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

- Li, Shu, and Barati Farimani. [Scalable Transformer for PDE Surrogate Modeling](https://arxiv.org/abs/2305.17560).
- The Kolmogorov-flow data is used for two-dimensional vorticity prediction at Re=1000.
- This model package is released under the Apache-2.0 license and retains attribution to the original paper and data sources.