Upload folder using huggingface_hub

README.md

@@ -0,0 +1,136 @@
+---
+frameworks:
+- ""
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- fluid dynamics
+- external flow prediction
+- unstructured-mesh simulation
+
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">MeshGraphNet</span>
+  </strong>
+</p>
+
+# Model Overview
+
+MeshGraphNets is a graph neural network developed by DeepMind for mesh-based physical simulation. It rapidly predicts the dynamics of complex physical systems, including fluids, structures, and cloth.
+
+Paper: Learning Mesh-Based Simulation with Graph Networks
+https://arxiv.org/abs/2010.03409
+
+# Model Description
+MeshGraphNets uses an encoder–processor–decoder graph-network architecture trained on trajectories from fluid, structural, and cloth simulations to perform long-horizon dynamical simulation of complex physical systems.
+
+## Use Cases
+
+| Use Case | Description |
+|---|---|
+| External flow prediction | Predict velocity, pressure, and other flow variables at mesh nodes |
+| Structural deformation simulation | Predict the displacement, stress, and deformation of loaded structures |
+| Cloth dynamics | Simulate the motion of deformable objects such as flexible membranes and cloth |
+| ModelScope/OneCode execution | Download the standalone model package, install its dependencies, and run the provided scripts |
+
+
+# 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.
+- 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/MeshGraphNet --local_dir ./MeshGraphNet 
+cd MeshGraphNet 
+```
+
+### 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 `cylinder_flow` 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/cylinder_flow --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 `.pth` files under `weight/checkpoints`.
+
+### Model Weights
+This repository will provide weights trained on the `cylinder_flow` dataset 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 MeshGraphNet paper: [Learning Mesh-Based Simulation with Graph Networks](https://arxiv.org/abs/2010.03409).
+- This repository retains source attribution and has been adapted for automated execution through OneScience and ModelScope.

config/config.yaml

@@ -0,0 +1,63 @@
+paths:
+  data_dir: data/cylinder_flow
+  fake_data_path: data/cylinder_flow/fake_cylinder_flow.pt
+  checkpoint_dir: weight/checkpoints
+  result_dir: results
+
+datapipe:
+  source:
+    data_dir: data/cylinder_flow
+    stats_dir: data/cylinder_flow/stats
+    fake_data: true
+    fake_data_path: data/cylinder_flow/fake_cylinder_flow.pt
+  verbose: true
+  data:
+    train_samples: 4
+    train_steps: 5
+    val_samples: 2
+    val_steps: 5
+    test_samples: 2
+    test_steps: 5
+    noise_std: 0.0
+  dataloader:
+    batch_size: 1
+    num_workers: 0
+
+model:
+  name: MeshGraphNet
+  specific_params:
+    MeshGraphNet:
+      num_input_features: 6
+      num_edge_features: 3
+      num_output_features: 3
+      processor_size: 2
+      hidden_dim_processor: 32
+      num_layers_node_processor: 2
+      num_layers_edge_processor: 2
+      hidden_dim_node_encoder: 32
+      hidden_dim_edge_encoder: 32
+      hidden_dim_node_decoder: 32
+      do_concat_trick: false
+      num_processor_checkpoint_segments: 0
+      recompute_activation: false
+
+training:
+  device: cpu
+  max_epoch: 1
+  lr: 0.0001
+  lr_decay_rate: 0.999
+  amp: false
+  jit: false
+  checkpoint_dir: weight/checkpoints
+  loss_criterion: MSE
+  patience: 10
+  log_interval: 1
+  gpuid: 0
+
+inference:
+  device: cpu
+  checkpoint_dir: weight/checkpoints
+  output_path: results/inference_results.npz
+  viz_vars: ["u", "v", "p"]
+  frame_skip: 1
+  frame_interval: 1

configuration.json

@@ -0,0 +1 @@
+{"framework":"Pytorch","task":"other","model":"MeshGraphNet","domain":"cfd"}

model/meshgraphnet.py

@@ -0,0 +1,306 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+try:
+    import dgl  # noqa: F401 for docs
+    from dgl import DGLGraph
+except ImportError:
+    raise ImportError(
+        "Mesh Graph Net requires the DGL library. Install the "
+    )
+from dataclasses import dataclass
+from itertools import chain
+from typing import Callable, List, Tuple, Union
+
+import onescience  # noqa: F401 for docs
+from onescience.modules.edge.mesh_edge_block import MeshEdgeBlock
+from onescience.modules.mlp.mesh_graph_mlp import MeshGraphMLP
+from onescience.modules.node.mesh_node_block import MeshNodeBlock
+
+from onescience.modules.utils.gnnlayer_utils import CuGraphCSC, set_checkpoint_fn
+from onescience.modules.layer.activations import get_activation
+from onescience.modules.meta import ModelMetaData
+from onescience.modules.module import Module
+
+
+@dataclass
+class MetaData(ModelMetaData):
+    name: str = "MeshGraphNet"
+    # Optimization, no JIT as DGLGraph causes trouble
+    jit: bool = False
+    cuda_graphs: bool = False
+    amp_cpu: bool = False
+    amp_gpu: bool = True
+    torch_fx: bool = False
+    # Inference
+    onnx: bool = False
+    # Physics informed
+    func_torch: bool = True
+    auto_grad: bool = True
+
+
+class MeshGraphNet(Module):
+    """
+        MeshGraphNet 网络架构。
+
+        该模型基于 "Learning mesh-based simulation with graph networks" (Pfaff et al., 2020) 实现。
+        它采用 Encode-Process-Decode 架构：
+        1. **Encoder**: 将节点和边的物理特征映射到高维隐空间。
+        2. **Processor**: 通过多层消息传递（Message Passing）在图中传播信息，更新节点和边的隐状态。
+        3. **Decoder**: 将处理后的节点特征解码回物理空间（例如加速度或速度增量）。
+
+        本实现使用 MeshGraphMLP、MeshEdgeBlock 和 MeshNodeBlock 构建。
+
+        Args:
+            input_dim_nodes (int): 输入节点特征的维度。
+            input_dim_edges (int): 输入边特征的维度。
+            output_dim (int): 输出特征的维度（通常是节点状态的更新量）。
+            processor_size (int, optional): 消息传递块（Processor Block）的数量。默认值: 15。
+            mlp_activation_fn (Union[str, List[str]], optional): MLP 中使用的激活函数。默认值: 'relu'。
+            num_layers_node_processor (int, optional): 处理器中节点更新 MLP 的层数。默认值: 2。
+            num_layers_edge_processor (int, optional): 处理器中边更新 MLP 的层数。默认值: 2。
+            hidden_dim_processor (int, optional): 处理器中隐层的特征维度。默认值: 128。
+            hidden_dim_node_encoder (int, optional): 节点编码器的隐层维度。默认值: 128。
+            num_layers_node_encoder (Union[int, None], optional): 节点编码器的层数。如果为 None，则不使用编码器。默认值: 2。
+            hidden_dim_edge_encoder (int, optional): 边编码器的隐层维度。默认值: 128。
+            num_layers_edge_encoder (Union[int, None], optional): 边编码器的层数。如果为 None，则不使用编码器。默认值: 2。
+            hidden_dim_node_decoder (int, optional): 节点解码器的隐层维度。默认值: 128。
+            num_layers_node_decoder (Union[int, None], optional): 节点解码器的层数。如果为 None，则不使用解码器。默认值: 2。
+            aggregation (str, optional): 消息聚合方式，可选 "sum", "mean" 等。默认值: "sum"。
+            do_concat_trick (bool, optional): 是否使用拼接优化技巧 (MLP+idx+sum) 以节省显存。默认值: False。
+            num_processor_checkpoint_segments (int, optional): 梯度检查点 (Gradient Checkpointing) 的分段数。0 表示禁用。默认值: 0。
+            recompute_activation (bool, optional): 是否重计算激活函数以节省显存。默认值: False。
+
+        形状:
+            输入 node_features: (N, input_dim_nodes)，其中 N 为节点总数。
+            输入 edge_features: (M, input_dim_edges)，其中 M 为边总数。
+            输入 graph: DGLGraph 或 CuGraphCSC，定义图拓扑结构。
+            输出: (N, output_dim)，解码后的节点物理量。
+
+    """
+
+    def __init__(
+        self,
+        input_dim_nodes: int,
+        input_dim_edges: int,
+        output_dim: int,
+        processor_size: int = 15,
+        mlp_activation_fn: Union[str, List[str]] = "relu",
+        num_layers_node_processor: int = 2,
+        num_layers_edge_processor: int = 2,
+        hidden_dim_processor: int = 128,
+        hidden_dim_node_encoder: int = 128,
+        num_layers_node_encoder: Union[int, None] = 2,
+        hidden_dim_edge_encoder: int = 128,
+        num_layers_edge_encoder: Union[int, None] = 2,
+        hidden_dim_node_decoder: int = 128,
+        num_layers_node_decoder: Union[int, None] = 2,
+        aggregation: str = "sum",
+        do_concat_trick: bool = False,
+        num_processor_checkpoint_segments: int = 0,
+        recompute_activation: bool = False,
+    ):
+        super().__init__(meta=MetaData())
+
+        activation_fn = get_activation(mlp_activation_fn)
+
+        # 1. Edge Encoder
+        self.edge_encoder = MeshGraphMLP(
+            input_dim=input_dim_edges,
+            output_dim=hidden_dim_processor,
+            hidden_dim=hidden_dim_edge_encoder,
+            hidden_layers=num_layers_edge_encoder,
+            activation_fn=activation_fn,
+            norm_type="LayerNorm",
+            recompute_activation=recompute_activation,
+        )
+
+        # 2. Node Encoder
+        self.node_encoder = MeshGraphMLP(
+            input_dim=input_dim_nodes,
+            output_dim=hidden_dim_processor,
+            hidden_dim=hidden_dim_node_encoder,
+            hidden_layers=num_layers_node_encoder,
+            activation_fn=activation_fn,
+            norm_type="LayerNorm",
+            recompute_activation=recompute_activation,
+        )
+
+        # 3. Node Decoder
+        self.node_decoder = MeshGraphMLP(
+            input_dim=hidden_dim_processor,
+            output_dim=output_dim,
+            hidden_dim=hidden_dim_node_decoder,
+            hidden_layers=num_layers_node_decoder,
+            activation_fn=activation_fn,
+            norm_type=None,
+            recompute_activation=recompute_activation,
+        )
+
+        # 4. Processor (Core GNN)
+        self.processor = MeshGraphNetProcessor(
+            processor_size=processor_size,
+            input_dim_node=hidden_dim_processor,
+            input_dim_edge=hidden_dim_processor,
+            num_layers_node=num_layers_node_processor,
+            num_layers_edge=num_layers_edge_processor,
+            aggregation=aggregation,
+            norm_type="LayerNorm",
+            activation_fn=activation_fn,
+            do_concat_trick=do_concat_trick,
+            num_processor_checkpoint_segments=num_processor_checkpoint_segments,
+        )
+
+    def forward(
+        self,
+        node_features: Tensor,
+        edge_features: Tensor,
+        graph: Union[DGLGraph, List[DGLGraph], CuGraphCSC],
+    ) -> Tensor:
+        edge_features = self.edge_encoder(edge_features)
+        node_features = self.node_encoder(node_features)
+        x = self.processor(node_features, edge_features, graph)
+        x = self.node_decoder(x)
+        return x
+
+
+class MeshGraphNetProcessor(nn.Module):
+    """
+        MeshGraphNet 核心处理器 (Processor)。
+
+        该模块由一系列堆叠的消息传递块 (Message Passing Blocks) 组成。
+        每个块包含两个步骤：
+        1. **Edge Block**: 使用 MeshEdgeBlock 更新边特征。
+        2. **Node Block**: 使用 MeshNodeBlock 聚合边信息并更新节点特征。
+
+        支持梯度检查点 (Gradient Checkpointing) 以减少大规模图训练时的显存占用。
+
+        Args:
+            processor_size (int, optional): 处理器包含的消息传递层数。默认值: 15。
+            input_dim_node (int, optional): 输入节点特征维度。默认值: 128。
+            input_dim_edge (int, optional): 输入边特征维度。默认值: 128。
+            num_layers_node (int, optional): 节点更新 MLP 的层数。默认值: 2。
+            num_layers_edge (int, optional): 边更新 MLP 的层数。默认值: 2。
+            aggregation (str, optional): 消息聚合方式 ("sum", "mean" 等)。默认值: "sum"。
+            norm_type (str, optional): 归一化类型。默认值: "LayerNorm"。
+            activation_fn (nn.Module, optional): 激活函数。默认值: nn.ReLU()。
+            do_concat_trick (bool, optional): 是否启用显存优化技巧。默认值: False。
+            num_processor_checkpoint_segments (int, optional): 梯度检查点分段数。默认值: 0 (禁用)。
+
+        形状:
+            输入 node_features: (N, input_dim_node)
+            输入 edge_features: (M, input_dim_edge)
+            输入 graph: DGLGraph
+            输出: (N, input_dim_node) - 仅返回更新后的节点特征。
+    
+    """
+
+    def __init__(
+        self,
+        processor_size: int = 15,
+        input_dim_node: int = 128,
+        input_dim_edge: int = 128,
+        num_layers_node: int = 2,
+        num_layers_edge: int = 2,
+        aggregation: str = "sum",
+        norm_type: str = "LayerNorm",
+        activation_fn: nn.Module = nn.ReLU(),
+        do_concat_trick: bool = False,
+        num_processor_checkpoint_segments: int = 0,
+    ):
+        super().__init__()
+        self.processor_size = processor_size
+        self.num_processor_checkpoint_segments = num_processor_checkpoint_segments
+
+        edge_blocks = []
+        node_blocks = []
+
+        for _ in range(self.processor_size):
+            edge_blocks.append(
+                MeshEdgeBlock(
+                    input_dim_nodes=input_dim_node,
+                    input_dim_edges=input_dim_edge,
+                    output_dim=input_dim_edge,
+                    hidden_dim=input_dim_edge,
+                    hidden_layers=num_layers_edge,
+                    activation_fn=activation_fn,
+                    norm_type=norm_type,
+                    do_concat_trick=do_concat_trick,
+                    recompute_activation=False
+                )
+            )
+            node_blocks.append(
+                MeshNodeBlock(
+                    aggregation=aggregation,
+                    input_dim_nodes=input_dim_node,
+                    input_dim_edges=input_dim_edge,
+                    output_dim=input_dim_node,
+                    hidden_dim=input_dim_node,
+                    hidden_layers=num_layers_node,
+                    activation_fn=activation_fn,
+                    norm_type=norm_type,
+                    recompute_activation=False
+                )
+            )
+
+        # 按照 Edge -> Node 的顺序交替排列
+        layers = list(chain(*zip(edge_blocks, node_blocks)))
+
+        self.processor_layers = nn.ModuleList(layers)
+        self.num_processor_layers = len(self.processor_layers)
+        self.set_checkpoint_segments(self.num_processor_checkpoint_segments)
+
+    def set_checkpoint_segments(self, checkpoint_segments: int):
+        if checkpoint_segments > 0:
+            if self.num_processor_layers % checkpoint_segments != 0:
+                raise ValueError(
+                    "Processor layers must be a multiple of checkpoint_segments"
+                )
+            segment_size = self.num_processor_layers // checkpoint_segments
+            self.checkpoint_segments = []
+            for i in range(0, self.num_processor_layers, segment_size):
+                self.checkpoint_segments.append((i, i + segment_size))
+            self.checkpoint_fn = set_checkpoint_fn(True)
+        else:
+            self.checkpoint_fn = set_checkpoint_fn(False)
+            self.checkpoint_segments = [(0, self.num_processor_layers)]
+
+    def run_function(
+        self, segment_start: int, segment_end: int
+    ) -> Callable[
+        [Tensor, Tensor, Union[DGLGraph, List[DGLGraph]]], Tuple[Tensor, Tensor]
+    ]:
+        segment = self.processor_layers[segment_start:segment_end]
+
+        def custom_forward(
+            node_features: Tensor,
+            edge_features: Tensor,
+            graph: Union[DGLGraph, List[DGLGraph]],
+        ) -> Tuple[Tensor, Tensor]:
+            for module in segment:
+                edge_features, node_features = module(
+                    edge_features, node_features, graph
+                )
+            return edge_features, node_features
+
+        return custom_forward
+
+    @torch.jit.unused
+    def forward(
+        self,
+        node_features: Tensor,
+        edge_features: Tensor,
+        graph: Union[DGLGraph, List[DGLGraph], CuGraphCSC],
+    ) -> Tensor:
+        for segment_start, segment_end in self.checkpoint_segments:
+            edge_features, node_features = self.checkpoint_fn(
+                self.run_function(segment_start, segment_end),
+                node_features,
+                edge_features,
+                graph,
+                use_reentrant=False,
+                preserve_rng_state=False,
+            )
+
+        return node_features

scripts/fake_data.py

@@ -0,0 +1,167 @@
+import os
+import sys
+from pathlib import Path
+
+import dgl
+import torch
+from dgl.dataloading import GraphDataLoader
+from torch.utils.data import Dataset
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "model"))
+
+from onescience.utils.YParams import YParams
+
+
+def make_graph(num_nodes: int = 12):
+    src = torch.arange(num_nodes, dtype=torch.int32)
+    dst = torch.roll(src, shifts=-1)
+    graph = dgl.to_bidirected(dgl.graph((src, dst), num_nodes=num_nodes, idtype=torch.int32))
+
+    pos = torch.stack(
+        (
+            torch.linspace(0.0, 1.0, num_nodes),
+            torch.sin(torch.linspace(0.0, 3.14159, num_nodes)) * 0.2,
+        ),
+        dim=1,
+    )
+    row, col = graph.edges()
+    disp = pos[row.long()] - pos[col.long()]
+    graph.edata["x"] = torch.cat(
+        (disp, torch.linalg.norm(disp, dim=-1, keepdim=True)),
+        dim=1,
+    )
+
+    velocity = torch.randn(num_nodes, 2) * 0.1
+    node_type = torch.zeros(num_nodes, 4)
+    node_type[:, 0] = 1.0
+    graph.ndata["x"] = torch.cat((velocity, node_type), dim=1)
+    graph.ndata["y"] = torch.cat(
+        (torch.randn(num_nodes, 2) * 0.01, torch.randn(num_nodes, 1) * 0.01),
+        dim=1,
+    )
+    graph.ndata["mesh_pos"] = pos
+
+    cells = torch.tensor(
+        [[i, i + 1, min(i + 2, num_nodes - 1)] for i in range(num_nodes - 2)],
+        dtype=torch.int64,
+    )
+    mask = torch.ones(num_nodes, 1, dtype=torch.bool)
+    return {"graph": graph, "cells": cells, "mask": mask}
+
+
+class FakeGraphDataset(Dataset):
+    def __init__(self, samples):
+        self.samples = samples
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        sample = self.samples[index]
+        if isinstance(sample, dict) and "graph" in sample:
+            return sample["graph"]
+        return sample
+
+
+def _resolve_path(project_root: Path, path):
+    path = Path(path)
+    return path if path.is_absolute() else project_root / path
+
+
+def _torch_load(path: Path):
+    try:
+        return torch.load(path, map_location="cpu", weights_only=False)
+    except TypeError:
+        return torch.load(path, map_location="cpu")
+
+
+class FakeCylinderFlowDatapipe:
+    def __init__(self, params, project_root: Path):
+        self.params = params
+        fake_data_path = _resolve_path(project_root, params.source.fake_data_path)
+        if not fake_data_path.exists():
+            raise FileNotFoundError(
+                f"Fake data file not found: {fake_data_path}. Run scripts/fake_data.py first."
+            )
+
+        payload = _torch_load(fake_data_path)
+        self.train_dataset = FakeGraphDataset(payload["train"])
+        self.val_dataset = FakeGraphDataset(payload["val"])
+        self.test_dataset = FakeGraphDataset(payload["test"])
+        self.stats = payload.get("stats", {})
+
+    def _loader(self, dataset, shuffle=False, drop_last=False):
+        return GraphDataLoader(
+            dataset,
+            batch_size=self.params.dataloader.batch_size,
+            drop_last=drop_last,
+            num_workers=self.params.dataloader.num_workers,
+            pin_memory=True,
+            shuffle=shuffle,
+        )
+
+    def train_dataloader(self):
+        return self._loader(self.train_dataset, shuffle=True), None
+
+    def val_dataloader(self):
+        return self._loader(self.val_dataset), None
+
+    def test_dataloader(self):
+        return self._loader(self.test_dataset)
+
+
+def use_fake_data(params):
+    return bool(getattr(params.source, "fake_data", False))
+
+
+def build_cylinder_flow_datapipe(params, distributed: bool, project_root: Path):
+    if use_fake_data(params):
+        return FakeCylinderFlowDatapipe(params=params, project_root=project_root)
+
+    from onescience.datapipes.cfd import DeepMind_CylinderFlowDatapipe
+
+    return DeepMind_CylinderFlowDatapipe(params=params, distributed=distributed)
+
+
+def main():
+    os.chdir(PROJECT_ROOT)
+    config_path = PROJECT_ROOT / "config" / "config.yaml"
+    cfg_data = YParams(config_path, "datapipe")
+    output_path = PROJECT_ROOT / cfg_data.source.fake_data_path
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+
+    payload = {
+        "train": [
+            make_graph()
+            for _ in range(cfg_data.data.train_samples * (cfg_data.data.train_steps - 1))
+        ],
+        "val": [
+            make_graph()
+            for _ in range(cfg_data.data.val_samples * (cfg_data.data.val_steps - 1))
+        ],
+        "test": [
+            make_graph()
+            for _ in range(cfg_data.data.test_samples * (cfg_data.data.test_steps - 1))
+        ],
+        "stats": {
+            "edge_stats": {
+                "edge_mean": torch.zeros(3),
+                "edge_std": torch.ones(3),
+            },
+            "node_stats": {
+                "velocity_mean": torch.zeros(2),
+                "velocity_std": torch.ones(2),
+                "velocity_diff_mean": torch.zeros(2),
+                "velocity_diff_std": torch.ones(2),
+                "pressure_mean": torch.zeros(1),
+                "pressure_std": torch.ones(1),
+            },
+        },
+    }
+    torch.save(payload, output_path)
+    print(f"Fake data saved to {output_path.relative_to(PROJECT_ROOT)}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,92 @@
+import logging
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+
+from model.meshgraphnet import MeshGraphNet
+from onescience.utils.YParams import YParams
+from onescience.launch.utils import load_checkpoint 
+from fake_data import build_cylinder_flow_datapipe
+
+
+def build_model(model_params, device):
+    mlp_act = "silu" if model_params.recompute_activation else "relu"
+    return MeshGraphNet(
+        input_dim_nodes=model_params.num_input_features,
+        input_dim_edges=model_params.num_edge_features,
+        output_dim=model_params.num_output_features,
+        processor_size=model_params.processor_size,
+        hidden_dim_processor=model_params.hidden_dim_processor,
+        num_layers_node_processor=model_params.num_layers_node_processor,
+        num_layers_edge_processor=model_params.num_layers_edge_processor,
+        hidden_dim_node_encoder=model_params.hidden_dim_node_encoder,
+        hidden_dim_edge_encoder=model_params.hidden_dim_edge_encoder,
+        hidden_dim_node_decoder=model_params.hidden_dim_node_decoder,
+        mlp_activation_fn=mlp_act,
+        do_concat_trick=model_params.do_concat_trick,
+        num_processor_checkpoint_segments=model_params.num_processor_checkpoint_segments,
+        recompute_activation=model_params.recompute_activation,
+    ).to(device)
+
+
+def resolve_device(device_name: str):
+    if device_name == "cpu":
+        return torch.device("cpu")
+    if device_name in ("cuda", "gpu"):
+        if not torch.cuda.is_available():
+            raise RuntimeError("Config requested cuda device, but torch.cuda.is_available() is false.")
+        return torch.device("cuda:0")
+    return torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+def main():
+    os.chdir(PROJECT_ROOT)
+    logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
+    logger = logging.getLogger("mesh_graph_net.inference")
+
+    config_path = PROJECT_ROOT / "config" / "config.yaml"
+    cfg_model = YParams(config_path, "model")
+    cfg_data = YParams(config_path, "datapipe")
+    cfg_train = YParams(config_path, "training")
+    cfg_inference = YParams(config_path, "inference")
+    model_params = cfg_model.specific_params[cfg_model.name]
+
+    device = resolve_device(getattr(cfg_inference, "device", "auto"))
+    logger.info("Using device: %s", device)
+    datapipe = build_cylinder_flow_datapipe(
+        params=cfg_data,
+        distributed=False,
+        project_root=PROJECT_ROOT,
+    )
+    loader = datapipe.test_dataloader()
+    model = build_model(model_params, device)
+    checkpoint_dir = PROJECT_ROOT / getattr(cfg_inference, "checkpoint_dir", cfg_train.checkpoint_dir)
+    epoch = load_checkpoint(checkpoint_dir, models=model, device=device)
+    if epoch == 0:
+        logger.warning("No checkpoint found in %s; running with randomly initialized weights", checkpoint_dir)
+    model.eval()
+
+    predictions, targets = [], []
+    with torch.no_grad():
+        for batch in loader:
+            graph = batch[0] if isinstance(batch, (tuple, list)) else batch
+            graph = graph.to(device)
+            pred = model(graph.ndata["x"], graph.edata["x"], graph)
+            predictions.append(pred.cpu().numpy())
+            targets.append(graph.ndata["y"].cpu().numpy())
+
+    output_path = PROJECT_ROOT / cfg_inference.output_path
+    output_path.parent.mkdir(parents=True, exist_ok=True)
+    np.savez(output_path, prediction=np.concatenate(predictions, axis=0), target=np.concatenate(targets, axis=0))
+    logger.info("Saved inference results to %s", output_path)
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,30 @@
+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"))
+
+from onescience.utils.YParams import YParams
+
+
+def main():
+    os.chdir(PROJECT_ROOT)
+    cfg = YParams(PROJECT_ROOT / "config" / "config.yaml", "inference")
+    result_path = PROJECT_ROOT / cfg.output_path
+    if not result_path.exists():
+        raise FileNotFoundError(f"Result file not found: {result_path}. Run scripts/inference.py first.")
+    data = np.load(result_path)
+    prediction = data["prediction"]
+    target = data["target"]
+    mse = float(np.mean((prediction - target) ** 2))
+    print(f"result_path: {result_path.relative_to(PROJECT_ROOT)}")
+    print(f"prediction_shape: {prediction.shape}")
+    print(f"target_shape: {target.shape}")
+    print(f"mse: {mse:.8f}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,154 @@
+import logging
+import os
+import sys
+import time
+from pathlib import Path
+
+import torch
+import torch.nn as nn
+from torch.amp import GradScaler, autocast
+from torch.nn.parallel import DistributedDataParallel
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.meshgraphnet import MeshGraphNet
+from onescience.distributed import DistributedManager
+from onescience.utils.YParams import YParams
+from onescience.launch.utils import load_checkpoint, save_checkpoint
+from fake_data import build_cylinder_flow_datapipe
+
+
+def setup_logging(rank: int):
+    level = logging.INFO if rank == 0 else logging.WARNING
+    logging.basicConfig(level=level, format="%(asctime)s - %(levelname)s - %(message)s")
+    return logging.getLogger("mesh_graph_net.train")
+
+
+def build_model(model_params, device):
+    mlp_act = "silu" if model_params.recompute_activation else "relu"
+    return MeshGraphNet(
+        input_dim_nodes=model_params.num_input_features,
+        input_dim_edges=model_params.num_edge_features,
+        output_dim=model_params.num_output_features,
+        processor_size=model_params.processor_size,
+        hidden_dim_processor=model_params.hidden_dim_processor,
+        num_layers_node_processor=model_params.num_layers_node_processor,
+        num_layers_edge_processor=model_params.num_layers_edge_processor,
+        hidden_dim_node_encoder=model_params.hidden_dim_node_encoder,
+        hidden_dim_edge_encoder=model_params.hidden_dim_edge_encoder,
+        hidden_dim_node_decoder=model_params.hidden_dim_node_decoder,
+        mlp_activation_fn=mlp_act,
+        do_concat_trick=model_params.do_concat_trick,
+        num_processor_checkpoint_segments=model_params.num_processor_checkpoint_segments,
+        recompute_activation=model_params.recompute_activation,
+    ).to(device)
+
+
+def graph_from_batch(batch):
+    return batch[0] if isinstance(batch, (tuple, list)) else batch
+
+
+def resolve_device(device_name: str, manager: DistributedManager, gpuid: int):
+    if manager.world_size > 1:
+        return manager.device
+    if device_name == "cpu":
+        return torch.device("cpu")
+    if device_name in ("cuda", "gpu"):
+        if not torch.cuda.is_available():
+            raise RuntimeError("Config requested cuda device, but torch.cuda.is_available() is false.")
+        return torch.device(f"cuda:{gpuid}")
+    return torch.device(f"cuda:{gpuid}" if torch.cuda.is_available() else "cpu")
+
+
+def main():
+    os.chdir(PROJECT_ROOT)
+    DistributedManager.initialize()
+    manager = DistributedManager()
+    logger = setup_logging(manager.rank)
+
+    config_path = PROJECT_ROOT / "config" / "config.yaml"
+    cfg_model = YParams(config_path, "model")
+    cfg_data = YParams(config_path, "datapipe")
+    cfg_train = YParams(config_path, "training")
+    model_params = cfg_model.specific_params[cfg_model.name]
+
+    datapipe = build_cylinder_flow_datapipe(
+        params=cfg_data,
+        distributed=(manager.world_size > 1),
+        project_root=PROJECT_ROOT,
+    )
+    train_loader, train_sampler = datapipe.train_dataloader()
+    val_loader, val_sampler = datapipe.val_dataloader()
+
+    device = resolve_device(getattr(cfg_train, "device", "auto"), manager, cfg_train.gpuid)
+    logger.info("Using device: %s", device)
+    model = build_model(model_params, 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=cfg_train.lr)
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda step: cfg_train.lr_decay_rate**step)
+    loss_criterion = nn.MSELoss() if cfg_train.loss_criterion == "MSE" else nn.L1Loss()
+    scaler = GradScaler(enabled=bool(cfg_train.amp))
+
+    checkpoint_dir = PROJECT_ROOT / cfg_train.checkpoint_dir
+    epoch_init = load_checkpoint(checkpoint_dir, models=model, optimizer=optimizer, scheduler=scheduler, scaler=scaler, device=device)
+    best_valid_loss = float("inf")
+    best_loss_epoch = epoch_init
+
+    logger.info("Starting training")
+    for epoch in range(epoch_init, cfg_train.max_epoch):
+        if train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+        start = time.time()
+        model.train()
+        train_loss = 0.0
+        for idx, batch in enumerate(train_loader):
+            graph = graph_from_batch(batch).to(device)
+            optimizer.zero_grad(set_to_none=True)
+            with autocast(device_type=device.type, enabled=bool(cfg_train.amp)):
+                pred = model(graph.ndata["x"], graph.edata["x"], graph)
+                loss = loss_criterion(pred, graph.ndata["y"])
+            scaler.scale(loss).backward()
+            scaler.step(optimizer)
+            scaler.update()
+            scheduler.step()
+            train_loss += loss.item()
+            if manager.rank == 0 and (idx + 1) % cfg_train.log_interval == 0:
+                logger.info("Epoch %s/%s batch %s/%s loss %.6f", epoch + 1, cfg_train.max_epoch, idx + 1, len(train_loader), loss.item())
+
+        train_loss /= max(len(train_loader), 1)
+        model.eval()
+        valid_loss = 0.0
+        with torch.no_grad():
+            for batch in val_loader:
+                graph = graph_from_batch(batch).to(device)
+                with autocast(device_type=device.type, enabled=bool(cfg_train.amp)):
+                    pred = model(graph.ndata["x"], graph.edata["x"], graph)
+                    loss = loss_criterion(pred, graph.ndata["y"])
+                valid_loss += loss.item()
+        valid_loss /= max(len(val_loader), 1)
+
+        if manager.rank == 0:
+            logger.info(
+                "Epoch %s finished in %.2fs train_loss %.6f valid_loss %.6f",
+                epoch + 1,
+                time.time() - start,
+                train_loss,
+                valid_loss,
+            )
+            if valid_loss < best_valid_loss:
+                best_valid_loss = valid_loss
+                best_loss_epoch = epoch
+                save_checkpoint(checkpoint_dir, models=model, optimizer=optimizer, scheduler=scheduler, scaler=scaler, epoch=epoch + 1)
+                logger.info("Checkpoint saved to %s", checkpoint_dir)
+            if (epoch - best_loss_epoch) > cfg_train.patience:
+                logger.warning("Early stopping after %s stale epochs", cfg_train.patience)
+                break
+
+    logger.info("Training finished")
+
+
+if __name__ == "__main__":
+    main()