Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -0,0 +1,140 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- fluid dynamics
+- unsteady turbulent flow prediction
+- long-range mesh dependency modeling
+frameworks: PyTorch
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">EagleMeshTransformer</span>
+  </strong>
+</p>
+
+# Model Overview
+
+EagleMeshTransformer is a multiscale Mesh Transformer developed by the LIRIS research laboratory in Lyon, France, for fluid prediction on dynamic unstructured meshes. It is particularly well suited to unsteady turbulent flows and problems involving long-range dependencies in flow fields.
+
+Paper: [EAGLE: Large-scale Learning of Turbulent Fluid Dynamics with Mesh Transformers](https://arxiv.org/abs/2302.10803).
+
+# Model Description
+
+EagleMeshTransformer uses a multiscale Mesh Transformer architecture trained on the EAGLE dataset to predict velocity and pressure fields in complex unsteady flows.
+   
+
+
+## Use Cases
+
+| Use Case | Description |
+|---|---|
+| Unsteady turbulent flow prediction | Predict velocity and pressure fields in complex, aperiodic turbulent flows involving drones, jets, wakes, and similar systems |
+| Unstructured-mesh simulation | Process irregular mesh data defined on complex geometries |
+| CFD surrogate acceleration | Provide fast approximations of conventional Navier–Stokes and CFD simulations |
+| Long-horizon physical prediction | Predict the evolution of physical states through autoregressive rollouts |
+
+# 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.
+- A CPU can be used for import checks and small-scale pipeline validation, but full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or later, or the OneScience-recommended version for the target cluster, is recommended.
+
+### Download the Model Package
+
+```bash
+modelscope download --model OneScience/EagleMeshTransformer --local_dir ./EagleMeshTransformer
+cd EagleMeshTransformer
+```
+
+### Set Up the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Installation with uv is also supported
+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
+# Installation with uv is also supported
+pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Synthetic Data Validation
+
+The default configuration points to the synthetic data directory in this repository and sets `training.max_epoch` to `1`. Generate a minimal EAGLE NPZ dataset to validate the training and inference pipelines:
+
+```bash
+python scripts/fake_data.py
+```
+### Training Data
+The OneScience community provides the EAGLE dataset for training. Download it with the command below and verify that the data path in `conf/config.yaml` is configured correctly.
+
+```bash
+modelscope download --dataset OneScience/eagle --local_dir ./data
+```
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multiple GPUs:
+
+```bash
+torchrun --nproc_per_node=8 --nnodes=1 --rdzv_id=1000 --rdzv_backend=c10d --max_restarts=0 --master_addr="localhost" --master_port=29500 scripts/train.py
+```
+
+Training saves `best_model.pth` in the `weight/` directory.
+
+### Model Weights
+This repository will provide pretrained EagleMeshTransformer weights in the `weights/` directory. The weights will be uploaded soon.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference results are saved to `result/output/`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+
+# 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
+- Original EagleMeshTransformer paper: [EAGLE: Large-scale Learning of Turbulent Fluid Dynamics with Mesh Transformers](https://arxiv.org/abs/2302.10803).
+- This repository retains source attribution and has been adapted for automated execution through OneScience and ModelScope.

config/config.yaml

@@ -0,0 +1,41 @@
+model:
+  name: "EagleMeshTransformer"
+  architecture: "GraphViT"
+  state_size: 4
+  w_size: 32
+
+datapipe:
+  source:
+    data_dir: "./data/fake/Eagle_dataset"
+    cluster_dir: "./data/fake/Eagle_dataset"
+    splits_dir: "./data/fake/splits"
+
+  verbose: false
+
+  data:
+    window_length_train: 3
+    window_length_val: 3
+    window_length_test: 3
+    n_cluster: 1
+    normalized: true
+    type_as_onehot: true
+    with_cells: true
+
+  dataloader:
+    batch_size: 1
+    num_workers: 0
+
+training:
+  max_epoch: 1
+  lr: 1.0e-4
+  loss_alpha: 0.1
+  patience: 10
+  checkpoint_dir: "./weight"
+  seed: 0
+  device: "cpu"
+
+inference:
+  checkpoint_path: "./weight/best_model.pth"
+  output_dir: "./result/output"
+  batch_size: 1
+  device: "cpu"

configuration.json

@@ -0,0 +1,9 @@
+{
+  "framework": "Pytorch",
+  "task": "other",
+  "model": {
+    "type": "EagleMeshTransformer",
+    "architecture": "GraphViT"
+  },
+  "allow_remote": false
+}

data/splits/test.txt

@@ -0,0 +1,118 @@
+Spl/184/2
+Tri/194/2
+Cre/185/1
+Tri/199/1
+Spl/188/1
+Tri/182/2
+Spl/181/2
+Spl/199/1
+Spl/187/2
+Tri/189/2
+Cre/182/1
+Spl/181/1
+Spl/182/2
+Tri/197/2
+Tri/181/2
+Spl/199/2
+Spl/188/2
+Tri/185/1
+Spl/182/1
+Spl/183/1
+Spl/185/2
+Spl/191/2
+Spl/191/1
+Cre/191/1
+Tri/195/1
+Tri/191/1
+Spl/189/1
+Cre/191/2
+Spl/185/1
+Tri/200/2
+Cre/183/1
+Tri/187/2
+Cre/200/1
+Tri/186/2
+Tri/198/2
+Cre/188/1
+Cre/192/1
+Tri/197/1
+Tri/185/2
+Cre/189/2
+Cre/193/1
+Spl/200/1
+Cre/182/2
+Spl/198/2
+Spl/200/2
+Tri/193/1
+Tri/194/1
+Cre/181/1
+Tri/182/1
+Tri/184/2
+Cre/190/1
+Tri/198/1
+Spl/190/1
+Cre/194/1
+Tri/184/1
+Cre/186/2
+Spl/193/1
+Tri/199/2
+Spl/183/2
+Spl/197/1
+Cre/184/2
+Spl/195/2
+Cre/195/2
+Cre/198/1
+Cre/197/2
+Tri/190/2
+Tri/193/2
+Spl/186/1
+Cre/199/2
+Cre/193/2
+Tri/195/2
+Cre/199/1
+Spl/184/1
+Cre/187/2
+Tri/188/2
+Tri/181/1
+Tri/183/2
+Cre/198/2
+Cre/189/1
+Cre/195/1
+Tri/192/1
+Cre/196/2
+Cre/192/2
+Spl/197/2
+Tri/186/1
+Tri/196/2
+Tri/188/1
+Spl/196/1
+Tri/200/1
+Spl/190/2
+Cre/194/2
+Spl/192/1
+Cre/190/2
+Cre/187/1
+Spl/198/1
+Spl/189/2
+Cre/186/1
+Cre/188/2
+Spl/193/2
+Spl/195/1
+Tri/187/1
+Spl/192/2
+Cre/184/1
+Tri/189/1
+Spl/187/1
+Tri/192/2
+Tri/190/1
+Cre/181/2
+Cre/185/2
+Tri/196/1
+Cre/200/2
+Cre/183/2
+Cre/196/1
+Spl/186/2
+Spl/196/2
+Tri/183/1
+Tri/191/2
+Cre/197/1

data/splits/train.txt

@@ -0,0 +1,947 @@
+Spl/20/2
+Spl/26/1
+Spl/85/1
+Cre/55/2
+Spl/38/2
+Tri/124/1
+Tri/78/1
+Cre/103/2
+Cre/148/2
+Tri/96/2
+Tri/66/2
+Cre/63/2
+Spl/22/2
+Tri/110/2
+Tri/35/2
+Tri/26/2
+Cre/108/1
+Spl/124/2
+Spl/73/1
+Tri/15/2
+Tri/125/1
+Spl/53/2
+Tri/79/2
+Spl/57/2
+Cre/109/2
+Spl/32/1
+Spl/51/2
+Cre/145/2
+Tri/102/1
+Tri/8/1
+Cre/28/2
+Tri/17/1
+Tri/73/2
+Spl/137/1
+Spl/2/1
+Tri/1/2
+Spl/106/2
+Tri/49/2
+Cre/73/1
+Cre/25/1
+Tri/106/2
+Tri/140/2
+Cre/143/1
+Tri/10/2
+Tri/40/1
+Tri/122/1
+Cre/111/2
+Cre/41/2
+Spl/71/1
+Spl/146/1
+Spl/50/2
+Spl/136/1
+Cre/104/1
+Cre/152/1
+Spl/30/1
+Tri/105/1
+Tri/118/2
+Spl/149/1
+Cre/138/2
+Cre/122/2
+Spl/43/2
+Cre/31/2
+Spl/120/2
+Cre/69/2
+Spl/153/1
+Spl/44/1
+Spl/155/2
+Spl/155/1
+Tri/103/2
+Tri/42/1
+Tri/101/2
+Tri/74/1
+Cre/122/1
+Spl/34/1
+Spl/34/2
+Cre/141/1
+Cre/35/1
+Spl/129/1
+Spl/70/1
+Tri/130/1
+Tri/133/2
+Tri/44/2
+Spl/121/1
+Cre/13/2
+Tri/60/2
+Cre/23/2
+Cre/134/2
+Tri/100/1
+Tri/6/1
+Tri/99/2
+Spl/116/2
+Tri/100/2
+Tri/6/2
+Tri/124/2
+Cre/101/2
+Cre/63/1
+Tri/147/2
+Spl/103/2
+Cre/18/1
+Cre/131/1
+Cre/81/1
+Spl/74/1
+Cre/18/2
+Spl/139/2
+Cre/74/2
+Spl/20/1
+Cre/15/2
+Cre/153/1
+Cre/70/2
+Tri/104/1
+Spl/58/2
+Spl/102/2
+Cre/136/1
+Tri/25/1
+Cre/28/1
+Spl/48/1
+Tri/30/1
+Cre/35/2
+Tri/103/1
+Spl/59/1
+Spl/84/2
+Tri/5/2
+Cre/151/2
+Tri/1/1
+Cre/65/1
+Cre/117/1
+Spl/10/1
+Spl/117/2
+Spl/101/2
+Cre/152/2
+Tri/140/1
+Cre/62/2
+Tri/133/1
+Spl/133/2
+Tri/19/2
+Spl/45/2
+Tri/109/1
+Cre/44/2
+Cre/126/1
+Tri/134/1
+Cre/23/1
+Spl/150/2
+Cre/88/2
+Tri/147/1
+Cre/58/1
+Tri/37/2
+Spl/150/1
+Tri/159/2
+Spl/96/2
+Spl/66/1
+Spl/158/2
+Spl/122/1
+Spl/35/2
+Spl/126/1
+Tri/157/1
+Tri/18/1
+Spl/58/1
+Tri/52/1
+Tri/73/1
+Tri/151/2
+Cre/53/1
+Spl/35/1
+Spl/28/1
+Spl/118/2
+Cre/119/1
+Spl/4/1
+Cre/57/2
+Spl/81/2
+Spl/139/1
+Cre/46/1
+Spl/12/2
+Cre/131/2
+Cre/55/1
+Spl/60/1
+Tri/9/2
+Spl/36/2
+Cre/71/1
+Spl/32/2
+Tri/81/1
+Spl/100/2
+Tri/122/2
+Spl/33/2
+Spl/37/1
+Spl/91/2
+Spl/27/2
+Spl/156/1
+Spl/47/1
+Cre/6/2
+Tri/46/1
+Cre/54/2
+Tri/137/1
+Spl/114/2
+Tri/78/2
+Spl/23/2
+Spl/84/1
+Spl/101/1
+Spl/40/1
+Cre/29/2
+Tri/65/1
+Spl/60/2
+Tri/142/1
+Spl/154/1
+Cre/40/1
+Tri/38/2
+Spl/87/1
+Cre/14/1
+Tri/95/1
+Tri/118/1
+Spl/81/1
+Spl/152/1
+Tri/61/2
+Tri/129/2
+Cre/119/2
+Cre/142/2
+Tri/96/1
+Cre/47/1
+Spl/86/2
+Spl/147/2
+Tri/26/1
+Tri/49/1
+Cre/91/2
+Cre/45/1
+Spl/156/2
+Tri/19/1
+Cre/59/1
+Cre/146/1
+Spl/25/2
+Cre/110/2
+Spl/127/2
+Tri/136/1
+Cre/4/1
+Tri/148/1
+Cre/48/1
+Spl/2/2
+Spl/14/2
+Spl/99/2
+Tri/22/2
+Cre/116/1
+Tri/114/1
+Cre/94/2
+Spl/109/1
+Cre/85/1
+Cre/27/1
+Tri/97/1
+Tri/77/2
+Cre/132/2
+Cre/137/1
+Cre/77/1
+Tri/144/1
+Tri/127/1
+Spl/135/1
+Cre/3/2
+Spl/96/1
+Cre/90/1
+Cre/89/1
+Tri/88/1
+Tri/92/2
+Tri/37/1
+Tri/125/2
+Tri/86/2
+Spl/46/1
+Cre/65/2
+Cre/84/1
+Spl/134/1
+Cre/137/2
+Tri/17/2
+Spl/133/1
+Spl/11/2
+Tri/155/1
+Tri/57/1
+Cre/140/1
+Cre/26/1
+Spl/113/2
+Spl/63/2
+Spl/105/1
+Tri/51/2
+Tri/88/2
+Tri/36/2
+Tri/80/2
+Tri/92/1
+Cre/79/1
+Tri/28/1
+Spl/56/2
+Tri/59/2
+Tri/98/1
+Spl/6/2
+Tri/120/1
+Spl/75/1
+Cre/147/1
+Tri/27/2
+Tri/71/1
+Spl/33/1
+Tri/27/1
+Spl/77/2
+Tri/48/2
+Spl/15/2
+Spl/89/1
+Spl/140/1
+Cre/117/2
+Spl/148/1
+Tri/43/2
+Spl/119/1
+Cre/101/1
+Spl/152/2
+Cre/49/1
+Tri/39/1
+Spl/157/1
+Spl/11/1
+Tri/32/2
+Spl/21/1
+Cre/76/2
+Spl/72/1
+Cre/87/1
+Cre/158/1
+Cre/120/2
+Tri/66/1
+Cre/16/1
+Cre/29/1
+Spl/3/1
+Tri/23/1
+Cre/50/2
+Cre/89/2
+Tri/144/2
+Tri/76/1
+Spl/95/1
+Tri/136/2
+Tri/104/2
+Cre/159/1
+Tri/36/1
+Spl/77/1
+Spl/61/2
+Tri/94/2
+Spl/88/2
+Tri/158/1
+Cre/13/1
+Tri/79/1
+Cre/99/2
+Cre/120/1
+Cre/157/2
+Spl/79/2
+Tri/138/1
+Cre/91/1
+Tri/76/2
+Tri/13/1
+Tri/138/2
+Tri/58/1
+Cre/126/2
+Tri/116/2
+Tri/112/2
+Tri/105/2
+Tri/60/1
+Cre/15/1
+Cre/86/2
+Cre/74/1
+Spl/63/1
+Spl/18/2
+Tri/83/1
+Tri/145/2
+Cre/16/2
+Spl/67/1
+Spl/13/2
+Spl/83/1
+Cre/78/2
+Cre/68/1
+Cre/127/1
+Tri/61/1
+Tri/106/1
+Cre/66/1
+Cre/146/2
+Cre/130/2
+Tri/129/1
+Spl/62/1
+Spl/125/2
+Tri/115/1
+Cre/6/1
+Spl/75/2
+Tri/120/2
+Cre/156/2
+Cre/100/1
+Cre/124/1
+Cre/107/2
+Spl/68/1
+Cre/138/1
+Cre/75/1
+Cre/53/2
+Tri/34/1
+Tri/152/2
+Tri/90/2
+Tri/160/2
+Cre/154/2
+Tri/115/2
+Spl/86/1
+Spl/10/2
+Spl/24/1
+Cre/147/2
+Cre/9/2
+Cre/129/2
+Cre/136/2
+Spl/1/2
+Cre/130/1
+Cre/96/1
+Cre/118/2
+Tri/8/2
+Cre/39/1
+Spl/15/1
+Spl/65/1
+Cre/115/2
+Tri/35/1
+Cre/4/2
+Spl/134/2
+Spl/31/2
+Cre/34/1
+Cre/58/2
+Tri/77/1
+Cre/11/1
+Cre/133/1
+Tri/135/1
+Tri/45/2
+Spl/99/1
+Cre/121/1
+Tri/47/2
+Tri/62/2
+Spl/19/2
+Spl/143/2
+Tri/130/2
+Spl/78/1
+Tri/56/2
+Cre/40/2
+Tri/41/2
+Spl/49/2
+Tri/84/1
+Tri/74/2
+Cre/10/1
+Cre/83/1
+Tri/50/1
+Spl/40/2
+Tri/32/1
+Tri/21/2
+Spl/73/2
+Cre/3/1
+Cre/57/1
+Cre/157/1
+Tri/128/2
+Tri/110/1
+Tri/64/2
+Tri/107/1
+Tri/3/2
+Spl/79/1
+Cre/128/2
+Tri/54/1
+Spl/37/2
+Tri/68/1
+Tri/14/2
+Cre/127/2
+Spl/43/1
+Tri/65/2
+Cre/109/1
+Tri/15/1
+Tri/4/1
+Cre/104/2
+Spl/149/2
+Spl/66/2
+Spl/122/2
+Cre/153/2
+Spl/18/1
+Cre/39/2
+Tri/52/2
+Cre/140/2
+Cre/93/2
+Spl/64/2
+Cre/118/1
+Tri/113/1
+Cre/27/2
+Cre/38/1
+Tri/159/1
+Cre/56/1
+Spl/128/1
+Cre/11/2
+Tri/5/1
+Cre/54/1
+Cre/66/2
+Spl/148/2
+Cre/128/1
+Cre/31/1
+Spl/45/1
+Tri/33/2
+Cre/116/2
+Tri/72/1
+Tri/152/1
+Tri/20/2
+Tri/117/1
+Tri/62/1
+Tri/51/1
+Spl/135/2
+Cre/30/1
+Spl/117/1
+Cre/50/1
+Spl/88/1
+Tri/97/2
+Spl/16/2
+Cre/150/2
+Tri/155/2
+Spl/159/2
+Spl/74/2
+Tri/29/2
+Tri/156/1
+Tri/57/2
+Cre/112/1
+Tri/91/1
+Tri/70/2
+Spl/112/1
+Tri/42/2
+Tri/150/2
+Cre/5/1
+Spl/127/1
+Spl/47/2
+Spl/51/1
+Cre/102/1
+Cre/44/1
+Cre/34/2
+Cre/45/2
+Tri/131/1
+Cre/86/1
+Cre/80/2
+Spl/24/2
+Spl/41/2
+Cre/20/1
+Spl/72/2
+Tri/95/2
+Tri/143/2
+Spl/57/1
+Spl/65/2
+Tri/16/2
+Spl/138/2
+Tri/83/2
+Tri/87/1
+Spl/144/1
+Spl/160/2
+Spl/83/2
+Cre/154/1
+Tri/9/1
+Tri/47/1
+Tri/148/2
+Cre/37/1
+Tri/108/2
+Cre/92/2
+Cre/158/2
+Cre/25/2
+Spl/50/1
+Tri/85/2
+Tri/141/1
+Cre/151/1
+Cre/67/1
+Tri/156/2
+Tri/43/1
+Tri/25/2
+Cre/149/1
+Spl/94/1
+Cre/84/2
+Cre/139/1
+Cre/22/1
+Tri/14/1
+Spl/129/2
+Tri/29/1
+Tri/13/2
+Spl/123/1
+Spl/76/1
+Cre/78/1
+Cre/71/2
+Cre/5/2
+Tri/33/1
+Cre/107/1
+Cre/92/1
+Tri/116/1
+Tri/31/2
+Tri/63/2
+Tri/75/2
+Tri/160/1
+Tri/139/1
+Spl/7/2
+Tri/28/2
+Tri/2/1
+Spl/4/2
+Spl/158/1
+Cre/22/2
+Cre/10/2
+Cre/143/2
+Cre/80/1
+Cre/59/2
+Tri/18/2
+Spl/132/2
+Cre/69/1
+Cre/132/1
+Spl/61/1
+Tri/72/2
+Spl/85/2
+Spl/59/2
+Cre/42/2
+Tri/44/1
+Spl/97/1
+Cre/52/1
+Cre/94/1
+Cre/144/1
+Cre/155/2
+Cre/102/2
+Cre/61/2
+Spl/120/1
+Cre/51/1
+Tri/111/2
+Cre/2/1
+Spl/130/1
+Cre/110/1
+Cre/2/2
+Spl/159/1
+Tri/54/2
+Spl/8/1
+Cre/8/2
+Spl/94/2
+Cre/43/2
+Spl/130/2
+Cre/97/1
+Tri/53/2
+Cre/26/2
+Spl/46/2
+Cre/160/2
+Spl/121/2
+Tri/94/1
+Spl/93/2
+Cre/1/1
+Tri/150/1
+Spl/26/2
+Spl/80/1
+Tri/111/1
+Tri/69/2
+Cre/81/2
+Spl/157/2
+Cre/33/2
+Tri/69/1
+Cre/85/2
+Cre/24/2
+Tri/55/1
+Spl/9/2
+Tri/46/2
+Cre/17/1
+Spl/16/1
+Spl/119/2
+Tri/34/2
+Spl/140/2
+Spl/7/1
+Cre/115/1
+Spl/22/1
+Cre/82/2
+Tri/12/1
+Cre/99/1
+Tri/145/1
+Spl/125/1
+Spl/131/1
+Spl/14/1
+Tri/53/1
+Spl/55/1
+Spl/143/1
+Cre/68/2
+Spl/56/1
+Spl/5/1
+Tri/143/1
+Spl/128/2
+Spl/154/2
+Tri/153/1
+Tri/21/1
+Tri/87/2
+Cre/123/2
+Cre/70/1
+Cre/144/2
+Spl/54/1
+Spl/69/2
+Spl/78/2
+Cre/41/1
+Tri/75/1
+Cre/7/1
+Tri/128/1
+Cre/79/2
+Spl/36/1
+Spl/17/2
+Tri/39/2
+Spl/126/2
+Spl/13/1
+Cre/160/1
+Spl/114/1
+Tri/132/1
+Tri/50/2
+Spl/29/2
+Tri/82/1
+Cre/7/2
+Cre/90/2
+Spl/21/2
+Tri/67/1
+Tri/146/2
+Cre/61/1
+Spl/146/2
+Tri/112/1
+Tri/16/1
+Cre/83/2
+Cre/82/1
+Spl/27/1
+Cre/30/2
+Spl/108/1
+Spl/25/1
+Spl/70/2
+Cre/87/2
+Tri/132/2
+Tri/81/2
+Tri/24/2
+Spl/52/1
+Cre/103/1
+Spl/62/2
+Spl/92/2
+Spl/52/2
+Spl/104/1
+Cre/21/1
+Tri/80/1
+Tri/98/2
+Spl/76/2
+Cre/93/1
+Tri/64/1
+Tri/20/1
+Spl/112/2
+Spl/12/1
+Cre/37/2
+Cre/114/1
+Cre/36/2
+Spl/97/2
+Spl/71/2
+Tri/158/2
+Tri/58/2
+Tri/119/2
+Cre/12/2
+Spl/87/2
+Cre/32/1
+Spl/53/1
+Cre/108/2
+Tri/86/1
+Cre/64/1
+Spl/48/2
+Cre/95/2
+Cre/100/2
+Cre/12/1
+Spl/64/1
+Spl/80/2
+Tri/30/2
+Spl/123/2
+Spl/91/1
+Cre/62/1
+Spl/102/1
+Tri/151/1
+Spl/100/1
+Spl/6/1
+Spl/110/2
+Cre/14/2
+Spl/141/2
+Spl/29/1
+Spl/144/2
+Tri/114/2
+Spl/5/2
+Tri/134/2
+Cre/24/1
+Cre/113/1
+Spl/38/1
+Cre/75/2
+Spl/138/1
+Spl/131/2
+Tri/59/1
+Spl/109/2
+Cre/88/1
+Tri/31/1
+Cre/113/2
+Spl/90/2
+Cre/60/2
+Cre/139/2
+Spl/55/2
+Cre/9/1
+Cre/125/1
+Cre/98/1
+Tri/12/2
+Spl/3/2
+Spl/106/1
+Tri/90/1
+Cre/60/1
+Spl/98/2
+Tri/7/1
+Cre/43/1
+Tri/23/2
+Spl/39/2
+Spl/9/1
+Spl/19/1
+Tri/93/1
+Tri/71/2
+Cre/141/2
+Tri/137/2
+Tri/55/2
+Tri/142/2
+Spl/145/1
+Cre/111/1
+Cre/121/2
+Tri/108/1
+Cre/142/1
+Tri/109/2
+Spl/151/2
+Tri/41/1
+Spl/30/2
+Spl/108/2
+Spl/39/1
+Spl/115/1
+Cre/20/2
+Cre/77/2
+Cre/67/2
+Spl/68/2
+Spl/28/2
+Cre/46/2
+Cre/36/1
+Tri/84/2
+Spl/132/1
+Spl/17/1
+Tri/154/2
+Cre/17/2
+Cre/96/2
+Cre/97/2
+Cre/95/1
+Cre/149/2
+Cre/124/2
+Tri/135/2
+Tri/3/1
+Spl/160/1
+Spl/89/2
+Tri/146/1
+Spl/115/2
+Tri/68/2
+Cre/148/1
+Tri/141/2
+Spl/110/1
+Tri/149/2
+Spl/147/1
+Tri/117/2
+Spl/98/1
+Tri/11/1
+Tri/121/1
+Cre/64/2
+Spl/42/2
+Cre/150/1
+Tri/123/1
+Cre/155/1
+Cre/56/2
+Tri/91/2
+Tri/101/1
+Tri/4/2
+Cre/42/1
+Cre/47/2
+Tri/70/1
+Tri/119/1
+Spl/118/1
+Cre/73/2
+Spl/92/1
+Cre/123/1
+Cre/114/2
+Tri/131/2
+Cre/98/2
+Cre/33/1
+Spl/41/1
+Tri/7/2
+Spl/142/1
+Cre/125/2
+Tri/11/2
+Spl/104/2
+Tri/93/2
+Tri/154/1
+Cre/1/2
+Spl/1/1
+Spl/23/1
+Spl/82/2
+Spl/103/1
+Spl/44/2
+Tri/45/1
+Spl/49/1
+Cre/48/2
+Cre/8/1
+Tri/121/2
+Spl/93/1
+Cre/32/2
+Cre/134/1
+Tri/107/2
+Spl/124/1
+Tri/67/2
+Spl/153/2
+Tri/123/2
+Spl/8/2
+Tri/153/2
+Tri/2/2
+Tri/63/1
+Tri/22/1
+Cre/52/2
+Tri/127/2
+Tri/82/2
+Tri/10/1
+Spl/31/1
+Spl/116/1
+Spl/141/1
+Cre/38/2
+Spl/54/2
+Cre/51/2
+Cre/159/2
+Cre/76/1
+Spl/142/2
+Tri/149/1
+Spl/95/2
+Spl/42/1
+Tri/113/2
+Cre/129/1
+Spl/90/1
+Tri/157/2
+Cre/49/2
+Spl/67/2
+Cre/72/2
+Cre/145/1
+Cre/72/1
+Tri/99/1
+Cre/112/2
+Spl/145/2
+Spl/137/2
+Tri/139/2
+Tri/24/1
+Tri/85/1
+Cre/156/1
+Spl/105/2
+Spl/82/1
+Spl/151/1
+Tri/48/1
+Tri/89/2
+Tri/102/2
+Cre/133/2
+Cre/21/2
+Tri/56/1
+Spl/136/2
+Spl/69/1
+Tri/38/1
+Spl/113/1
+Tri/89/1
+Cre/105/1
+Cre/105/2

data/splits/valid.txt

@@ -0,0 +1,118 @@
+Spl/178/1
+Cre/167/1
+Tri/177/1
+Spl/172/1
+Tri/173/1
+Cre/165/2
+Spl/177/1
+Cre/177/1
+Tri/175/1
+Spl/171/1
+Tri/175/2
+Cre/171/1
+Cre/168/2
+Cre/163/2
+Spl/161/2
+Cre/161/2
+Tri/166/1
+Spl/172/2
+Tri/174/1
+Tri/171/2
+Tri/165/1
+Tri/180/1
+Tri/178/1
+Spl/163/1
+Tri/162/2
+Cre/178/1
+Tri/172/2
+Cre/170/2
+Spl/170/1
+Tri/180/2
+Tri/166/2
+Tri/170/2
+Spl/162/2
+Cre/168/1
+Spl/171/2
+Tri/179/2
+Cre/165/1
+Cre/180/2
+Spl/180/1
+Cre/161/1
+Tri/176/2
+Spl/161/1
+Spl/164/1
+Tri/171/1
+Cre/162/1
+Cre/169/2
+Spl/173/2
+Spl/176/2
+Cre/176/1
+Spl/168/2
+Spl/180/2
+Cre/176/2
+Tri/176/1
+Cre/179/2
+Cre/174/2
+Tri/164/2
+Cre/178/2
+Tri/161/2
+Spl/173/1
+Cre/177/2
+Spl/165/2
+Spl/168/1
+Cre/162/2
+Spl/175/2
+Tri/167/1
+Tri/163/2
+Spl/177/2
+Tri/173/2
+Cre/166/2
+Cre/180/1
+Tri/179/1
+Cre/164/2
+Tri/164/1
+Spl/175/1
+Cre/166/1
+Spl/179/2
+Cre/163/1
+Spl/163/2
+Cre/164/1
+Tri/178/2
+Spl/167/1
+Spl/169/2
+Cre/173/2
+Cre/175/1
+Spl/165/1
+Tri/172/1
+Tri/168/1
+Spl/174/1
+Tri/163/1
+Tri/169/1
+Cre/174/1
+Cre/169/1
+Tri/161/1
+Tri/174/2
+Spl/170/2
+Spl/166/1
+Tri/169/2
+Spl/179/1
+Spl/176/1
+Spl/164/2
+Spl/162/1
+Tri/167/2
+Spl/174/2
+Cre/171/2
+Tri/165/2
+Tri/162/1
+Spl/166/2
+Tri/170/1
+Tri/168/2
+Cre/167/2
+Cre/173/1
+Cre/179/1
+Tri/177/2
+Spl/169/1
+Cre/170/1
+Cre/175/2
+Spl/178/2
+Spl/167/2

model/graphViT.py

@@ -0,0 +1,173 @@
+import torch
+import torch.nn as nn
+
+from onescience.modules.decoder.graphvit_decoder import GraphViTDecoder
+from onescience.modules.embedding.fourier_pos_embedding import FourierPosEmbedding
+from onescience.modules.encoder.graphvit_encoder import GraphViTEncoder
+from onescience.modules.pooling.rnn_cluster_pooling import RNNClusterPooling
+from onescience.modules.transformer.preln_transformer_block import (
+    PreLNTransformerBlock,
+)
+
+# 节点类型常量
+NODE_NORMAL = 0
+NODE_INPUT = 4
+NODE_OUTPUT = 5
+NODE_WALL = 6
+NODE_DISABLE = 2
+
+class GraphViT(nn.Module):
+    """
+
+    该模型利用图神经网络提取局部特征，通过聚类池化在潜在空间进行全局 Transformer 交互，
+    最后解码回物理空间进行下一时刻的状态预测。
+
+    Args:
+        state_size (int): 物理状态的维度（如速度 u, v, w）。
+        w_size (int): 潜在空间（Cluster Level）的特征维度。默认值: 512。
+        n_attention (int): Transformer 注意力层的数量。默认值: 4。
+        nb_gn (int): Encoder 中 GNN 的层数。默认值: 4。
+        n_heads (int): Transformer 注意力头数。默认值: 4。
+    """
+    def __init__(self, state_size, w_size=512, n_attention=4, nb_gn=4, n_heads=4):
+        super(GraphViT, self).__init__()
+        
+        # 傅里叶位置编码参数
+        pos_start = -3
+        pos_length = 8
+
+        # 1. 编码器: Mesh Space -> Latent Graph
+        self.encoder = GraphViTEncoder(
+            nb_gn=nb_gn, 
+            state_size=state_size, 
+            pos_length=pos_length
+        )
+
+        # 2. 池化层: Node Features -> Cluster Features
+        self.graph_pooling = RNNClusterPooling(
+            w_size=w_size, 
+            pos_length=pos_length
+        )
+
+        # 3. 解码器: Cluster Features -> Node State Update
+        self.graph_retrieve = GraphViTDecoder(
+            w_size=w_size, 
+            pos_length=pos_length, 
+            state_size=state_size
+        )
+
+        # 4. 潜在空间 Transformer 交互层
+        self.attention = nn.ModuleList([
+            PreLNTransformerBlock(
+                w_size=w_size, 
+                pos_length=pos_length, 
+                n_heads=n_heads
+            )
+            for _ in range(n_attention)
+        ])
+        
+        self.ln = nn.LayerNorm(w_size)
+        self.noise_std = 0.0
+        
+        # 5. 坐标嵌入层
+        self.positional_encoder = FourierPosEmbedding(
+            pos_start=pos_start, 
+            pos_length=pos_length
+        )
+
+    def forward(
+        self,
+        mesh_pos,
+        edges,
+        state,
+        node_type,
+        clusters,
+        clusters_mask,
+        apply_noise=False,
+    ):
+        """
+        前向传播 (自回归滚动预测)。
+        
+        参数:
+            mesh_pos: (B, T, N, 3) 节点坐标
+            edges: (B, T, M, 2) 边索引
+            state: (B, T, N, state_size) 节点物理状态
+            node_type: (B, T, N, 9) 节点类型 (One-hot)
+            clusters: (B, T, K, C_max) 簇索引
+            clusters_mask: (B, T, K, C_max) 簇掩码
+            apply_noise: 是否在初始状态施加噪声 (Training 时通常为 True)
+        """
+        if apply_noise:
+            mask = torch.logical_or(
+                node_type[:, 0, :, NODE_NORMAL] == 1,
+                node_type[:, 0, :, NODE_OUTPUT] == 1,
+            )
+            noise = torch.randn_like(state[:, 0]).to(state.device) * self.noise_std
+            state[:, 0][mask] = state[:, 0][mask] + noise[mask]
+
+        state_hat = [state[:, 0]] 
+        output_hat = []
+        target = []
+
+        # 从 t=1 开始预测，基于 t-1 的状态
+        for t in range(1, state.shape[1]):
+            # 1. 生成位置编码
+            mesh_posenc, cluster_posenc = self.positional_encoder(
+                mesh_pos[:, t - 1], clusters[:, t - 1], clusters_mask[:, t - 1]
+            )
+
+            # 2. 图节点和边编码
+            V, E = self.encoder(
+                mesh_pos[:, t - 1],
+                edges[:, t - 1],
+                state_hat[-1],
+                node_type[:, t - 1],
+                mesh_posenc,
+            )
+
+            # 3. 聚类池化
+            W = self.graph_pooling(
+                V, clusters[:, t - 1], mesh_posenc, clusters_mask[:, t - 1]
+            )
+
+            # 4. 构建 Attention Mask
+            attention_mask = clusters_mask[:, t - 1].sum(-1, keepdim=True) == 0
+            attention_mask = (
+                attention_mask.unsqueeze(1)
+                .repeat(1, len(self.attention), 1, W.shape[1])
+                .view(-1, W.shape[1], W.shape[1])
+            )
+            attention_mask[:, torch.eye(W.shape[1], dtype=torch.bool)] = False
+            attention_mask = attention_mask.transpose(-1, -2)
+
+            # 5. 潜在空间 Transformer 交互
+            for a in self.attention:
+                W = a(W, attention_mask, cluster_posenc)
+            W = self.ln(W)
+
+            # 6. 解码预测更新量
+            next_output = self.graph_retrieve(
+                W, V, clusters[:, t - 1], mesh_posenc, edges[:, t - 1], E
+            )
+
+            # 7. 更新下一时刻状态
+            next_state = state_hat[-1] + next_output
+            target.append(state[:, t] - state_hat[-1])
+
+            # 8. 强制边界条件 (Mask 覆盖)
+            mask = torch.logical_or(
+                node_type[:, t, :, NODE_INPUT] == 1, node_type[:, t, :, NODE_WALL] == 1
+            )
+            mask = torch.logical_or(mask, node_type[:, t, :, NODE_DISABLE] == 1)
+            next_state[mask, :] = state[:, t][mask, :]
+
+            state_hat.append(next_state)
+            output_hat.append(next_output)
+
+        # 堆叠时间步
+        velocity_hat = torch.stack(state_hat, dim=1)
+        output_hat = torch.stack(output_hat, dim=1)
+        target = torch.stack(target, dim=1)
+
+        # velocity_hat: [B, T, N, S], output_hat: [B, T-1, N, S]
+        return velocity_hat, output_hat, target

scripts/fake_data.py

@@ -0,0 +1,86 @@
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "model"))
+
+
+def write_split(split_dir: Path, name: str, samples):
+    split_dir.mkdir(parents=True, exist_ok=True)
+    with open(split_dir / f"{name}.txt", "w", encoding="utf-8") as handle:
+        for sample in samples:
+            handle.write(f"{sample}\n")
+
+
+def make_sample(path: Path, seed: int, steps: int = 990, nodes: int = 6):
+    rng = np.random.default_rng(seed)
+    path.mkdir(parents=True, exist_ok=True)
+
+    base = np.array(
+        [
+            [0.0, 0.0],
+            [1.0, 0.0],
+            [0.0, 1.0],
+            [1.0, 1.0],
+            [0.5, 0.2],
+            [0.5, 0.8],
+        ],
+        dtype=np.float32,
+    )
+    drift = np.linspace(0.0, 0.04, steps, dtype=np.float32).reshape(steps, 1, 1)
+    pointcloud = base.reshape(1, nodes, 2) + drift
+
+    t = np.arange(steps, dtype=np.float32).reshape(steps, 1)
+    x = pointcloud[..., 0]
+    y = pointcloud[..., 1]
+    vx = np.sin(x + t * 0.1).astype(np.float32)
+    vy = np.cos(y + t * 0.1).astype(np.float32)
+    ps = (x + y + 0.01 * rng.standard_normal((steps, nodes))).astype(np.float32)
+    pg = (x - y + 0.01 * rng.standard_normal((steps, nodes))).astype(np.float32)
+    mask = np.zeros((steps, nodes), dtype=np.int64)
+
+    triangles = np.array(
+        [
+            [0, 1, 4],
+            [0, 4, 2],
+            [1, 3, 4],
+            [2, 4, 5],
+            [3, 5, 4],
+        ],
+        dtype=np.int64,
+    )
+    triangles = np.repeat(triangles.reshape(1, -1, 3), steps, axis=0)
+
+    np.savez(
+        path / "sim.npz",
+        pointcloud=pointcloud,
+        VX=vx,
+        VY=vy,
+        PS=ps,
+        PG=pg,
+        mask=mask,
+    )
+    np.save(path / "triangles.npy", triangles)
+
+
+def main():
+    os.chdir(PROJECT_ROOT)
+    data_root = PROJECT_ROOT / "data" / "fake" / "Eagle_dataset"
+    split_dir = PROJECT_ROOT / "data" / "fake" / "splits"
+    samples = ["Cre/fake_case_000/1", "Cre/fake_case_001/1", "Cre/fake_case_002/1"]
+
+    for idx, sample in enumerate(samples):
+        make_sample(data_root / sample, seed=idx)
+
+    write_split(split_dir, "train", samples[:2])
+    write_split(split_dir, "valid", samples[2:])
+    write_split(split_dir, "test", samples[2:])
+    print(f"Fake Eagle data written to: {data_root}")
+    print(f"Fake splits written to: {split_dir}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,121 @@
+import ctypes
+import os
+import sys
+import sysconfig
+from pathlib import Path
+
+
+def preload_python_shared_library():
+    """Make libpython visible to native extensions loaded with ctypes."""
+    libdir = sysconfig.get_config_var("LIBDIR")
+    version = sysconfig.get_config_var("VERSION")
+    if not libdir or not version:
+        return
+
+    candidates = [
+        Path(libdir) / f"libpython{version}.so.1.0",
+        Path(libdir) / f"libpython{version}.so",
+    ]
+    for libpython in candidates:
+        if libpython.exists():
+            ctypes.CDLL(str(libpython), mode=ctypes.RTLD_GLOBAL)
+            return
+
+
+preload_python_shared_library()
+
+import numpy as np
+import torch
+from tqdm import tqdm
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.insert(0, str(root_path))
+
+from model.graphViT import GraphViT
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+from onescience.datapipes.cfd import EagleDatapipe
+
+
+def resolve_project_path(path):
+    path = Path(path)
+    return path if path.is_absolute() else root_path / path
+
+
+def fix_single_cluster_path(datapipe, cfg_data):
+    if int(cfg_data.data.n_cluster) != 1:
+        return
+
+    cluster_path = Path(cfg_data.source.cluster_dir)
+    for dataset_name in ("train_dataset", "val_dataset", "test_dataset"):
+        dataset = getattr(datapipe, dataset_name, None)
+        if dataset is not None and getattr(dataset, "cluster_path", None) is None:
+            dataset.cluster_path = cluster_path
+
+
+def main():
+    os.chdir(root_path)
+    DistributedManager.initialize()
+    manager = DistributedManager()
+
+    config_path = root_path / "config" / "config.yaml"
+    cfg_model = YParams(config_path, "model")
+    cfg_data = YParams(config_path, "datapipe")
+    cfg_infer = YParams(config_path, "inference")
+
+    checkpoint_path = resolve_project_path(cfg_infer.checkpoint_path)
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(
+            f"Checkpoint not found: {checkpoint_path}. Run `python scripts/train.py` first."
+        )
+
+    datapipe = EagleDatapipe(params=cfg_data, distributed=False)
+    fix_single_cluster_path(datapipe, cfg_data)
+    dataloader, _ = datapipe.test_dataloader(batch_size=int(cfg_infer.batch_size))
+    device_name = cfg_infer.get("device", "auto")
+    device = manager.device if device_name == "auto" else torch.device(device_name)
+    model = GraphViT(state_size=cfg_model.state_size, w_size=cfg_model.w_size).to(device)
+    checkpoint = torch.load(checkpoint_path, map_location=device, weights_only=False)
+    model.load_state_dict(checkpoint.get("model_state_dict", checkpoint))
+
+    output_dir = resolve_project_path(cfg_infer.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    model.eval()
+    saved = 0
+    with torch.no_grad():
+        for idx, x in enumerate(tqdm(dataloader, desc="Inference")):
+            if not x:
+                continue
+            mesh_pos = x["mesh_pos"].to(device)
+            edges = x["edges"].to(device).long()
+            velocity = x["velocity"].to(device)
+            pressure = x["pressure"].to(device)
+            node_type = x["node_type"].to(device)
+            clusters = x["cluster"].to(device).long()
+            clusters_mask = x["cluster_mask"].to(device).long()
+
+            state = torch.cat([velocity, pressure], dim=-1)
+            state_hat, output, target = model(
+                mesh_pos,
+                edges,
+                state,
+                node_type,
+                clusters,
+                clusters_mask,
+                apply_noise=False,
+            )
+            velocity_hat, pressure_hat = dataloader.dataset.denormalize(
+                state_hat[..., :2], state_hat[..., 2:]
+            )
+            pred = torch.cat([velocity_hat, pressure_hat], dim=-1).cpu().numpy()
+            np.save(output_dir / f"prediction_{idx:04d}.npy", pred)
+            saved += 1
+
+    print(f"Saved {saved} prediction file(s) to {output_dir}")
+    manager.cleanup()
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,23 @@
+import sys
+from pathlib import Path
+
+import numpy as np
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def main():
+    output_dir = PROJECT_ROOT / "result" / "output"
+    files = sorted(output_dir.glob("prediction_*.npy"))
+    if not files:
+        print(f"No prediction files found in {output_dir}")
+        return
+
+    print(f"Found {len(files)} prediction file(s) in {output_dir}")
+    for file in files[:5]:
+        arr = np.load(file)
+        print(f"{file.name}: shape={arr.shape}, mean={arr.mean():.6f}, std={arr.std():.6f}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,246 @@
+import ctypes
+import logging
+import os
+import random
+import sys
+import sysconfig
+import time
+from pathlib import Path
+
+
+def preload_python_shared_library():
+    """Make libpython visible to native extensions loaded with ctypes."""
+    libdir = sysconfig.get_config_var("LIBDIR")
+    version = sysconfig.get_config_var("VERSION")
+    if not libdir or not version:
+        return
+
+    candidates = [
+        Path(libdir) / f"libpython{version}.so.1.0",
+        Path(libdir) / f"libpython{version}.so",
+    ]
+    for libpython in candidates:
+        if libpython.exists():
+            ctypes.CDLL(str(libpython), mode=ctypes.RTLD_GLOBAL)
+            return
+
+
+preload_python_shared_library()
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel
+from tqdm import tqdm
+
+# 获取项目根目录（train.py上级的上级）
+root_path = Path(__file__).parent.parent
+sys.path.insert(0, str(root_path))
+
+from model.graphViT import GraphViT
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+from onescience.datapipes.cfd import EagleDatapipe
+
+
+
+def save_best_model(model, optimizer, checkpoint_dir: str):
+    Path(checkpoint_dir).mkdir(parents=True, exist_ok=True)
+    model_to_save = model.module if hasattr(model, "module") else model
+    torch.save(
+        {
+            "model_state_dict": model_to_save.state_dict(),
+            "optimizer_state_dict": optimizer.state_dict(),
+        },
+        Path(checkpoint_dir) / "best_model.pth",
+    )
+
+
+def load_best_model(model, checkpoint_dir: str, device: torch.device):
+    ckpt_path = Path(checkpoint_dir) / "best_model.pth"
+    checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False)
+    state_dict = checkpoint.get("model_state_dict", checkpoint)
+    model.load_state_dict(state_dict)
+
+
+def setup_logging(rank: int):
+    logging.basicConfig(
+        level=logging.INFO if rank == 0 else logging.WARNING,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+        stream=sys.stdout,
+        force=True,
+    )
+    return logging.getLogger("train")
+
+
+def get_loss(velocity, pressure, output, state_hat, target, mask, alpha):
+    velocity = velocity[:, 1:]
+    pressure = pressure[:, 1:]
+    velocity_hat = state_hat[:, 1:, :, :2]
+    pressure_hat = state_hat[:, 1:, :, 2:]
+    mask = mask[:, 1:].unsqueeze(-1)
+
+    loss_velocity = torch.sqrt(((velocity * mask - velocity_hat * mask) ** 2).mean(dim=-1)).mean()
+    loss_pressure = torch.sqrt(((pressure * mask - pressure_hat * mask) ** 2).mean(dim=-1)).mean()
+    mse = nn.MSELoss()
+    loss = mse(target[..., :2] * mask, output[..., :2] * mask)
+    loss = loss + alpha * mse(target[..., 2:] * mask, output[..., 2:] * mask)
+    return {"loss": loss, "MSE_velocity": loss_velocity, "MSE_pressure": loss_pressure}
+
+
+def move_batch(x, device):
+    return {
+        "mesh_pos": x["mesh_pos"].to(device),
+        "edges": x["edges"].to(device).long(),
+        "velocity": x["velocity"].to(device),
+        "pressure": x["pressure"].to(device),
+        "node_type": x["node_type"].to(device),
+        "mask": x["mask"].to(device),
+        "cluster": x["cluster"].to(device).long(),
+        "cluster_mask": x["cluster_mask"].to(device).long(),
+    }
+
+
+def fix_single_cluster_path(datapipe, cfg_data):
+    if int(cfg_data.data.n_cluster) != 1:
+        return
+
+    cluster_path = Path(cfg_data.source.cluster_dir)
+    for dataset_name in ("train_dataset", "val_dataset", "test_dataset"):
+        dataset = getattr(datapipe, dataset_name, None)
+        if dataset is not None and getattr(dataset, "cluster_path", None) is None:
+            dataset.cluster_path = cluster_path
+
+
+def validate(model, dataloader, device, alpha, manager):
+    model.eval()
+    total_loss, count = 0.0, 0
+    with torch.no_grad():
+        for x in dataloader:
+            if not x:
+                continue
+            batch = move_batch(x, device)
+            state = torch.cat([batch["velocity"], batch["pressure"]], dim=-1)
+            state_hat, output, target = model(
+                batch["mesh_pos"],
+                batch["edges"],
+                state,
+                batch["node_type"],
+                batch["cluster"],
+                batch["cluster_mask"],
+                apply_noise=False,
+            )
+            dataset = dataloader.dataset
+            state_hat[..., :2], state_hat[..., 2:] = dataset.denormalize(
+                state_hat[..., :2], state_hat[..., 2:]
+            )
+            velocity, pressure = dataset.denormalize(batch["velocity"], batch["pressure"])
+            costs = get_loss(velocity, pressure, output, state_hat, target, batch["mask"], alpha)
+            if manager.world_size > 1:
+                dist.all_reduce(costs["loss"], op=dist.ReduceOp.AVG)
+            total_loss += costs["loss"].item()
+            count += 1
+    return total_loss / max(count, 1)
+
+
+def main():
+    os.chdir(root_path)
+    DistributedManager.initialize()
+    manager = DistributedManager()
+    logger = setup_logging(manager.rank)
+
+    config_path = root_path / "config" / "config.yaml"
+    cfg_model = YParams(config_path, "model")
+    cfg_data = YParams(config_path, "datapipe")
+    cfg_train = YParams(config_path, "training")
+
+    seed = int(cfg_train.get("seed", 0))
+    torch.manual_seed(seed)
+    random.seed(seed)
+    np.random.seed(seed)
+
+    datapipe = EagleDatapipe(params=cfg_data, distributed=(manager.world_size > 1))
+    fix_single_cluster_path(datapipe, cfg_data)
+    train_loader, train_sampler = datapipe.train_dataloader()
+    val_loader, val_sampler = datapipe.val_dataloader()
+
+    device_name = cfg_train.get("device", "auto")
+    device = manager.device if device_name == "auto" else torch.device(device_name)
+    model = GraphViT(state_size=cfg_model.state_size, w_size=cfg_model.w_size).to(device)
+    if manager.world_size > 1:
+        model = DistributedDataParallel(model, device_ids=[manager.local_rank], output_device=manager.local_rank)
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=float(cfg_train.lr))
+    best_valid_loss = float("inf")
+    best_epoch = 0
+
+    for epoch in range(int(cfg_train.max_epoch)):
+        start = time.time()
+        if manager.world_size > 1:
+            train_sampler.set_epoch(epoch)
+            if val_sampler:
+                val_sampler.set_epoch(epoch)
+
+        model.train()
+        train_loss, count = 0.0, 0
+        pbar = tqdm(train_loader, desc=f"Epoch {epoch + 1}", disable=(manager.rank != 0))
+        for x in pbar:
+            if not x:
+                continue
+            batch = move_batch(x, device)
+            state = torch.cat([batch["velocity"], batch["pressure"]], dim=-1)
+            state_hat, output, target = model(
+                batch["mesh_pos"],
+                batch["edges"],
+                state,
+                batch["node_type"],
+                batch["cluster"],
+                batch["cluster_mask"],
+                apply_noise=True,
+            )
+            state_hat[..., :2], state_hat[..., 2:] = train_loader.dataset.denormalize(
+                state_hat[..., :2], state_hat[..., 2:]
+            )
+            velocity, pressure = train_loader.dataset.denormalize(batch["velocity"], batch["pressure"])
+            costs = get_loss(velocity, pressure, output, state_hat, target, batch["mask"], cfg_train.loss_alpha)
+
+            optimizer.zero_grad()
+            costs["loss"].backward()
+            optimizer.step()
+
+            train_loss += costs["loss"].item()
+            count += 1
+            pbar.set_postfix(loss=f"{costs['loss'].item():.6f}")
+
+        train_loss /= max(count, 1)
+        valid_loss = validate(model, val_loader, device, cfg_train.loss_alpha, manager)
+        if manager.rank == 0:
+            logger.info(
+                "Epoch %s/%s | %.2fs | train %.6f | valid %.6f",
+                epoch + 1,
+                cfg_train.max_epoch,
+                time.time() - start,
+                train_loss,
+                valid_loss,
+            )
+            if valid_loss < best_valid_loss:
+                best_valid_loss = valid_loss
+                best_epoch = epoch
+                save_best_model(model, optimizer, cfg_train.checkpoint_dir)
+                logger.info("Saved checkpoint to %s/best_model.pth", cfg_train.checkpoint_dir)
+            if epoch - best_epoch > int(cfg_train.patience):
+                break
+
+    if manager.rank == 0:
+        final_model = GraphViT(state_size=cfg_model.state_size, w_size=cfg_model.w_size).to(device)
+        load_best_model(final_model, cfg_train.checkpoint_dir, device)
+        test_loader, _ = datapipe.test_dataloader()
+        test_loss = validate(final_model, test_loader, device, cfg_train.loss_alpha, manager)
+        logger.info("Final test loss: %.6f", test_loss)
+
+    manager.cleanup()
+
+
+if __name__ == "__main__":
+    main()