OneScience-Group/BCAT-NS2D-Rollout

BCAT-NS2D

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

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

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

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

# 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

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

python scripts/train.py

The best checkpoint is saved to:

weight/best_model.pt

Pretrained weights

The model package includes weights produced with the standard configuration:

weight/best_model.pt

Inference

python scripts/inference.py

Outputs are saved to:

results/metrics.json
results/predictions.npz

Metrics for the included checkpoint:

Metric	Value
relative_l2	1.2696276903152466
mse	0.713721513748169

Evaluation and visualization

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