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