---
license: mit
language:
- en
library_name: pytorch
tags:
- OneScience
- fluid-dynamics
- navier-stokes
- neural-operator
- transformer
- pde-rollout
---

<p align="center">
  <strong><span style="font-size: 30px;">BCAT-NS2D</span></strong>
</p>

# Model Introduction

BCAT-NS2D is a two-dimensional fluid-sequence prediction model based on BCAT, the Block Causal Transformer developed jointly by researchers from the UCLA Department of Mathematics and the CMU Department of Mathematical Sciences. It uses historical flow fields as context and autoregressively predicts the evolution of velocity and vorticity at subsequent time steps.

Paper: [BCAT: A Block Causal Transformer for PDE Foundation Models for Fluid Dynamics](https://arxiv.org/abs/2501.18972)

# Model Description

BCAT-NS2D takes a sequence of two-dimensional flow fields containing the velocity components and vorticity `[u, v, omega]` as input. Each time step is divided into spatial patches and combined with temporal and patch-position embeddings. A block-causal attention mask prevents information leakage from future temporal blocks. The Transformer encodes the historical spatiotemporal features and predicts the velocity and vorticity fields at later time steps. This implementation is trained with mean squared error on synthetic, divergence-free, two-dimensional vortex sequences with periodic boundaries.

## Intended Uses

| Use case | Description |
| :--- | :--- |
| 2D flow-sequence prediction | Predict the velocity components and vorticity distribution at future time steps from historical flow fields. |
| Navier–Stokes rollout | Approximate the temporal evolution of two-dimensional incompressible Navier–Stokes flow. |
| Transformer PDE validation | Validate flow-patch encoding, spatiotemporal embeddings, block-causal attention, and rollout behavior. |

# Usage

## 1. OneCode

Try one-click AI-for-Science programming in the OneCode online environment:

[Launch OneCode](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)

## 2. Manual Setup

**Hardware requirements**

- A GPU or DCU is recommended.
- A CPU is sufficient for imports, forward passes, and small connectivity checks.
- DCU users must install DTK in advance and use a OneScience Python environment compatible with the cluster.

### Download the model repository from Hugging Face

```bash
pip install -U huggingface_hub
hf download OneScience-Group/BCAT-NS2D-Rollout --local-dir ./BCAT-NS2D-Rollout
cd BCAT-NS2D-Rollout
```

### Install the runtime environment

**DCU environment**

```bash
# Activate DTK first, then create and activate the Conda environment.
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
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

No external dataset is required. `scripts/fake_data.py` generates divergence-free two-dimensional vortex sequences with periodic boundaries for training and inference.

The data configuration is stored in `config/config.yaml`. The defaults are:

- Grid size: \(16\times16\)
- Sequence length: 5 time steps
- Physical channels: horizontal velocity \(u\), vertical velocity \(v\), and vorticity \(\omega\)
- Kinematic viscosity: \(0.001\)
- Sample variation: randomized initial vortex phases and horizontal/vertical translation speeds
- Temporal evolution: gradual decay caused by translation and viscosity
- Objective: predict the next `[u, v, omega]` field from the current and previous flow fields

The generated data is intended primarily to validate the BCAT rollout pipeline.

### Train

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

The best checkpoint is saved to:

```text
weight/best_model.pt
```

### Pretrained weights

The model package includes weights produced with the standard configuration:

```text
weight/best_model.pt
```

### Inference

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

Outputs are saved to:

```text
results/metrics.json
results/predictions.npz
```

Metrics for the included checkpoint:

| Metric | Value |
| :--- | ---: |
| relative_l2 | 1.2696276903152466 |
| mse | 0.713721513748169 |

### Evaluation and visualization

```bash
python scripts/result.py
```

The result summary is saved to `results/summary.json`.

# OneScience

| Platform | OneScience repository | OneSkills 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 |

# Citation and License

- Paper: [BCAT: A Block Causal Transformer for PDE Foundation Models for Fluid Dynamics](https://arxiv.org/abs/2501.18972).
- This reproduction package is released under the MIT License, as declared in the Hugging Face model-card metadata. The paper and any third-party materials remain subject to their respective terms.