Upload folder using huggingface_hub

README.md

@@ -0,0 +1,125 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- fluid dynamics
+- automotive aerodynamic design
+- CFD surrogate modeling
+frameworks: PyTorch
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Transolver-Car-Design</span>
+  </strong>
+</p>
+
+# Model Overview
+
+Transolver-Car-Design is a three-dimensional automotive external-flow prediction model built on Transolver and Transolver++ by Tsinghua University's THUML group. It provides surrogate modeling of vehicle flow fields and rapid prediction of drag coefficients.
+
+Paper: [Transolver: A Fast Transformer Solver for PDEs on General Geometries](https://arxiv.org/pdf/2402.02366)
+
+# Model Description
+Transolver-Car-Design uses a Transformer architecture with Physics-Attention and is trained on ShapeNet-Car automotive aerodynamic simulation data. It predicts velocity fields, pressure fields, and drag coefficients for complex vehicle geometries.
+
+## Use Cases
+
+| Use Case | Description |
+| :--- | :--- |
+| Automotive aerodynamic design | Rapidly predict external-flow velocities and surface pressures around vehicles |
+| CFD surrogate modeling | Approximate fluid simulation on complex unstructured meshes with a neural network |
+| Simulation acceleration | Provide a lightweight evaluation pipeline for large-scale candidate design screening |
+
+# 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-Car-Design --local_dir ./Transolver-Car-Design
+cd Transolver-Car-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 `ShapeNetCar` 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/ShapeNetCar --local_dir ./data
+```
+
+### Training
+
+```bash
+python scripts/train.py
+```
+
+Training saves `Transolver_plus.pth` under `weight/`:
+
+```text
+./weight/Transolver_plus.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).
+- Original Transolver++ paper: [Transolver++: An Accurate Neural Solver for PDEs on Million-Scale Geometries](https://arxiv.org/abs/2502.02414)
+- This repository retains source attribution and has been adapted for automated execution through OneScience and ModelScope.

conf/config.yaml

@@ -0,0 +1,49 @@
+model:
+  name: "Transolver_plus"
+  specific_params:
+    Transolver: &params
+      n_hidden: 16
+      n_layers: 1
+      space_dim: 7
+      fun_dim: 0
+      n_head: 2
+      mlp_ratio: 2
+      out_dim: 4
+      slice_num: 4
+      unified_pos: 0
+      cfd_mesh: true
+      r: 0.2
+    Transolver_plus: *params
+
+datapipe:
+  name: "ShapeNetCar"
+  task: "Transolver-Car-Design"
+  verbose: true
+  source:
+    data_dir: "./data/mlcfd_data/training_data"
+    preprocessed_save_dir: "./data/mlcfd_data/preprocessed_data"
+    stats_dir: "./data/mlcfd_data/stats"
+    preprocessed: 1
+  data:
+    splits:
+      fold_id: 0
+  dataloader:
+    batch_size: 1
+    num_workers: 0
+    pin_memory: false
+
+training:
+  gpuid: -1
+  max_epoch: 1
+  patience: 50
+  lr: 0.001
+  loss_criterion: "MSE"
+  loss_weight: 0.5
+  checkpoint_dir: "./weight"
+  val_iter: 1
+
+inference:
+  gpuid: -1
+  result_dir: "./results/ShapeNetCar"
+  save_vtk: false
+  visualize: false

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other"}

model/Transolver3D.py

@@ -0,0 +1,139 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from timm.layers import trunc_normal_
+from einops import rearrange, repeat
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.Transolver_block import Transolver_block
+
+class Transolver3D(nn.Module):
+    """
+    Transolver3D 模型。
+
+    该模型专为处理三维物理场问题（如 CFD）而设计。它利用 TransolverBlock 堆叠而成的深层网络来捕捉复杂的物理依赖关系。
+    模型首先通过一个预处理 MLP 将输入的物理状态（和可选的统一位置编码）映射到隐空间，然后经过多层 Transolver Block 进行特征提取和交互，
+    最后输出预测的物理场。
+
+    Args:
+        space_dim (int): 空间维度（例如 1, 2, 3）。默认值: 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(Transolver3D, self).__init__()
+        self.__name__ = 'Transolver3D'
+        self.ref = ref
+        self.unified_pos = unified_pos
+
+        if self.unified_pos:
+            input_dim = fun_dim + self.ref * 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):
+        # my_pos 1 N 3
+        batchsize = my_pos.shape[0]
+
+        gridx = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float)
+        gridx = gridx.reshape(1, self.ref, 1, 1, 1).repeat([batchsize, 1, self.ref, self.ref, 1])
+        gridy = torch.tensor(np.linspace(0, 2, self.ref), dtype=torch.float)
+        gridy = gridy.reshape(1, 1, self.ref, 1, 1).repeat([batchsize, self.ref, 1, self.ref, 1])
+        gridz = torch.tensor(np.linspace(-4, 4, self.ref), dtype=torch.float)
+        gridz = gridz.reshape(1, 1, 1, self.ref, 1).repeat([batchsize, self.ref, self.ref, 1, 1])
+        
+        grid_ref = torch.cat((gridx, gridy, gridz), dim=-1).to(my_pos.device).reshape(batchsize, self.ref ** 3, 3)  # B 4 4 4 3
+
+        pos = torch.sqrt(
+            torch.sum((my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2,
+                      dim=-1)). \
+            reshape(batchsize, my_pos.shape[1], self.ref * self.ref * self.ref).contiguous()
+        return pos
+
+    def forward(self, data):
+        cfd_data = data
+        x, fx, T = cfd_data.x, None, None
+        x = x[None, :, :] # [1, N, C]
+        
+        if self.unified_pos:
+            new_pos = self.get_grid(cfd_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/Transolver3D_plus.py

@@ -0,0 +1,135 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from timm.layers import trunc_normal_
+from einops import rearrange, repeat
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.Transolver_block import Transolver_block
+
+class Transolver3D_plus(nn.Module):
+    """
+    Transolver3D_plus 模型。
+
+    这是 Transolver3D 的增强版本。
+    它使用了 '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(Transolver3D_plus, self).__init__()
+        self.__name__ = 'Transolver3D'
+        self.ref = ref
+        self.unified_pos = unified_pos
+        
+        if self.unified_pos:
+            input_dim = fun_dim + self.ref * 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]
+
+        # 3D Grid Generation Logic
+        gridx = torch.tensor(np.linspace(-1.5, 1.5, self.ref), dtype=torch.float)
+        gridx = gridx.reshape(1, self.ref, 1, 1, 1).repeat([batchsize, 1, self.ref, self.ref, 1])
+        gridy = torch.tensor(np.linspace(0, 2, self.ref), dtype=torch.float)
+        gridy = gridy.reshape(1, 1, self.ref, 1, 1).repeat([batchsize, self.ref, 1, self.ref, 1])
+        gridz = torch.tensor(np.linspace(-4, 4, self.ref), dtype=torch.float)
+        gridz = gridz.reshape(1, 1, 1, self.ref, 1).repeat([batchsize, self.ref, self.ref, 1, 1])
+        
+        grid_ref = torch.cat((gridx, gridy, gridz), dim=-1).to(my_pos.device).reshape(batchsize, self.ref ** 3, 3)
+
+        pos = torch.sqrt(
+            torch.sum((my_pos[:, :, None, :] - grid_ref[:, None, :, :]) ** 2,
+                      dim=-1)). \
+            reshape(batchsize, my_pos.shape[1], self.ref * self.ref * self.ref).contiguous()
+        return pos
+
+    def forward(self, data):
+        cfd_data = data
+        x, fx, T = cfd_data.x, None, None
+        x = x[None, :, :]
+        
+        if self.unified_pos:
+            new_pos = self.get_grid(cfd_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 .Transolver3D import Transolver3D
+from .Transolver3D_plus import Transolver3D_plus
+
+__all__ = ["Transolver3D", "Transolver3D_plus"]

scripts/fake_data.py

@@ -0,0 +1,97 @@
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+os.chdir(ROOT)
+
+from onescience.utils.YParams import YParams
+
+
+def make_edges(num_nodes: int) -> np.ndarray:
+    src = np.arange(num_nodes, dtype=np.int64)
+    dst = (src + 1) % num_nodes
+    edges = np.stack(
+        [np.concatenate([src, dst]), np.concatenate([dst, src])],
+        axis=0,
+    )
+    return edges
+
+
+def write_placeholder_training_files(param_dir: Path, rng: np.random.Generator) -> None:
+    param_dir.mkdir(parents=True, exist_ok=True)
+    np.save(param_dir / "Cd.npy", rng.normal(size=(1,)).astype(np.float32))
+    np.save(param_dir / "I1.npy", rng.normal(size=(8, 3)).astype(np.float32))
+    np.save(param_dir / "I2.npy", rng.normal(size=(8, 3)).astype(np.float32))
+    np.save(param_dir / "Press.npy", rng.normal(size=(8, 1)).astype(np.float32))
+    np.save(param_dir / "Velo.npy", rng.normal(size=(8, 3)).astype(np.float32))
+
+
+def write_sample(sample_dir: Path, num_nodes: int, rng: np.random.Generator) -> tuple[np.ndarray, np.ndarray]:
+    sample_dir.mkdir(parents=True, exist_ok=True)
+
+    pos = rng.uniform(low=-1.0, high=1.0, size=(num_nodes, 3)).astype(np.float32)
+    sdf = rng.uniform(low=0.0, high=0.2, size=(num_nodes, 1)).astype(np.float32)
+    normals = rng.normal(size=(num_nodes, 3)).astype(np.float32)
+    normals /= np.linalg.norm(normals, axis=1, keepdims=True) + 1e-6
+    x = np.concatenate([pos, sdf, normals], axis=1).astype(np.float32)
+
+    y = np.concatenate(
+        [
+            0.1 * pos + rng.normal(scale=0.01, size=(num_nodes, 3)),
+            rng.normal(scale=0.05, size=(num_nodes, 1)),
+        ],
+        axis=1,
+    ).astype(np.float32)
+
+    surf = np.zeros((num_nodes,), dtype=np.bool_)
+    surf[num_nodes // 2 :] = True
+    edge_index = make_edges(num_nodes)
+
+    np.save(sample_dir / "x.npy", x)
+    np.save(sample_dir / "y.npy", y)
+    np.save(sample_dir / "pos.npy", pos)
+    np.save(sample_dir / "surf.npy", surf)
+    np.save(sample_dir / "edge_index.npy", edge_index)
+    return x, y
+
+
+def main() -> None:
+    cfg = YParams(str(ROOT / "conf/config.yaml"), "datapipe")
+    data_dir = ROOT / cfg.source.data_dir
+    preprocessed_dir = ROOT / cfg.source.preprocessed_save_dir
+    stats_dir = ROOT / cfg.source.stats_dir
+    stats_dir.mkdir(parents=True, exist_ok=True)
+
+    rng = np.random.default_rng(42)
+    all_x: list[np.ndarray] = []
+    all_y: list[np.ndarray] = []
+
+    for fold_id in range(9):
+        param_name = f"param{fold_id}"
+        param_dir = data_dir / param_name
+        write_placeholder_training_files(param_dir, rng)
+
+        sample_name = "sample_000"
+        (param_dir / sample_name).mkdir(parents=True, exist_ok=True)
+        x, y = write_sample(preprocessed_dir / param_name / sample_name, num_nodes=8, rng=rng)
+        all_x.append(x)
+        all_y.append(y)
+
+    x_all = np.concatenate(all_x, axis=0)
+    y_all = np.concatenate(all_y, axis=0)
+    np.save(stats_dir / "mean_in.npy", x_all.mean(axis=0).astype(np.float32))
+    np.save(stats_dir / "std_in.npy", (x_all.std(axis=0) + 1e-6).astype(np.float32))
+    np.save(stats_dir / "mean_out.npy", y_all.mean(axis=0).astype(np.float32))
+    np.save(stats_dir / "std_out.npy", (y_all.std(axis=0) + 1e-6).astype(np.float32))
+
+    print(f"Fake ShapeNetCar data generated under {ROOT / 'data/mlcfd_data'}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,201 @@
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import time
+import importlib.util
+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))
+os.chdir(ROOT)
+
+from model import Transolver3D, Transolver3D_plus
+import onescience
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+from onescience.utils.transolver import cal_coefficient, save_prediction_to_vtk, visualize_prediction
+
+
+def load_shapenet_car_datapipe():
+    module_path = Path(onescience.__file__).resolve().parent / "datapipes/cfd/ShapeNetCar.py"
+    spec = importlib.util.spec_from_file_location("_onescience_shapenetcar", module_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Unable to load ShapeNetCarDatapipe from {module_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.ShapeNetCarDatapipe
+
+
+def setup_logging(rank: int) -> logging.Logger:
+    level = logging.INFO if rank == 0 else logging.WARNING
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    logging.getLogger().setLevel(level)
+    return logging.getLogger(__name__)
+
+
+def build_model(model_name: str, model_params, device: torch.device) -> torch.nn.Module:
+    model_cls = {
+        "Transolver": Transolver3D,
+        "Transolver_plus": Transolver3D_plus,
+    }.get(model_name)
+    if model_cls is None:
+        raise NotImplementedError(f"Model {model_name} initialization not implemented.")
+    return model_cls(
+        n_hidden=model_params.n_hidden,
+        n_layers=model_params.n_layers,
+        space_dim=model_params.space_dim,
+        fun_dim=model_params.fun_dim,
+        n_head=model_params.n_head,
+        mlp_ratio=model_params.mlp_ratio,
+        out_dim=model_params.out_dim,
+        slice_num=model_params.slice_num,
+        unified_pos=model_params.unified_pos,
+    ).to(device)
+
+
+def resolve_device(gpuid: int) -> torch.device:
+    if torch.cuda.is_available() and int(gpuid) >= 0:
+        return torch.device(f"cuda:{gpuid}")
+    return torch.device("cpu")
+
+
+def maybe_calculate_coefficient(data_dir: Path, pred_press: np.ndarray, pred_velo: np.ndarray, gt_press: np.ndarray, gt_velo: np.ndarray):
+    if not (data_dir / "quadpress_smpl.vtk").exists() or not (data_dir / "hexvelo_smpl.vtk").exists():
+        return None, None
+    pred_coef = cal_coefficient(str(data_dir), pred_press[:, None], pred_velo)
+    gt_coef = cal_coefficient(str(data_dir), gt_press[:, None], gt_velo)
+    return pred_coef, gt_coef
+
+
+def main() -> None:
+    DistributedManager.initialize()
+    manager = DistributedManager()
+    logger = setup_logging(manager.rank)
+    if manager.rank != 0:
+        logger.warning("Inference should run on a single process; exiting non-zero rank.")
+        return
+
+    config_file_path = str(ROOT / "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+    cfg_data = YParams(config_file_path, "datapipe")
+    cfg_train = YParams(config_file_path, "training")
+    cfg_test = YParams(config_file_path, "inference")
+
+    model_name = cfg.name
+    model_params = cfg.specific_params[model_name]
+    cfg_data.model_hparams = model_params
+
+    device = resolve_device(cfg_test.gpuid)
+    logger.info("Using device: %s", device)
+
+    ShapeNetCarDatapipe = load_shapenet_car_datapipe()
+    datapipe = ShapeNetCarDatapipe(params=cfg_data, distributed=False)
+    val_dataset = datapipe.val_dataset
+    coef_norm = datapipe.coef_norm
+    val_names = val_dataset.data_list_names
+    test_loader, _ = datapipe.val_dataloader()
+    logger.info("Loaded %d validation samples.", len(val_dataset))
+
+    model = build_model(model_name, model_params, device)
+    checkpoint_path = Path(cfg_train.checkpoint_dir) / f"{model_name}.pth"
+    if not checkpoint_path.exists():
+        raise FileNotFoundError(f"Checkpoint not found: {checkpoint_path}")
+    checkpoint = torch.load(checkpoint_path, map_location=device)
+    model.load_state_dict(checkpoint["model_state_dict"])
+    model.eval()
+
+    result_root = Path(cfg_test.result_dir) / model_name
+    npy_dir = result_root / "npy"
+    vtk_dir = result_root / "vtk"
+    vis_dir = result_root / "vis"
+    npy_dir.mkdir(parents=True, exist_ok=True)
+    if cfg_test.save_vtk:
+        vtk_dir.mkdir(parents=True, exist_ok=True)
+    if cfg_test.visualize:
+        vis_dir.mkdir(parents=True, exist_ok=True)
+
+    criterion_func = nn.MSELoss(reduction="none")
+    l2errs_press, l2errs_velo, mses_press, mses_velo_var, times = [], [], [], [], []
+    gt_coef_list, pred_coef_list = [], []
+
+    mean = torch.tensor(coef_norm[2], dtype=torch.float32, device=device)
+    std = torch.tensor(coef_norm[3], dtype=torch.float32, device=device)
+
+    with torch.no_grad():
+        for index, data in enumerate(test_loader):
+            if index >= len(val_names):
+                break
+            sample_name = val_names[index]
+            data = data.to(device)
+            tic = time.time()
+            out = model(data)
+            times.append(time.time() - tic)
+            targets = data.y
+
+            pred_press = out[data.surf, -1] * std[-1] + mean[-1]
+            gt_press = targets[data.surf, -1] * std[-1] + mean[-1]
+            pred_velo = out[~data.surf, :-1] * std[:-1] + mean[:-1]
+            gt_velo = targets[~data.surf, :-1] * std[:-1] + mean[:-1]
+            out_denorm = out * std + mean
+            y_denorm = targets * std + mean
+
+            safe_name = sample_name.replace("/", "_")
+            np.save(npy_dir / f"{index}_{safe_name}_pred.npy", out_denorm.cpu().numpy())
+            np.save(npy_dir / f"{index}_{safe_name}_gt.npy", y_denorm.cpu().numpy())
+
+            data_dir = ROOT / cfg_data.source.data_dir / sample_name
+            pred_coef, gt_coef = maybe_calculate_coefficient(
+                data_dir,
+                pred_press.cpu().numpy(),
+                pred_velo.cpu().numpy(),
+                gt_press.cpu().numpy(),
+                gt_velo.cpu().numpy(),
+            )
+            if pred_coef is not None and gt_coef is not None:
+                pred_coef_list.append(pred_coef)
+                gt_coef_list.append(gt_coef)
+
+            l2errs_press.append((torch.norm(pred_press - gt_press) / (torch.norm(gt_press) + 1e-8)).cpu().numpy())
+            l2errs_velo.append((torch.norm(pred_velo - gt_velo) / (torch.norm(gt_velo) + 1e-8)).cpu().numpy())
+            mses_press.append(criterion_func(out[data.surf, -1], targets[data.surf, -1]).mean().cpu().numpy())
+            mses_velo_var.append(criterion_func(out[~data.surf, :-1], targets[~data.surf, :-1]).mean().cpu().numpy())
+
+            if cfg_test.save_vtk and (data_dir / "quadpress_smpl.vtk").exists():
+                save_prediction_to_vtk(
+                    out_denorm=out_denorm,
+                    targets=targets,
+                    cfd_data=data,
+                    sample_name=sample_name,
+                    output_dir=str(vtk_dir),
+                    index=index,
+                    data_dir=str(ROOT / cfg_data.source.data_dir),
+                )
+            if cfg_test.visualize and cfg_test.save_vtk:
+                visualize_prediction(output_dir=str(vtk_dir), vis_dir=str(vis_dir), index=index)
+
+    logger.info("Results saved to: %s", result_root)
+    logger.info("Relative L2 pressure: %.6f", float(np.mean(l2errs_press)))
+    logger.info("Relative L2 velocity: %.6f", float(np.mean(l2errs_velo)))
+    logger.info("RMSE pressure: %.6f", float(np.sqrt(np.mean(mses_press)) * coef_norm[3][-1]))
+    rmse_velo = np.sqrt(np.mean(mses_velo_var, axis=0)) * coef_norm[3][:-1]
+    logger.info("Combined velocity RMSE: %.6f", float(np.sqrt(np.mean(np.square(rmse_velo)))))
+    logger.info("Mean inference time (s): %.6f", float(np.mean(times)))
+    if gt_coef_list:
+        coef_error = np.mean(np.abs(np.array(pred_coef_list) - np.array(gt_coef_list)) / (np.array(gt_coef_list) + 1e-8))
+        logger.info("Mean relative CD error: %.6f", float(coef_error))
+    else:
+        logger.info("Skipped drag coefficient metrics because VTK geometry files were not present.")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,30 @@
+from __future__ import annotations
+
+import os
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+os.chdir(ROOT)
+
+from onescience.utils.YParams import YParams
+
+
+def main() -> None:
+    cfg = YParams(str(ROOT / "conf/config.yaml"), "model")
+    cfg_test = YParams(str(ROOT / "conf/config.yaml"), "inference")
+    result_root = ROOT / cfg_test.result_dir / cfg.name
+    npy_dir = result_root / "npy"
+    pred_files = sorted(npy_dir.glob("*_pred.npy")) if npy_dir.exists() else []
+    gt_files = sorted(npy_dir.glob("*_gt.npy")) if npy_dir.exists() else []
+
+    print(f"Result directory: {result_root}")
+    print(f"Prediction files: {len(pred_files)}")
+    print(f"Ground-truth files: {len(gt_files)}")
+    if pred_files:
+        print(f"First prediction: {pred_files[0]}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,221 @@
+from __future__ import annotations
+
+import logging
+import os
+import sys
+import time
+import importlib.util
+from pathlib import Path
+
+import torch
+import torch.distributed as dist
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel
+
+ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(ROOT))
+os.chdir(ROOT)
+
+from model import Transolver3D, Transolver3D_plus
+import onescience
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+
+
+def load_shapenet_car_datapipe():
+    module_path = Path(onescience.__file__).resolve().parent / "datapipes/cfd/ShapeNetCar.py"
+    spec = importlib.util.spec_from_file_location("_onescience_shapenetcar", module_path)
+    if spec is None or spec.loader is None:
+        raise ImportError(f"Unable to load ShapeNetCarDatapipe from {module_path}")
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module.ShapeNetCarDatapipe
+
+
+def setup_logging(rank: int) -> logging.Logger:
+    level = logging.INFO if rank == 0 else logging.WARNING
+    logging.basicConfig(
+        level=level,
+        format="%(asctime)s - %(levelname)s - %(message)s",
+        datefmt="%Y-%m-%d %H:%M:%S",
+    )
+    logging.getLogger().setLevel(level)
+    return logging.getLogger(__name__)
+
+
+def build_model(model_name: str, model_params, device: torch.device) -> torch.nn.Module:
+    model_cls = {
+        "Transolver": Transolver3D,
+        "Transolver_plus": Transolver3D_plus,
+    }.get(model_name)
+    if model_cls is None:
+        raise NotImplementedError(f"Model {model_name} initialization not implemented.")
+
+    return model_cls(
+        n_hidden=model_params.n_hidden,
+        n_layers=model_params.n_layers,
+        space_dim=model_params.space_dim,
+        fun_dim=model_params.fun_dim,
+        n_head=model_params.n_head,
+        mlp_ratio=model_params.mlp_ratio,
+        out_dim=model_params.out_dim,
+        slice_num=model_params.slice_num,
+        unified_pos=model_params.unified_pos,
+    ).to(device)
+
+
+def resolve_device(gpuid: int) -> torch.device:
+    if torch.cuda.is_available() and int(gpuid) >= 0:
+        return torch.device(f"cuda:{gpuid}")
+    return torch.device("cpu")
+
+
+def save_checkpoint(model, optimizer, scheduler, epoch: int, loss: float, ckp_dir: str, model_name: str) -> None:
+    Path(ckp_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,
+        },
+        Path(ckp_dir) / f"{model_name}.pth",
+    )
+
+
+def main() -> None:
+    DistributedManager.initialize()
+    manager = DistributedManager()
+    logger = setup_logging(manager.rank)
+
+    config_file_path = str(ROOT / "conf/config.yaml")
+    cfg = YParams(config_file_path, "model")
+    cfg_data = YParams(config_file_path, "datapipe")
+    cfg_train = YParams(config_file_path, "training")
+
+    model_name = cfg.name
+    if model_name not in cfg.specific_params:
+        raise ValueError(f"Model '{model_name}' not found in model.specific_params.")
+    model_params = cfg.specific_params[model_name]
+    cfg_data.model_hparams = model_params
+
+    logger.info("Initializing ShapeNetCar datapipe...")
+    ShapeNetCarDatapipe = load_shapenet_car_datapipe()
+    datapipe = ShapeNetCarDatapipe(params=cfg_data, distributed=(manager.world_size > 1))
+    train_dataloader, train_sampler = datapipe.train_dataloader()
+    val_dataloader, val_sampler = datapipe.val_dataloader()
+
+    if manager.world_size > 1:
+        device = torch.device(f"cuda:{manager.local_rank}" if torch.cuda.is_available() else "cpu")
+    else:
+        device = resolve_device(cfg_train.gpuid)
+
+    model = build_model(model_name, model_params, device)
+    if manager.rank == 0:
+        total_params = sum(p.numel() for p in model.parameters() if p.requires_grad)
+        logger.info("Model: %s, trainable params: %.2fM", model_name, total_params / 1e6)
+
+    if manager.world_size > 1:
+        model = DistributedDataParallel(
+            model,
+            device_ids=[manager.local_rank],
+            output_device=manager.local_rank,
+            find_unused_parameters=True,
+        )
+
+    optimizer = torch.optim.Adam(model.parameters(), lr=cfg_train.lr)
+    scheduler = torch.optim.lr_scheduler.OneCycleLR(
+        optimizer,
+        max_lr=cfg_train.lr,
+        total_steps=max(1, len(train_dataloader) * cfg_train.max_epoch),
+    )
+    if cfg_train.loss_criterion == "MSE":
+        loss_criterion = nn.MSELoss(reduction="none")
+    elif cfg_train.loss_criterion == "MAE":
+        loss_criterion = nn.L1Loss(reduction="none")
+    else:
+        raise ValueError(f"Unknown loss criterion: {cfg_train.loss_criterion}")
+
+    best_valid_loss = 1.0e6
+    best_loss_epoch = 0
+    logger.info("Starting training...")
+
+    for epoch in range(cfg_train.max_epoch):
+        epoch_start_time = time.time()
+        if manager.world_size > 1:
+            train_sampler.set_epoch(epoch)
+            if val_sampler is not None:
+                val_sampler.set_epoch(epoch)
+
+        model.train()
+        train_loss = train_loss_press = train_loss_velo = 0.0
+        for data in train_dataloader:
+            data = data.to(device)
+            optimizer.zero_grad()
+            out = model(data)
+            targets = data.y
+            loss_press = loss_criterion(out[data.surf, -1], targets[data.surf, -1]).mean()
+            loss_velo = loss_criterion(out[:, :-1], targets[:, :-1]).mean()
+            loss = loss_velo + cfg_train.loss_weight * loss_press
+            loss.backward()
+            optimizer.step()
+            scheduler.step()
+            train_loss += loss.item()
+            train_loss_press += loss_press.item()
+            train_loss_velo += loss_velo.item()
+
+        train_loss /= max(1, len(train_dataloader))
+        train_loss_press /= max(1, len(train_dataloader))
+        train_loss_velo /= max(1, len(train_dataloader))
+
+        valid_loss = valid_loss_press = valid_loss_velo = 0.0
+        if (epoch + 1) % cfg_train.val_iter == 0 or epoch == cfg_train.max_epoch - 1:
+            model.eval()
+            with torch.no_grad():
+                for data in val_dataloader:
+                    data = data.to(device)
+                    out = model(data)
+                    targets = data.y
+                    loss_press = loss_criterion(out[data.surf, -1], targets[data.surf, -1]).mean()
+                    loss_velo = loss_criterion(out[:, :-1], targets[:, :-1]).mean()
+                    loss = loss_velo + cfg_train.loss_weight * loss_press
+                    if manager.world_size > 1:
+                        dist.all_reduce(loss, op=dist.ReduceOp.AVG)
+                        dist.all_reduce(loss_press, op=dist.ReduceOp.AVG)
+                        dist.all_reduce(loss_velo, op=dist.ReduceOp.AVG)
+                    valid_loss += loss.item()
+                    valid_loss_press += loss_press.item()
+                    valid_loss_velo += loss_velo.item()
+
+            valid_loss /= max(1, len(val_dataloader))
+            valid_loss_press /= max(1, len(val_dataloader))
+            valid_loss_velo /= max(1, len(val_dataloader))
+
+        if manager.rank == 0:
+            logger.info(
+                "Epoch [%d/%d] | Time: %.2fs | Train: %.6f (velo %.6f, press %.6f) | "
+                "Valid: %.6f (velo %.6f, press %.6f)",
+                epoch + 1,
+                cfg_train.max_epoch,
+                time.time() - epoch_start_time,
+                train_loss,
+                train_loss_velo,
+                train_loss_press,
+                valid_loss,
+                valid_loss_velo,
+                valid_loss_press,
+            )
+            if valid_loss > 0 and valid_loss < best_valid_loss:
+                best_valid_loss = valid_loss
+                best_loss_epoch = epoch
+                save_checkpoint(model, optimizer, scheduler, epoch, valid_loss, cfg_train.checkpoint_dir, model_name)
+                logger.info("New best checkpoint saved to %s/%s.pth", cfg_train.checkpoint_dir, model_name)
+            if epoch - best_loss_epoch > cfg_train.patience:
+                logger.warning("Validation loss has not improved for %d epochs. Stopping.", cfg_train.patience)
+                break
+
+
+if __name__ == "__main__":
+    main()

weight/.gitkeep

@@ -0,0 +1 @@
+