Upload folder using huggingface_hub

README.md

@@ -0,0 +1,122 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- fluid dynamics
+- airfoil aerodynamic prediction
+frameworks: PyTorch
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Transolver-Airfoil-Design</span>
+  </strong>
+</p>
+
+# Model Overview
+
+Transolver-Airfoil-Design is a two-dimensional external-flow prediction model for airfoils, built on Transolver by Tsinghua University's THUML group. It rapidly predicts flow-field distributions and aerodynamic performance around airfoils.
+
+Paper: [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://arxiv.org/abs/2402.02366)
+
+# Model Description
+Transolver-Airfoil-Design uses a Transformer architecture with Physics-Attention and is adapted for training on the unstructured-mesh AirfRANS airfoil dataset. It predicts velocity fields, pressure fields, and drag coefficients for airfoil geometries.
+
+## Use Cases
+
+| Use Case | Description |
+| :--- | :--- |
+| Airfoil aerodynamic design | Rapidly predict two-dimensional external flow around airfoils to screen candidate geometries |
+| Industrial simulation acceleration | Accelerate large-scale simulations using PDE surrogate modeling on complex geometries |
+
+# 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/Transolver-Airfoil-Design --local_dir ./Transolver-Airfoil-Design
+cd Transolver-Airfoil-Design
+```
+
+### 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
+```
+
+### Training Data
+The OneScience community provides the `airfrans` dataset for training. Download it with the command below and verify that the data path in `config/config.yaml` is configured correctly:
+
+```bash
+modelscope download --dataset OneScience/airfrans --local_dir ./data
+```
+
+### Training
+
+```bash
+python scripts/train.py
+```
+
+By default, training saves the following checkpoint:
+
+```text
+./weight/Transolver.pth
+```
+
+### Model Weights
+This repository will provide model weights pretrained on ShapeNetCar data in the `weights/` directory. The weights will be uploaded soon.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+### 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 Transolver paper: [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://arxiv.org/pdf/2402.02366).
+- This repository retains source attribution and has been adapted for automated execution through OneScience and ModelScope.

conf/config.yaml

@@ -0,0 +1,64 @@
+model:
+  name: "Transolver"
+  specific_params:
+    Transolver: &transolver_params
+      n_hidden: 32
+      n_layers: 2
+      space_dim: 7
+      fun_dim: 0
+      n_head: 4
+      mlp_ratio: 2
+      out_dim: 4
+      slice_num: 8
+      unified_pos: 0
+      build_graph: false
+      r: 0.05
+      max_neighbors: 32
+    Transolver_plus: *transolver_params
+
+datapipe:
+  name: "AirfRANS"
+  backend: "fake_airfrans"
+  verbose: false
+  source:
+    data_dir: "./data/fake_airfrans"
+    stats_dir: "./data/stats"
+  data:
+    splits:
+      task: "fake"
+      train_name: "train"
+      val_name: "val"
+      test_name: "test"
+      val_split_ratio: 0.25
+    sampling:
+      sample_strategy: null
+      n_boot: 128
+      surf_ratio: 0.25
+    crop: null
+    subsampling: 64
+  dataloader:
+    batch_size: 1
+    num_workers: 0
+
+training:
+  device: "cpu"
+  max_epoch: 1
+  lr: 0.001
+  gpuid: 0
+  patience: 20
+  progress_log_interval: 5
+  loss_criterion: "MSE_weighted"
+  loss_weight: 1.0
+  checkpoint_dir: "./weight"
+  result_dir: "./result"
+  n_test: 2
+  seed: 42
+
+fake_data:
+  num_train: 4
+  num_val: 2
+  num_test: 2
+  num_nodes: 64
+  input_dim: 7
+  output_dim: 4
+  surface_ratio: 0.25

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other","model":"Transolver-Airfoil-Design","pipeline_tag":"cfd-airfoil-surrogate"}

model/Transolver2D.py

@@ -0,0 +1,136 @@
+import torch
+import torch.nn as nn
+import numpy as np
+try:
+    from timm.layers import trunc_normal_
+except ImportError:
+    from torch.nn.init import trunc_normal_
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.Transolver_block import Transolver_block
+
+class Transolver2D(nn.Module):
+    """
+    Transolver2D 模型。
+
+    该模型专为处理二维物理场问题（如 2D CFD, 弹性力学等）而设计。
+    架构与 Transolver3D 类似，但针对 2D 空间特性进行了配置。
+    它包含预处理 MLP 映射和多层 Transolver Block 堆叠，用于提取非结构化或结构化网格上的物理特征。
+
+    Args:
+        space_dim (int): 空间维度。默认值: 1。
+        n_layers (int): Transolver Block 的层数。默认值: 5。
+        n_hidden (int): 隐藏层特征维度。默认值: 256。
+        dropout (float): Dropout 概率。默认值: 0。
+        n_head (int): 注意力头数。默认值: 8。
+        act (str): 激活函数类型。默认值: 'gelu'。
+        mlp_ratio (float): MLP 膨胀比率。默认值: 1。
+        fun_dim (int): 输入物理场的特征维度。默认值: 1。
+        out_dim (int): 输出特征维度。默认值: 1。
+        slice_num (int): 物理注意力中的切片数量。默认值: 32。
+        ref (int): 统一位置编码的参考网格分辨率。默认值: 8。
+        unified_pos (bool): 是否使用统一位置编码。默认值: False。
+
+    形状:
+        输入 data: 包含 x (物理状态) 和 pos (坐标) 的数据对象。
+        输出: (N, out_dim)，预测的物理场。
+    """
+    def __init__(self,
+                 space_dim=1,
+                 n_layers=5,
+                 n_hidden=256,
+                 dropout=0,
+                 n_head=8,
+                 act='gelu',
+                 mlp_ratio=1,
+                 fun_dim=1,
+                 out_dim=1,
+                 slice_num=32,
+                 ref=8,
+                 unified_pos=False
+                 ):
+        super(Transolver2D, self).__init__()
+        self.__name__ = 'Transolver'
+        self.ref = ref
+        self.unified_pos = unified_pos
+    
+        if self.unified_pos:
+            input_dim = fun_dim + space_dim + self.ref * self.ref
+        else:
+            input_dim = fun_dim + space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            hidden_dims=[n_hidden * 2], 
+            output_dim=n_hidden,
+            activation=act,
+            use_bias=True,
+            use_skip_connection=False
+        )
+        self.n_hidden = n_hidden
+        self.space_dim = space_dim
+
+        self.blocks = nn.ModuleList([
+            Transolver_block(
+                num_heads=n_head,
+                hidden_dim=n_hidden,
+                dropout=dropout,
+                act=act,
+                mlp_ratio=mlp_ratio,
+                out_dim=out_dim,
+                slice_num=slice_num,
+                last_layer=(_ == n_layers - 1),
+                geotype='unstructured' 
+            )
+            for _ in range(n_layers)
+        ])
+        
+        self.initialize_weights()
+        self.placeholder = nn.Parameter((1 / (n_hidden)) * torch.rand(n_hidden, dtype=torch.float))
+
+    def initialize_weights(self):
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_grid(self, my_pos):
+        batchsize = my_pos.shape[0]
+
+        gridx = torch.tensor(np.linspace(-2, 4, self.ref), dtype=torch.float)
+        gridx = gridx.reshape(1, self.ref, 1, 1).repeat([batchsize, 1, self.ref, 1])
+        gridy = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float)
+        gridy = gridy.reshape(1, 1, self.ref, 1).repeat([batchsize, self.ref, 1, 1])
+        
+        grid_ref = torch.cat((gridx, gridy), dim=-1).to(my_pos.device).reshape(batchsize, self.ref ** 2, 2)  # B Ref^2 2
+
+        pos = torch.sqrt(
+            torch.sum((my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2,
+                      dim=-1)). \
+            reshape(batchsize, my_pos.shape[1], self.ref * self.ref).contiguous()
+        return pos
+
+    def forward(self, data):
+        x, fx, T = data.x, None, None
+        x = x[None, :, :] # [1, N, C]
+        
+        if self.unified_pos:
+            new_pos = self.get_grid(data.pos[None, :, :])
+            x = torch.cat((x, new_pos), dim=-1)
+            
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+            fx = fx + self.placeholder[None, None, :]
+
+        for block in self.blocks:
+            fx = block(fx)
+
+        return fx[0]

model/Transolver2D_plus.py

@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+import numpy as np
+try:
+    from timm.layers import trunc_normal_
+except ImportError:
+    from torch.nn.init import trunc_normal_
+
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.Transolver_block import Transolver_block
+
+class Transolver2D_plus(nn.Module):
+    """
+    Transolver2D_plus 模型。
+
+    这是 Transolver2D 的增强版本。
+    主要的区别在于它使用了 'unstructured_plus' 类型的注意力机制，这种机制通常包含更复杂的切片或特征交互逻辑（如 Gumbel Softmax 采样等），以增强模型对不规则网格的捕捉能力。
+
+    Args:
+        space_dim (int): 空间维度。默认值: 1。
+        n_layers (int): Block 层数。默认值: 5。
+        n_hidden (int): 隐藏层维度。默认值: 256。
+        dropout (float): Dropout 概率。默认值: 0。
+        n_head (int): 注意力头数。默认值: 8。
+        act (str): 激活函数。默认值: 'gelu'。
+        mlp_ratio (float): MLP 膨胀比率。默认值: 1。
+        fun_dim (int): 输入物理场特征维度。默认值: 1。
+        out_dim (int): 输出特征维度。默认值: 1。
+        slice_num (int): 切片数量。默认值: 32。
+        ref (int): 统一位置编码参考分辨率。默认值: 8。
+        unified_pos (bool): 是否使用统一位置编码。默认值: False。
+    """
+    def __init__(self,
+                 space_dim=1,
+                 n_layers=5,
+                 n_hidden=256,
+                 dropout=0,
+                 n_head=8,
+                 act='gelu',
+                 mlp_ratio=1,
+                 fun_dim=1,
+                 out_dim=1,
+                 slice_num=32,
+                 ref=8,
+                 unified_pos=False
+                 ):
+        super(Transolver2D_plus, self).__init__()
+        self.__name__ = 'Transolver'
+        self.ref = ref
+        self.unified_pos = unified_pos
+        
+        if self.unified_pos:
+            input_dim = fun_dim + space_dim + self.ref * self.ref
+        else:
+            input_dim = fun_dim + space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            hidden_dims=[n_hidden * 2], 
+            output_dim=n_hidden,
+            activation=act,
+            use_bias=True,
+            use_skip_connection=False
+        )
+
+        self.n_hidden = n_hidden
+        self.space_dim = space_dim
+
+        self.blocks = nn.ModuleList([
+            Transolver_block(
+                num_heads=n_head,
+                hidden_dim=n_hidden,
+                dropout=dropout,
+                act=act,
+                mlp_ratio=mlp_ratio,
+                out_dim=out_dim,
+                slice_num=slice_num,
+                last_layer=(_ == n_layers - 1),
+                geotype='unstructured_plus' 
+            )
+            for _ in range(n_layers)
+        ])
+        
+        self.initialize_weights()
+        self.placeholder = nn.Parameter((1 / (n_hidden)) * torch.rand(n_hidden, dtype=torch.float))
+
+    def initialize_weights(self):
+        self.apply(self._init_weights)
+
+    def _init_weights(self, m):
+        if isinstance(m, nn.Linear):
+            trunc_normal_(m.weight, std=0.02)
+            if isinstance(m, nn.Linear) and m.bias is not None:
+                nn.init.constant_(m.bias, 0)
+        elif isinstance(m, (nn.LayerNorm, nn.BatchNorm1d)):
+            nn.init.constant_(m.bias, 0)
+            nn.init.constant_(m.weight, 1.0)
+
+    def get_grid(self, my_pos):
+        batchsize = my_pos.shape[0]
+
+        gridx = torch.tensor(np.linspace(-2, 4, self.ref), dtype=torch.float)
+        gridx = gridx.reshape(1, self.ref, 1, 1).repeat([batchsize, 1, self.ref, 1])
+        gridy = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float)
+        gridy = gridy.reshape(1, 1, self.ref, 1).repeat([batchsize, self.ref, 1, 1])
+        
+        grid_ref = torch.cat((gridx, gridy), dim=-1).to(my_pos.device).reshape(batchsize, self.ref ** 2, 2)
+
+        pos = torch.sqrt(
+            torch.sum((my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2,
+                      dim=-1)). \
+            reshape(batchsize, my_pos.shape[1], self.ref * self.ref).contiguous()
+        return pos
+
+    def forward(self, data):
+        x, fx, T = data.x, None, None
+        x = x[None, :, :] 
+        
+        if self.unified_pos:
+            new_pos = self.get_grid(data.pos[None, :, :])
+            x = torch.cat((x, new_pos), dim=-1)
+            
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+            fx = fx + self.placeholder[None, None, :]
+
+        for block in self.blocks:
+            fx = block(fx)
+
+        return fx[0]

model/__init__.py

@@ -0,0 +1,4 @@
+from .Transolver2D import Transolver2D
+from .Transolver2D_plus import Transolver2D_plus
+
+__all__ = ["Transolver2D", "Transolver2D_plus"]

scripts/fake_data.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+import json
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+from torch_geometric.data import Data
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from onescience.utils.YParams import YParams
+
+
+def resolve_path(path: str) -> Path:
+    value = Path(path).expanduser()
+    return value if value.is_absolute() else ROOT / value
+
+
+def make_sample(num_nodes: int, input_dim: int, output_dim: int, surface_ratio: float, seed: int) -> Data:
+    generator = torch.Generator().manual_seed(seed)
+    pos = torch.rand((num_nodes, 2), generator=generator)
+    freestream = torch.ones((num_nodes, 2))
+    freestream[:, 1] = torch.linspace(-0.2, 0.2, num_nodes)
+    sdf = torch.linalg.norm(pos - 0.5, dim=1, keepdim=True) - 0.25
+    normals = torch.nn.functional.normalize(pos - 0.5, dim=1)
+    x = torch.cat([pos, freestream, sdf, normals], dim=1)
+    if x.size(1) != input_dim:
+        raise ValueError(f"fake input_dim must be 7 for this AirfRANS feature layout, got {input_dim}")
+
+    y0 = torch.sin(np.pi * pos[:, :1]) + 0.1 * freestream[:, :1]
+    y1 = torch.cos(np.pi * pos[:, 1:2]) + 0.1 * freestream[:, 1:2]
+    pressure = pos[:, :1] - pos[:, 1:2]
+    viscosity = torch.abs(sdf) + 0.01
+    y = torch.cat([y0, y1, pressure, viscosity], dim=1)
+    if y.size(1) != output_dim:
+        raise ValueError(f"fake output_dim must be 4 for this AirfRANS target layout, got {output_dim}")
+
+    surf = torch.zeros(num_nodes, dtype=torch.bool)
+    surf_count = max(1, int(num_nodes * surface_ratio))
+    surf[:surf_count] = True
+    perm = torch.randperm(num_nodes, generator=generator)
+    return Data(pos=pos[perm], x=x[perm].float(), y=y[perm].float(), surf=surf[perm])
+
+
+def main() -> int:
+    cfg = YParams(str(ROOT / "conf" / "config.yaml"), "fake_data")
+    cfg_data = YParams(str(ROOT / "conf" / "config.yaml"), "datapipe")
+    cfg_train = YParams(str(ROOT / "conf" / "config.yaml"), "training")
+    data_dir = resolve_path(cfg_data.source.data_dir)
+    stats_dir = resolve_path(cfg_data.source.stats_dir)
+    data_dir.mkdir(parents=True, exist_ok=True)
+    stats_dir.mkdir(parents=True, exist_ok=True)
+
+    split_counts = {
+        cfg_data.data.splits.train_name: int(cfg.num_train),
+        cfg_data.data.splits.val_name: int(cfg.num_val),
+        cfg_data.data.splits.test_name: int(cfg.num_test),
+    }
+    manifest: dict[str, list[str]] = {}
+    all_x, all_y = [], []
+    seed = int(cfg_train.seed)
+
+    for split, count in split_counts.items():
+        names = []
+        for idx in range(count):
+            name = f"{split}_{idx:04d}"
+            sample = make_sample(
+                num_nodes=int(cfg.num_nodes),
+                input_dim=int(cfg.input_dim),
+                output_dim=int(cfg.output_dim),
+                surface_ratio=float(cfg.surface_ratio),
+                seed=seed + len(all_x),
+            )
+            torch.save(sample, data_dir / f"{name}.pt")
+            names.append(name)
+            if split == cfg_data.data.splits.train_name:
+                all_x.append(sample.x.numpy())
+                all_y.append(sample.y.numpy())
+        manifest[split] = names
+
+    (data_dir / "manifest.json").write_text(json.dumps(manifest, indent=2), encoding="utf-8")
+    train_x = np.concatenate(all_x, axis=0)
+    train_y = np.concatenate(all_y, axis=0)
+    np.save(stats_dir / "mean_in.npy", train_x.mean(axis=0).astype(np.float32))
+    np.save(stats_dir / "std_in.npy", (train_x.std(axis=0) + 1e-6).astype(np.float32))
+    np.save(stats_dir / "mean_out.npy", train_y.mean(axis=0).astype(np.float32))
+    np.save(stats_dir / "std_out.npy", (train_y.std(axis=0) + 1e-6).astype(np.float32))
+    print(f"Fake AirfRANS graph data written to {data_dir}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/inference.py

@@ -0,0 +1,177 @@
+from __future__ import annotations
+
+import json
+import logging
+import random
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+
+from train import build_datapipe, build_model, resolve_path
+
+
+def setup_logging() -> logging.Logger:
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    return logging.getLogger("transolver-inference")
+
+
+def load_checkpoint(path: Path, device: torch.device):
+    if not path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {path}. Run scripts/train.py first or provide a compatible weight.")
+    try:
+        return torch.load(path, map_location=device, weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location=device)
+
+
+def run_fake_inference(cfg_model, cfg_data, cfg_train, logger: logging.Logger, device: torch.device) -> int:
+    model_params = cfg_model.specific_params[cfg_model.name]
+    datapipe = build_datapipe(cfg_data, model_params, distributed=False)
+    test_loader = datapipe._loader(datapipe.test_dataset, False)[0]
+
+    model = build_model(cfg_model).to(device)
+    checkpoint_path = resolve_path(cfg_train.checkpoint_dir) / f"{cfg_model.name}.pth"
+    checkpoint = load_checkpoint(checkpoint_path, device)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    result_dir = resolve_path(cfg_train.result_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    mean_out, std_out = datapipe.coef_norm[2], datapipe.coef_norm[3]
+    preds, targets = [], []
+    criterion = nn.MSELoss()
+
+    with torch.no_grad():
+        for data in test_loader:
+            data = data.to(device)
+            out = model(data)
+            pred = out.cpu().numpy() * (std_out + 1e-8) + mean_out
+            target = data.y.cpu().numpy() * (std_out + 1e-8) + mean_out
+            preds.append(pred)
+            targets.append(target)
+
+    pred_arr = np.concatenate(preds, axis=0)
+    target_arr = np.concatenate(targets, axis=0)
+    mse = float(criterion(torch.from_numpy(pred_arr), torch.from_numpy(target_arr)))
+    np.savez(result_dir / "predictions.npz", pred=pred_arr, target=target_arr, mse=np.array(mse, dtype=np.float32))
+    (result_dir / "score.json").write_text(json.dumps({"model_name": cfg_model.name, "mse": mse}, indent=2), encoding="utf-8")
+    logger.info("Saved fake inference outputs to %s", result_dir)
+    logger.info("MSE: %.6f", mse)
+    return 0
+
+
+def run_airfrans_inference(cfg_model, cfg_data, cfg_train, logger: logging.Logger, device: torch.device) -> int:
+    import pyvista as pv
+    import scipy.stats as sc
+    from tqdm import tqdm
+    from onescience.utils.transolver.metrics import (
+        Airfoil_test,
+        Compute_coefficients,
+        Infer_test,
+        NumpyEncoder,
+        rel_err,
+    )
+
+    model_params = cfg_model.specific_params[cfg_model.name]
+    hparams = model_params
+    hparams["subsampling"] = cfg_data.data.subsampling
+    datapipe = build_datapipe(cfg_data, model_params, distributed=False)
+    coef_norm = datapipe.coef_norm
+    test_loader = datapipe.test_dataloader()
+    test_names = datapipe.test_dataset.data_list_names
+
+    model = build_model(cfg_model).to(device)
+    checkpoint_path = resolve_path(cfg_train.checkpoint_dir) / f"{cfg_model.name}.pth"
+    checkpoint = load_checkpoint(checkpoint_path, device)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    path_in = resolve_path(cfg_data.source.data_dir)
+    path_out = resolve_path(cfg_train.result_dir) / cfg_data.data.splits.task
+    path_out.mkdir(parents=True, exist_ok=True)
+    sample_count = min(int(cfg_train.n_test), len(test_names))
+    chosen = sorted(random.sample(range(len(test_names)), k=sample_count)) if sample_count else []
+
+    criterion = nn.MSELoss(reduction="none")
+    scores_vol, scores_surf, scores_force, scores_p, scores_wss = [], [], [], [], []
+    times, true_coefs, pred_coefs = [], [], []
+
+    for j, data in enumerate(tqdm(test_loader, desc="Testing")):
+        sim_name = test_names[j]
+        u_inf, angle = float(sim_name.split("_")[2]), float(sim_name.split("_")[3])
+        outs, elapsed = Infer_test(device, [model], [hparams], data, coef_norm=coef_norm)
+        times.append(elapsed)
+
+        intern = pv.read(path_in / sim_name / f"{sim_name}_internal.vtu")
+        aerofoil = pv.read(path_in / sim_name / f"{sim_name}_aerofoil.vtp")
+        tc, true_intern, true_airfoil = Compute_coefficients([intern], [aerofoil], data.surf.cpu(), u_inf, angle, keep_vtk=True)
+        tc, true_airfoil = tc[0], true_airfoil[0]
+        intern_pred, aerofoil_pred = Airfoil_test(intern, aerofoil, outs, coef_norm, data.surf.cpu())
+        pc, intern_pred_vtk, aerofoil_pred_vtk = Compute_coefficients(intern_pred, aerofoil_pred, data.surf.cpu(), u_inf, angle, keep_vtk=True)
+        true_coefs.append(tc)
+        pred_coefs.append(pc)
+
+        out = outs[0]
+        scores_vol.append(criterion(out[~data.surf], data.y[~data.surf]).mean(dim=0).cpu().numpy())
+        scores_surf.append(criterion(out[data.surf], data.y[data.surf]).mean(dim=0).cpu().numpy())
+        scores_force.append(rel_err(tc, pc[0]))
+        scores_wss.append(rel_err(true_airfoil.point_data["wallShearStress"], aerofoil_pred_vtk[0].point_data["wallShearStress"]).mean(axis=0))
+        scores_p.append(rel_err(true_airfoil.point_data["p"], aerofoil_pred_vtk[0].point_data["p"]).mean(axis=0))
+        if j in chosen:
+            intern_pred_vtk[0].save(path_out / f"{sim_name}_pred_internal.vtu")
+            aerofoil_pred_vtk[0].save(path_out / f"{sim_name}_pred_aerofoil.vtp")
+
+    true_coefs_arr = np.array(true_coefs)
+    pred_coefs_arr = np.array(pred_coefs)
+    spear = [
+        sc.spearmanr(true_coefs_arr[:, 0], pred_coefs_arr[:, 0, 0])[0],
+        sc.spearmanr(true_coefs_arr[:, 1], pred_coefs_arr[:, 0, 1])[0],
+    ]
+    score = {
+        "model_name": cfg_model.name,
+        "mean_time": np.array(times).mean(axis=0),
+        "std_time": np.array(times).std(axis=0),
+        "mean_score_vol": np.array(scores_vol).mean(axis=0),
+        "mean_score_surf": np.array(scores_surf).mean(axis=0),
+        "mean_score_force": np.array(scores_force).mean(axis=0),
+        "mean_rel_p": np.array(scores_p).mean(axis=0),
+        "mean_rel_wss": np.array(scores_wss).mean(axis=0),
+        "spearman_coef": np.array(spear),
+    }
+    (path_out / f"score_{cfg_model.name}.json").write_text(json.dumps(score, indent=2, cls=NumpyEncoder), encoding="utf-8")
+    logger.info("Saved AirfRANS inference outputs to %s", path_out)
+    return 0
+
+
+def main() -> int:
+    config_path = ROOT / "conf" / "config.yaml"
+    cfg_model = YParams(str(config_path), "model")
+    cfg_data = YParams(str(config_path), "datapipe")
+    cfg_train = YParams(str(config_path), "training")
+    logger = setup_logging()
+    DistributedManager.initialize()
+    device_name = getattr(cfg_train, "device", "auto")
+    if device_name == "cpu":
+        device = torch.device("cpu")
+    elif device_name.startswith("cuda"):
+        device = torch.device(device_name if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(f"cuda:{cfg_train.gpuid}" if torch.cuda.is_available() else "cpu")
+    if cfg_data.backend == "fake_airfrans":
+        return run_fake_inference(cfg_model, cfg_data, cfg_train, logger, device)
+    if cfg_data.backend == "airfrans":
+        return run_airfrans_inference(cfg_model, cfg_data, cfg_train, logger, device)
+    raise ValueError(f"Unsupported datapipe.backend: {cfg_data.backend}")
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/result.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parents[1]
+RESULT_DIR = ROOT / "result"
+
+
+def main() -> int:
+    score_path = RESULT_DIR / "score.json"
+    pred_path = RESULT_DIR / "predictions.npz"
+    if score_path.exists():
+        print(score_path.read_text(encoding="utf-8"))
+        return 0
+    if not pred_path.exists():
+        raise FileNotFoundError(f"No result found at {score_path} or {pred_path}. Run scripts/inference.py first.")
+    data = np.load(pred_path)
+    mse = float(np.mean((data["pred"] - data["target"]) ** 2))
+    summary = {"mse": mse, "num_nodes": int(data["pred"].shape[0])}
+    print(json.dumps(summary, indent=2))
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/train.py

@@ -0,0 +1,270 @@
+from __future__ import annotations
+
+import logging
+import random
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+from torch.utils.data.distributed import DistributedSampler
+from torch_geometric.loader import DataLoader as PyGDataLoader
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+
+from model import Transolver2D, Transolver2D_plus
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+
+
+def resolve_path(path: str) -> Path:
+    value = Path(path).expanduser()
+    return value if value.is_absolute() else ROOT / value
+
+
+def setup_logging(rank: int) -> logging.Logger:
+    logging.basicConfig(
+        level=logging.INFO if rank == 0 else logging.WARNING,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    return logging.getLogger("transolver-train")
+
+
+def load_graph(path: Path):
+    try:
+        return torch.load(path, weights_only=False)
+    except TypeError:
+        return torch.load(path)
+
+
+class FakeAirfRANSDataset(torch.utils.data.Dataset):
+    def __init__(self, data_dir: Path, stats_dir: Path, split: str, coef_norm=None):
+        manifest_path = data_dir / "manifest.json"
+        if not manifest_path.exists():
+            raise FileNotFoundError(f"Fake manifest not found: {manifest_path}. Run scripts/fake_data.py first.")
+        import json
+
+        manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
+        self.data_dir = data_dir
+        self.names = manifest[split]
+        self.coef_norm = coef_norm or (
+            np.load(stats_dir / "mean_in.npy"),
+            np.load(stats_dir / "std_in.npy"),
+            np.load(stats_dir / "mean_out.npy"),
+            np.load(stats_dir / "std_out.npy"),
+        )
+
+    def __len__(self) -> int:
+        return len(self.names)
+
+    def __getitem__(self, index: int):
+        data = load_graph(self.data_dir / f"{self.names[index]}.pt")
+        mean_in, std_in, mean_out, std_out = self.coef_norm
+        data.x = (data.x - torch.as_tensor(mean_in, dtype=data.x.dtype)) / torch.as_tensor(std_in + 1e-8, dtype=data.x.dtype)
+        data.y = (data.y - torch.as_tensor(mean_out, dtype=data.y.dtype)) / torch.as_tensor(std_out + 1e-8, dtype=data.y.dtype)
+        return data
+
+
+class FakeAirfRANSDatapipe:
+    def __init__(self, cfg_data, distributed: bool):
+        self.cfg_data = cfg_data
+        self.distributed = distributed
+        data_dir = resolve_path(cfg_data.source.data_dir)
+        stats_dir = resolve_path(cfg_data.source.stats_dir)
+        self.train_dataset = FakeAirfRANSDataset(data_dir, stats_dir, cfg_data.data.splits.train_name)
+        self.coef_norm = self.train_dataset.coef_norm
+        self.val_dataset = FakeAirfRANSDataset(data_dir, stats_dir, cfg_data.data.splits.val_name, self.coef_norm)
+        self.test_dataset = FakeAirfRANSDataset(data_dir, stats_dir, cfg_data.data.splits.test_name, self.coef_norm)
+
+    def _loader(self, dataset, shuffle: bool):
+        sampler = DistributedSampler(dataset, shuffle=shuffle) if self.distributed else None
+        return (
+            PyGDataLoader(
+                dataset,
+                batch_size=self.cfg_data.dataloader.batch_size,
+                num_workers=self.cfg_data.dataloader.num_workers,
+                pin_memory=torch.cuda.is_available(),
+                shuffle=shuffle and sampler is None,
+                sampler=sampler,
+            ),
+            sampler,
+        )
+
+    def train_dataloader(self):
+        return self._loader(self.train_dataset, True)
+
+    def val_dataloader(self):
+        return self._loader(self.val_dataset, False)
+
+
+def build_model(cfg_model):
+    model_name = cfg_model.name
+    if model_name not in cfg_model.specific_params:
+        raise ValueError(f"Model '{model_name}' is missing from model.specific_params")
+    params = cfg_model.specific_params[model_name]
+    model_cls = Transolver2D_plus if model_name == "Transolver_plus" else Transolver2D
+    if model_name not in ("Transolver", "Transolver_plus"):
+        raise NotImplementedError("This refactored package localizes only Transolver and Transolver_plus.")
+    return model_cls(
+        n_hidden=params.n_hidden,
+        n_layers=params.n_layers,
+        space_dim=params.space_dim,
+        fun_dim=params.fun_dim,
+        n_head=params.n_head,
+        mlp_ratio=params.mlp_ratio,
+        out_dim=params.out_dim,
+        slice_num=params.slice_num,
+        unified_pos=bool(params.unified_pos),
+    )
+
+
+def build_datapipe(cfg_data, model_params, distributed: bool):
+    cfg_data.model_hparams = model_params
+    if cfg_data.backend == "airfrans":
+        from onescience.datapipes.cfd import AirfRANSDatapipe
+
+        return AirfRANSDatapipe(params=cfg_data, distributed=distributed)
+    if cfg_data.backend == "fake_airfrans":
+        return FakeAirfRANSDatapipe(cfg_data, distributed=distributed)
+    raise ValueError(f"Unsupported datapipe.backend: {cfg_data.backend}")
+
+
+def masked_mean(losses: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
+    if mask.any():
+        return losses[mask].mean()
+    return losses.mean() * 0.0
+
+
+def save_checkpoint(model, optimizer, scheduler, epoch: int, loss: float, checkpoint_dir: Path, model_name: str) -> None:
+    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(),
+            "scheduler_state_dict": scheduler.state_dict(),
+            "epoch": epoch,
+            "loss": loss,
+        },
+        checkpoint_dir / f"{model_name}.pth",
+    )
+
+
+def main() -> int:
+    config_path = ROOT / "conf" / "config.yaml"
+    cfg_model = YParams(str(config_path), "model")
+    cfg_data = YParams(str(config_path), "datapipe")
+    cfg_train = YParams(str(config_path), "training")
+
+    torch.manual_seed(int(cfg_train.seed))
+    np.random.seed(int(cfg_train.seed))
+    random.seed(int(cfg_train.seed))
+
+    DistributedManager.initialize()
+    manager = DistributedManager()
+    logger = setup_logging(manager.rank)
+    device_name = getattr(cfg_train, "device", "auto")
+    if device_name == "cpu":
+        device = torch.device("cpu")
+    elif device_name.startswith("cuda"):
+        device = torch.device(device_name if torch.cuda.is_available() else "cpu")
+    else:
+        device = torch.device(f"cuda:{cfg_train.gpuid}" if torch.cuda.is_available() and manager.world_size == 1 else manager.device)
+
+    model_params = cfg_model.specific_params[cfg_model.name]
+    datapipe = build_datapipe(cfg_data, model_params, distributed=(manager.world_size > 1))
+    train_loader, train_sampler = datapipe.train_dataloader()
+    val_loader, val_sampler = datapipe.val_dataloader()
+    if len(train_loader) == 0 or len(val_loader) == 0:
+        raise RuntimeError("Train and validation loaders must both be non-empty.")
+
+    model = build_model(cfg_model).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))
+    scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer,
+        max_lr=float(cfg_train.lr),
+        total_steps=max(1, len(train_loader) * int(cfg_train.max_epoch)),
+    )
+    criterion = nn.L1Loss(reduction="none") if cfg_train.loss_criterion == "MAE" else nn.MSELoss(reduction="none")
+    use_weighted_loss = cfg_train.loss_criterion == "MSE_weighted"
+    checkpoint_dir = resolve_path(cfg_train.checkpoint_dir)
+    best_valid_loss = float("inf")
+    best_epoch = -1
+
+    if manager.rank == 0:
+        params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        logger.info("Training %s with %.3fM parameters on %s", cfg_model.name, params / 1e6, device)
+
+    for epoch in range(int(cfg_train.max_epoch)):
+        start = time.time()
+        if train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+        if val_sampler is not None:
+            val_sampler.set_epoch(epoch)
+
+        model.train()
+        train_loss = 0.0
+        for data in train_loader:
+            data = data.to(device)
+            optimizer.zero_grad(set_to_none=True)
+            out = model(data)
+            losses = criterion(out, data.y)
+            surf_loss = masked_mean(losses, data.surf)
+            vol_loss = masked_mean(losses, ~data.surf)
+            loss = vol_loss + float(cfg_train.loss_weight) * surf_loss if use_weighted_loss else losses.mean()
+            loss.backward()
+            optimizer.step()
+            scheduler.step()
+            train_loss += float(loss.detach().cpu())
+        train_loss /= len(train_loader)
+
+        model.eval()
+        valid_loss = 0.0
+        with torch.no_grad():
+            for data in val_loader:
+                data = data.to(device)
+                out = model(data)
+                losses = criterion(out, data.y)
+                surf_loss = masked_mean(losses, data.surf)
+                vol_loss = masked_mean(losses, ~data.surf)
+                loss = vol_loss + float(cfg_train.loss_weight) * surf_loss if use_weighted_loss else losses.mean()
+                if manager.world_size > 1:
+                    dist.all_reduce(loss, op=dist.ReduceOp.AVG)
+                valid_loss += float(loss.detach().cpu())
+        valid_loss /= len(val_loader)
+
+        if manager.rank == 0:
+            logger.info(
+                "Epoch %d/%d train=%.6f valid=%.6f time=%.2fs",
+                epoch + 1,
+                int(cfg_train.max_epoch),
+                train_loss,
+                valid_loss,
+                time.time() - start,
+            )
+            if valid_loss < best_valid_loss:
+                best_valid_loss = valid_loss
+                best_epoch = epoch
+                save_checkpoint(model, optimizer, scheduler, epoch, valid_loss, checkpoint_dir, cfg_model.name)
+                logger.info("Saved checkpoint to %s", checkpoint_dir / f"{cfg_model.name}.pth")
+            if epoch - best_epoch > int(cfg_train.patience):
+                logger.info("Early stopping after %d epochs without improvement.", int(cfg_train.patience))
+                break
+
+    if manager.world_size > 1 and dist.is_available() and dist.is_initialized():
+        dist.barrier()
+        dist.destroy_process_group()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())