Upload folder using huggingface_hub

README.md

@@ -0,0 +1,172 @@
+---
+frameworks:
+- ""
+language:
+- en
+license: apache-2.0
+tags:
+- OneScience
+- fluid dynamics
+- neural PDE solver evaluation
+- unstructured-mesh simulation
+- multiphysics modeling
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">CFD_Benchmark</span>
+  </strong>
+</p>
+
+# Model Overview
+CFD_Benchmark is an open-source deep-learning benchmark library for research on neural partial differential equation (PDE) solvers. It extends Tsinghua University's open-source Neural-Solver-Library with Distributed Data Parallel (DDP) training support, additional models, and new datasets, while retaining the original neural-operator and physical-field modeling framework. The library supports neural PDE solver evaluation, deep-learning research for CFD, multi-model performance comparisons, large-scale distributed training experiments, physical simulation dataset development, and algorithm benchmarking.
+
+The library currently supports the following benchmarks:
+
+- Six standard benchmarks from [[FNO]](https://arxiv.org/abs/2010.08895) and [[geo-FNO]](https://arxiv.org/abs/2207.05209)
+- PDEBench [[NeurIPS 2022 Track dataset and benchmark]](https://arxiv.org/abs/2210.07182) for autoregressive tasks
+- The ShapeNet-Car dataset [[TOG 2018]](https://dl.acm.org/doi/abs/10.1145/3197517.3201325) for industrial design benchmarks
+- The BubbleML [[Multiphase Multiphysics Dataset]](https://arxiv.org/abs/2307.14623) for studying multiphysics phase-transition phenomena
+
+---
+
+## Supported Neural Solvers
+
+The following neural PDE solvers are supported:
+
+
+- **Transolver** - Transolver: A Fast Transformer Solver for PDEs on General Geometries [[ICML 2024]](https://arxiv.org/abs/2402.02366) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/Transolver.py)
+- **ONO** - Improved Operator Learning by Orthogonal Attention [[ICML 2024]](https://arxiv.org/abs/2310.12487v3) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/ONO.py)
+- **Factformer** - Scalable Transformer for PDE Surrogate Modeling [[NeurIPS 2023]](https://arxiv.org/abs/2305.17560) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/Factformer.py)
+- **U-NO** - U-NO: U-shaped Neural Operators [[TMLR 2023]](https://openreview.net/pdf?id=j3oQF9coJd) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/U_NO.py)
+- **LSM** - Solving High-Dimensional PDEs with Latent Spectral Models [[ICML 2023]](https://arxiv.org/pdf/2301.12664) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/LSM.py)
+- **GNOT** - GNOT: A General Neural Operator Transformer for Operator Learning [[ICML 2023]](https://arxiv.org/abs/2302.14376) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/GNOT.py)
+- **F-FNO** - Factorized Fourier Neural Operators [[ICLR 2023]](https://arxiv.org/abs/2111.13802) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/F_FNO.py)
+- **U-FNO** - An enhanced Fourier neural operator-based deep-learning model for multiphase flow [[Advances in Water Resources 2022]](https://www.sciencedirect.com/science/article/pii/S0309170822000562) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/U_FNO.py)
+- **Galerkin Transformer** - Choose a Transformer: Fourier or Galerkin [[NeurIPS 2021]](https://arxiv.org/abs/2105.14995) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/Galerkin_Transformer.py)
+- **MWT** - Multiwavelet-based Operator Learning for Differential Equations [[NeurIPS 2021]](https://openreview.net/forum?id=LZDiWaC9CGL) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/MWT.py)
+- **FNO** - Fourier Neural Operator for Parametric Partial Differential Equations [[ICLR 2021]](https://arxiv.org/pdf/2010.08895) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/FNO.py)
+- **Transformer** - Attention Is All You Need [[NeurIPS 2017]](https://arxiv.org/pdf/1706.03762) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/Transformer.py)
+
+- **GFNO** - Group Equivariant Fourier Neural Operators for Partial Differential Equations[[2023 Poster]](https://arxiv.org/pdf/1706.03762)[[Code]](https://github.com/divelab/AIRS/blob/main/OpenPDE/G-FNO/models/GFNO.py)
+
+Several vision architectures also serve as effective baselines for structured-geometry tasks:
+
+- **Swin Transformer** - Swin Transformer: Hierarchical Vision Transformer using Shifted Windows [[ICCV 2021]](https://arxiv.org/abs/2103.14030) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/Swin_Transformer.py)
+- **U-Net** - U-Net: Convolutional Networks for Biomedical Image Segmentation [[MICCAI 2015]](https://arxiv.org/pdf/1505.04597) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/U_Net.py)
+
+Several established geometric deep-learning models are included for design tasks:
+
+- **Graph-UNet** - Graph U-Nets [[ICML 2019]](https://arxiv.org/pdf/1905.05178) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/Graph_UNet.py)
+- **GraphSAGE** - Inductive Representation Learning on Large Graphs [[NeurIPS 2017]](https://arxiv.org/pdf/1706.02216) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/GraphSAGE.py)
+- **PointNet** - PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation [[CVPR 2017]](https://arxiv.org/pdf/1612.00593) [[Code]](https://github.com/thuml/Neural-Solver-Library/blob/main/models/PointNet.py)
+
+The library also includes the following graph neural network:
+
+- **MeshGraphNet** LEARNING MESH-BASED SIMULATION WITH GRAPH NETWORKS[ICLR 2021](https://arxiv.org/abs/2010.03409) [[Code]](https://github.com/google-deepmind/deepmind-research/tree/master/meshgraphnets)
+
+
+
+## Use Cases
+
+| Use Case | Description |
+|---|---|
+| Neural PDE solver evaluation | Train, run inference with, and compare models such as FNO, Transolver, GNOT, ONO, and U-NO through a unified workflow |
+| Autoregressive physical prediction | Predict the temporal evolution of PDE states step by step using datasets such as PDEBench |
+| Multiphysics modeling | Study multiphase flows, multiphysics coupling, and phase-transition phenomena using datasets such as BubbleML |
+| 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.
+- 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/CFD_Benchmark --local_dir ./CFD_Benchmark 
+cd CFD_Benchmark 
+```
+
+### 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
+
+Use the dataset links in the benchmark overview above to download the required data.
+
+The six standard benchmark datasets from [[FNO]](https://arxiv.org/abs/2010.08895) and [[geo-FNO]](https://arxiv.org/abs/2207.05209) are available from [this link](https://drive.google.com/drive/folders/1YBuaoTdOSr_qzaow-G-iwvbUI7fiUzu8).
+
+The PDEBench [[NeurIPS 2022 Track dataset and benchmark]](https://arxiv.org/abs/2210.07182), used for benchmarking autoregressive tasks, is available from [this link](https://darus.uni-stuttgart.de/dataset.xhtml?persistentId=doi:10.18419/darus-2986).
+
+The ShapeNet-Car [[TOG 2018]](https://dl.acm.org/doi/abs/10.1145/3197517.3201325) benchmark dataset for industrial design tasks is available from [[this link]](http://www.nobuyuki-umetani.com/publication/mlcfd_data.zip).
+
+The BubbleML [[Multiphase Multiphysics Dataset]](https://arxiv.org/abs/2307.14623), designed for research on multiphysics phase-transition phenomena, is available from [[this link]](https://github.com/HPCForge/BubbleML/blob/main/bubbleml_data/README.md).
+
+The OneScience community also provides the `cfd_benchmark` dataset for training. Download it with the command below and verify that the data path in `conf/config.yaml` is configured correctly:
+```
+modelscope download --dataset OneScience/cfd_benchmark  --local_dir ./data
+```
+### Training
+
+```bash
+python scripts/train.py
+```
+
+### Model Weights
+This repository will provide weights trained on the OneScience `cfd_benchmark` dataset in the `weights/` directory. The weights will be uploaded soon.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+Inference loads the trained weights referenced by `paths.weight_path` and writes the metrics to:
+
+```text
+./results/{train.save_name}/metrics.json
+```
+
+
+# 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
+- Reference repository: [Neural-Solver-Library](https://github.com/thuml/Neural-Solver-Library).
+- This repository retains source attribution and has been adapted for automated execution through OneScience and ModelScope.

config/config.yaml

@@ -0,0 +1,56 @@
+project:
+  name: CFD_Benchmark
+  package: cfd_benchmark
+
+runtime:
+  device: cpu
+
+paths:
+  data_dir: ./data/fake_airfoil
+  checkpoint_dir: ./checkpoints
+  result_dir: ./results
+  weight_path: ./checkpoints/transolver_fake.pt
+
+data:
+  loader: airfoil
+  task: steady
+  ntrain: 2
+  ntest: 1
+  batch_size: 1
+  downsamplex: 20
+  downsampley: 10
+  normalize: false
+  norm_type: UnitTransformer
+  space_dim: 2
+  fun_dim: 0
+  out_dim: 1
+  geotype: unstructured
+
+model:
+  name: Transolver
+  n_hidden: 16
+  n_layers: 1
+  n_heads: 2
+  mlp_ratio: 1
+  dropout: 0.0
+  act: gelu
+  slice_num: 4
+  unified_pos: 0
+  ref: 4
+
+train:
+  epochs: 1
+  lr: 0.001
+  weight_decay: 0.00001
+  optimizer: AdamW
+  scheduler: CosineAnnealingLR
+  pct_start: 0.3
+  max_grad_norm: null
+  derivloss: false
+  use_checkpoint: false
+  checkpoint_layers: ""
+  save_name: transolver_fake
+
+inference:
+  eval: 1
+  vis_num: 0

configuration.json

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

model/DeepONet.py

@@ -0,0 +1,122 @@
+import torch
+import torch.nn as nn
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+
+class Model(nn.Module):
+    """
+    DeepONet 模型。
+    
+    采用双网络架构 (Branch Net 和 Trunk Net) 逼近非线性算子。
+    - Branch Net 处理输入函数 (如初始条件、边界条件)。
+    - Trunk Net 处理目标评估点的坐标 (物理位置)。
+    两者输出特征做内积后得到目标位置的物理场预测值。
+    """
+    def __init__(self, args, device):
+        super(Model, self).__init__()
+        self.__name__ = "DeepONet"
+        self.args = args
+
+        # 位置处理
+        if args.unified_pos and args.geotype != "unstructured":
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            trunk_in_dim = args.ref ** len(args.shapelist)
+        else:
+            trunk_in_dim = args.space_dim
+
+        # 分支网络输入维度（函数输入+时间嵌入）
+        branch_in_dim = args.fun_dim
+        if args.time_input:
+            branch_in_dim += args.n_hidden  # 增加时间嵌入维度
+
+        # 分支网络（函数空间）
+        self.branch_net = StandardMLP(
+            input_dim=branch_in_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden] * args.branch_depth,
+            activation=args.act,
+            use_bias=True
+        )
+
+        # 主干网络（物理空间）
+        self.trunk_net = StandardMLP(
+            input_dim=trunk_in_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden] * args.trunk_depth,
+            activation=args.act,
+            use_bias=True
+        )
+
+        # 时间处理层
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # 输出层 - 支持多个输出通道
+        self.out_layer = nn.Linear(args.n_hidden, args.out_dim)
+        self.bias = nn.Parameter(torch.zeros(1, 1, args.out_dim))
+
+    def structured_forward(self, x, fx, T=None, geo=None):
+        B, N, _ = x.shape  # x: [B, N, space_dim]
+        if self.args.unified_pos:
+            x = self.pos.repeat(B, 1, 1)  # [B, N, d]
+
+        # 处理时间信息
+        if T is not None and self.args.time_input:
+            # 确保时间输入维度正确 [B, 1] -> [B]
+            T = T.view(-1)
+
+            # 生成时间嵌入 [B, D]
+            T_emb = timestep_embedding(T, self.args.n_hidden)
+
+            # 通过时间处理层 [B, D] -> [B, D]
+            T_emb = self.time_fc(T_emb)
+
+            # 扩展时间嵌入以匹配空间点 [B, D] -> [B, N, D]
+            T_emb_expanded = T_emb.unsqueeze(1).expand(-1, N, -1)
+
+            # 将时间嵌入连接到函数输入
+            if fx is not None:
+                # [B, N, fun_dim] + [B, N, D] = [B, N, fun_dim + D]
+                fx = torch.cat([fx, T_emb_expanded], dim=-1)
+            else:
+                fx = T_emb_expanded  # [B, N, D]
+
+        # 分支网络处理
+        branch_feat = self.branch_net(fx)  # [B, N, D]
+
+        # 主干网络处理
+        trunk_feat = self.trunk_net(x)  # [B, N, D]
+        
+        # 点积操作 + 输出层
+        inner = branch_feat * trunk_feat
+        out = self.out_layer(inner) + self.bias
+        return out
+
+    def unstructured_forward(self, x, fx, T=None):
+        B, N, _ = x.shape
+        if T is not None and self.args.time_input:
+            T = T.view(-1)
+            T_emb = timestep_embedding(T, self.args.n_hidden)
+            T_emb = self.time_fc(T_emb)
+            T_emb_expanded = T_emb.unsqueeze(1).expand(-1, N, -1)
+            if fx is not None:
+                fx = torch.cat([fx, T_emb_expanded], dim=-1)
+            else:
+                fx = T_emb_expanded
+
+        branch_feat = self.branch_net(fx)
+        trunk_feat = self.trunk_net(x)
+
+        inner = branch_feat * trunk_feat
+        out = self.out_layer(inner) + self.bias
+        return out
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_forward(x, fx, T)
+        else:
+            return self.structured_forward(x, fx, T)

model/FNO.py

@@ -0,0 +1,194 @@
+import torch
+import math
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+
+from onescience.modules.fourier.fno_layers import (
+    SpectralConv1d,
+    SpectralConv2d,
+    SpectralConv3d,
+)
+from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+
+ConvList = [None, nn.Conv1d, nn.Conv2d, nn.Conv3d]
+
+
+class Model(nn.Module):
+    """
+    傅里叶神经算子 (Fourier Neural Operator, FNO)。
+    支持 1D/2D/3D 结构化网格，以及基于 Geo-FNO 的非结构化网格。
+    """
+    def __init__(self, args, device, s1=96, s2=96):
+        super(Model, self).__init__()
+        self.__name__ = "FNO"
+        self.args = args
+        
+        # ==========================================
+        # 1. Embedding & Preprocess
+        # ==========================================
+        if args.unified_pos and args.geotype != "unstructured":  # structured mesh
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            input_dim = args.fun_dim + args.ref ** len(args.shapelist)
+        else:
+            input_dim = args.fun_dim + args.space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            hidden_dims=[args.n_hidden * 2],
+            output_dim=args.n_hidden,
+            activation=args.act,
+            n_layers=0,
+            res=False,
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # ==========================================
+        # 2. Geometry Projection (GeoFNO 特有)
+        # ==========================================
+        if self.args.geotype == "unstructured":
+            # 明确传入具体参数
+            self.fftproject_in = GeoSpectralConv2d(
+                in_channels=args.n_hidden, 
+                out_channels=args.n_hidden, 
+                modes1=args.modes, 
+                modes2=args.modes, 
+                s1=s1, 
+                s2=s2
+            )
+            self.fftproject_out = GeoSpectralConv2d(
+                in_channels=args.n_hidden, 
+                out_channels=args.n_hidden, 
+                modes1=args.modes, 
+                modes2=args.modes, 
+                s1=s1, 
+                s2=s2
+            )
+            self.iphi = IPHI()
+            self.padding = [(16 - size % 16) % 16 for size in [s1, s2]]
+        else:
+            self.padding = [(16 - size % 16) % 16 for size in args.shapelist]
+
+        # ==========================================
+        # 3. FNO Blocks (显式参数实例化)
+        # ==========================================
+        dim = len(self.padding)
+        
+        if dim == 1:
+            self.conv0 = SpectralConv1d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes)
+            self.conv1 = SpectralConv1d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes)
+            self.conv2 = SpectralConv1d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes)
+            self.conv3 = SpectralConv1d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes)
+        elif dim == 2:
+            self.conv0 = SpectralConv2d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes)
+            self.conv1 = SpectralConv2d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes)
+            self.conv2 = SpectralConv2d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes)
+            self.conv3 = SpectralConv2d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes)
+        elif dim == 3:
+            self.conv0 = SpectralConv3d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes, modes3=args.modes)
+            self.conv1 = SpectralConv3d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes, modes3=args.modes)
+            self.conv2 = SpectralConv3d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes, modes3=args.modes)
+            self.conv3 = SpectralConv3d(in_channels=args.n_hidden, out_channels=args.n_hidden, modes1=args.modes, modes2=args.modes, modes3=args.modes)
+        else:
+            raise ValueError(f"Unsupported dimension: {dim}. Only 1D, 2D, and 3D are supported.")
+
+        # 对应的 1x1 卷积通道混合层
+        self.w0 = ConvList[dim](args.n_hidden, args.n_hidden, 1)
+        self.w1 = ConvList[dim](args.n_hidden, args.n_hidden, 1)
+        self.w2 = ConvList[dim](args.n_hidden, args.n_hidden, 1)
+        self.w3 = ConvList[dim](args.n_hidden, args.n_hidden, 1)
+        
+        # ==========================================
+        # 4. Projectors (输出层)
+        # ==========================================
+        self.fc1 = nn.Linear(args.n_hidden, args.n_hidden)
+        self.fc2 = nn.Linear(args.n_hidden, args.out_dim)
+
+    def structured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 1)
+        
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+            
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden).repeat(1, x.shape[1], 1)
+            Time_emb = self.time_fc(Time_emb)
+            fx = fx + Time_emb
+            
+        x = fx.permute(0, 2, 1).reshape(B, self.args.n_hidden, *self.args.shapelist)
+        
+        # Padding
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 2:
+                x = F.pad(x, [0, self.padding[1], 0, self.padding[0]])
+            elif len(self.args.shapelist) == 3:
+                x = F.pad(x, [0, self.padding[2], 0, self.padding[1], 0, self.padding[0]])
+                
+        # Spectral Convs + Res connections
+        x = F.gelu(self.conv0(x) + self.w0(x))
+        x = F.gelu(self.conv1(x) + self.w1(x))
+        x = F.gelu(self.conv2(x) + self.w2(x))
+        x = self.conv3(x) + self.w3(x)
+
+        # Unpadding
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 2:
+                x = x[..., : -self.padding[0], : -self.padding[1]]
+            elif len(self.args.shapelist) == 3:
+                x = x[..., : -self.padding[0], : -self.padding[1], : -self.padding[2]]
+                
+        x = x.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
+        x = F.gelu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+    def unstructured_geo(self, x, fx, T=None):
+        original_pos = x
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden).repeat(1, x.shape[1], 1)
+            Time_emb = self.time_fc(Time_emb)
+            fx = fx + Time_emb
+
+        # 透传参数到 GeoSpectralConv2d
+        x = self.fftproject_in(
+            fx.permute(0, 2, 1), x_in=original_pos, iphi=self.iphi, code=None
+        )
+
+        x = F.gelu(self.conv0(x) + self.w0(x))
+        x = F.gelu(self.conv1(x) + self.w1(x))
+        x = F.gelu(self.conv2(x) + self.w2(x))
+        x = self.conv3(x) + self.w3(x)
+
+        # 透传参数到 GeoSpectralConv2d
+        x = self.fftproject_out(
+            x, x_out=original_pos, iphi=self.iphi, code=None
+        ).permute(0, 2, 1)
+        
+        x = F.gelu(self.fc1(x))
+        x = self.fc2(x)
+        return x
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/F_FNO.py

@@ -0,0 +1,194 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from onescience.modules.fourier.ffno_layers import (
+    SpectralConv1d,
+    SpectralConv2d,
+    SpectralConv3d,
+)
+from onescience.modules.mlp import StandardMLP
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI
+
+SpectralConvList = [None, SpectralConv1d, SpectralConv2d, SpectralConv3d]
+
+class Model(nn.Module):
+    """
+    Factorized Fourier Neural Operator (F-FNO) 模型。
+    """
+    def __init__(self, args, device, s1=96, s2=96):
+        super(Model, self).__init__()
+        self.__name__ = "F-FNO"
+        self.args = args
+        self.device = device
+        
+        # 1. Embedding & Preprocessing
+        input_dim = args.fun_dim
+        
+        if args.unified_pos and args.geotype != "unstructured":
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            input_dim += args.ref ** len(args.shapelist)
+        else:
+            input_dim += args.space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2], 
+            activation=args.act,
+            use_bias=True
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # 2. Geometry Projection & Spectral Layers
+        self.spectral_layers = nn.ModuleList([])
+        
+        if self.args.geotype == "unstructured":
+            # --- 非结构化网格路径 (GeoFNO) ---
+            self.fftproject_in = GeoSpectralConv2d(
+                in_channels=args.n_hidden,
+                out_channels=args.n_hidden,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            
+            self.fftproject_out = GeoSpectralConv2d(
+                in_channels=args.n_hidden,
+                out_channels=args.n_hidden,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            
+            self.iphi = IPHI()
+            self.padding = [(16 - size % 16) % 16 for size in [s1, s2]]
+            
+            # 中间层使用 FFNO
+            spectral_class = SpectralConv2d
+            conv_args = {
+                "in_dim": args.n_hidden,
+                "out_dim": args.n_hidden,
+                "modes_x": args.modes,
+                "modes_y": args.modes
+            }
+            
+        else:
+            # --- 结构化网格路径 (FFNO) ---
+            self.padding = [(16 - size % 16) % 16 for size in args.shapelist]
+            
+            dim = len(self.padding)
+            spectral_class = SpectralConvList[dim]
+            
+            conv_args = {
+                "in_dim": args.n_hidden,
+                "out_dim": args.n_hidden,
+            }
+            
+            mode_names = ["modes_x", "modes_y", "modes_z"]
+            for i in range(dim):
+                if i < len(mode_names):
+                    conv_args[mode_names[i]] = args.modes
+
+        for _ in range(args.n_layers):
+            self.spectral_layers.append(
+                spectral_class(**conv_args)
+            )
+
+        # 3. Projectors (Decoder)
+        self.fc1 = nn.Linear(args.n_hidden, args.n_hidden)
+        self.fc2 = nn.Linear(args.n_hidden, args.out_dim)
+
+    def structured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+        
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 1)
+            
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            # 【修复】使用广播机制，避免 repeat 导致的维度错误
+            Time_emb = timestep_embedding(T, self.args.n_hidden) # (B, C)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1) # (B, 1, C)
+            fx = fx + Time_emb # Broadcasting: (B, N, C) + (B, 1, C)
+            
+        x = fx.permute(0, 2, 1).reshape(B, self.args.n_hidden, *self.args.shapelist)
+        
+        if not all(item == 0 for item in self.padding):
+            pad_arg = []
+            for p in reversed(self.padding):
+                pad_arg.extend([0, p]) 
+            x = F.pad(x, pad_arg)
+
+        for i in range(self.args.n_layers):
+            x = x + self.spectral_layers[i](x)
+
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 1:
+                x = x[..., :-self.padding[0]]
+            elif len(self.args.shapelist) == 2:
+                x = x[..., :-self.padding[0], :-self.padding[1]]
+            elif len(self.args.shapelist) == 3:
+                x = x[..., :-self.padding[0], :-self.padding[1], :-self.padding[2]]
+        
+        x = x.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def unstructured_geo(self, x, fx, T=None):
+        original_pos = x
+        
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            # 【修复】使用广播机制，避免 repeat 导致的维度错误
+            Time_emb = timestep_embedding(T, self.args.n_hidden) # (B, C)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1) # (B, 1, C)
+            fx = fx + Time_emb # Broadcasting: (B, N, C) + (B, 1, C)
+
+        x = self.fftproject_in(
+            fx.permute(0, 2, 1), x_in=original_pos, iphi=self.iphi, code=None
+        )
+        
+        for i in range(self.args.n_layers):
+            x = x + self.spectral_layers[i](x)
+            
+        x = self.fftproject_out(
+            x, x_out=original_pos, iphi=self.iphi, code=None
+        ).permute(0, 2, 1)
+        
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/Factformer.py

@@ -0,0 +1,115 @@
+import torch
+import torch.nn as nn
+from timm.layers import trunc_normal_
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.factformer_block import Factformer_block
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+
+class Model(nn.Module):
+    """
+    FactFormer 模型。
+    
+    基于 Factorized Attention 的 Transformer 架构，用于处理结构化网格数据。
+    
+    """
+    def __init__(self, args, device):
+        super(Model, self).__init__()
+        self.args = args
+        self.__name__ = "Factformer"
+        # 1. 几何类型检查
+        if args.geotype == "unstructured":
+            raise ValueError(
+                "Factformer does not support unstructured geometry, please try to integrate GeoFNO layer"
+            )
+            
+        # 2. Embedding & Preprocessing 
+        input_dim = args.fun_dim
+        if args.unified_pos:  # only for structured mesh
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            input_dim += args.ref ** len(args.shapelist)
+        else:
+            input_dim += args.space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True, 
+            res=False,
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # 3. Transformer Blocks
+        self.blocks = nn.ModuleList([
+            Factformer_block(
+                num_heads=args.n_heads,
+                hidden_dim=args.n_hidden,
+                dropout=args.dropout,
+                act=args.act,
+                mlp_ratio=args.mlp_ratio,
+                out_dim=args.out_dim,
+                last_layer=(_ == args.n_layers - 1),
+                geotype=args.geotype,
+                shapelist=args.shapelist,
+            )
+            for _ in range(args.n_layers)
+        ])
+        
+        self.placeholder = nn.Parameter(
+            (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
+        )
+        self.initialize_weights()
+
+    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 structured_geo(self, x, fx, T=None):
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 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, :]
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden) # (B, C)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1) # (B, 1, C)
+            fx = fx + Time_emb
+
+        for block in self.blocks:
+            fx = block(fx)
+        return fx
+
+    def unstructured_geo(self, x, fx, T=None):
+        # 保持对非结构化网格的异常抛出
+        raise ValueError(
+            "Factformer does not support unstructured geometry, please try to integrate GeoFNO layer"
+        )
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/GFNO.py

@@ -0,0 +1,264 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from onescience.modules.equivariant.group_conv import GroupEquivariantConv2d
+from onescience.modules.fourier.group_spectral import GSpectralConv2d
+from onescience.modules.mlp.GMLP import GroupEquivariantMLP2d
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI
+
+class GNorm(nn.Module):
+    def __init__(self, width, group_size):
+        super().__init__()
+        self.group_size = group_size
+        self.norm = torch.nn.InstanceNorm3d(width)
+
+    def forward(self, x):
+        # x shape: (B, C*Group, H, W) -> (B, C, Group, H, W)
+        x = x.view(x.shape[0], -1, self.group_size, x.shape[-2], x.shape[-1])
+        x = self.norm(x)
+        x = x.view(x.shape[0], -1, x.shape[-2], x.shape[-1])
+        return x
+
+class Model(nn.Module):
+    """
+    GFNO (Group Factorized Neural Operator) 模型。
+    
+    """
+    def __init__(self, args, device, s1=96, s2=96):
+        super(Model, self).__init__()
+        self.__name__ = "GFNO"
+        self.args = args
+        self.device = device
+        
+        self.in_channels = args.fun_dim
+        self.out_channels = args.out_dim
+        self.modes = args.modes
+        self.width = args.n_hidden
+
+        reflection = False
+
+        # 1. Embedding & Preprocessing
+        if args.unified_pos and args.geotype != "unstructured":
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            in_dim = args.fun_dim + args.ref ** len(args.shapelist)
+        else:
+            in_dim = args.fun_dim + args.space_dim
+        
+        self.preprocess = StandardMLP(
+            input_dim=in_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        # 2. Time Embedding
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+        else:
+            self.time_fc = None
+
+        # 3. Geometry Projection & Output Heads
+        if args.geotype == "unstructured":
+            self.group_size = 4 * (1 + reflection)
+            
+            self.fftproject_in = GeoSpectralConv2d(
+                in_channels=args.n_hidden,
+                out_channels=args.n_hidden,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2,
+            )
+            self.fftproject_out = GeoSpectralConv2d(
+                in_channels=self.width * self.group_size,
+                out_channels=self.width,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2,
+            )
+            self.iphi = IPHI()
+            
+            grid_h, grid_w = s1, s2
+            self.point_q = nn.Sequential(
+                nn.Linear(self.width, self.width * 4),
+                nn.GELU(),
+                nn.Linear(self.width * 4, self.out_channels),
+            )
+        else:
+            grid_h, grid_w = args.shapelist
+            # 结构化输出头 (Explicit instantiation)
+            self.q = GroupEquivariantMLP2d(
+                in_channels=self.width,
+                out_channels=self.out_channels,
+                mid_channels=self.width * 4,
+                reflection=reflection,
+                last_layer=True,
+            )
+
+        self.padding = [(16 - size % 16) % 16 for size in [grid_h, grid_w]]
+
+        # 4. GFNO Stem Layers
+        # Lifting Layer
+        self.p = GroupEquivariantConv2d(
+            in_channels=args.n_hidden,
+            out_channels=self.width,
+            kernel_size=1,
+            reflection=reflection,
+            first_layer=True,
+        )
+        
+        # Spectral Layers
+        self.conv0 = GSpectralConv2d(in_channels=self.width, out_channels=self.width, modes=self.modes, reflection=reflection)
+        self.conv1 = GSpectralConv2d(in_channels=self.width, out_channels=self.width, modes=self.modes, reflection=reflection)
+        self.conv2 = GSpectralConv2d(in_channels=self.width, out_channels=self.width, modes=self.modes, reflection=reflection)
+        self.conv3 = GSpectralConv2d(in_channels=self.width, out_channels=self.width, modes=self.modes, reflection=reflection)
+        
+        self.mlp0 = GroupEquivariantMLP2d(
+            in_channels=self.width,
+            out_channels=self.width,
+            mid_channels=self.width,
+            reflection=reflection
+        )
+        self.mlp1 = GroupEquivariantMLP2d(
+            in_channels=self.width,
+            out_channels=self.width,
+            mid_channels=self.width,
+            reflection=reflection
+        )
+        self.mlp2 = GroupEquivariantMLP2d(
+            in_channels=self.width,
+            out_channels=self.width,
+            mid_channels=self.width,
+            reflection=reflection
+        )
+        self.mlp3 = GroupEquivariantMLP2d(
+            in_channels=self.width,
+            out_channels=self.width,
+            mid_channels=self.width,
+            reflection=reflection
+        )
+        
+        # Residual Weights (1x1 GConv)
+        self.w0 = GroupEquivariantConv2d(
+            in_channels=self.width,
+            out_channels=self.width,
+            kernel_size=1,
+            reflection=reflection
+        )
+        self.w1 = GroupEquivariantConv2d(
+            in_channels=self.width,
+            out_channels=self.width,
+            kernel_size=1,
+            reflection=reflection
+        )
+        self.w2 = GroupEquivariantConv2d(
+            in_channels=self.width,
+            out_channels=self.width,
+            kernel_size=1,
+            reflection=reflection
+        )
+        self.w3 = GroupEquivariantConv2d(
+            in_channels=self.width,
+            out_channels=self.width,
+            kernel_size=1,
+            reflection=reflection
+        )
+        
+        self.norm = GNorm(self.width, group_size=4 * (1 + reflection))
+
+    def _gfno_stem_forward(self, x):
+        # x: (B, C, H, W)
+        x = self.p(x) # Lifting
+
+        # 4 Layers
+        x1 = self.norm(self.conv0(self.norm(x)))
+        x1 = self.mlp0(x1)
+        x = F.gelu(x1 + self.w0(x))
+        
+        x1 = self.norm(self.conv1(self.norm(x)))
+        x1 = self.mlp1(x1)
+        x = F.gelu(x1 + self.w1(x))
+        
+        x1 = self.norm(self.conv2(self.norm(x)))
+        x1 = self.mlp2(x1)
+        x = F.gelu(x1 + self.w2(x))
+        
+        x1 = self.norm(self.conv3(self.norm(x)))
+        x1 = self.mlp3(x1)
+        x = x1 + self.w3(x)
+        return x
+
+    def structured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+
+        if self.args.unified_pos:
+            pos = self.pos.repeat(B, 1, 1)
+            feats = torch.cat([pos, fx], dim=-1) if fx is not None else pos
+        else:
+            feats = torch.cat([x, fx], dim=-1) if fx is not None else x
+
+        feats = self.preprocess(feats) 
+
+        if (T is not None) and (self.time_fc is not None):
+            t_emb = timestep_embedding(T, self.args.n_hidden)
+            t_emb = self.time_fc(t_emb)
+            if t_emb.ndim == 2:
+                t_emb = t_emb.unsqueeze(1)
+            feats = feats + t_emb # Broadcast
+
+        H, W = self.args.shapelist
+        xg = feats.permute(0, 2, 1).reshape(B, self.args.n_hidden, H, W)
+
+        if not all(p == 0 for p in self.padding):
+            xg = F.pad(xg, [0, self.padding[1], 0, self.padding[0]])
+
+        xg = self._gfno_stem_forward(xg)
+
+        if not all(p == 0 for p in self.padding):
+            xg = xg[..., : -self.padding[0], : -self.padding[1]]
+
+        xg = self.q(xg) 
+        out = xg.reshape(B, self.out_channels, -1).permute(0, 2, 1)
+        return out
+
+    def unstructured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+
+        feats = torch.cat([x, fx], dim=-1) if fx is not None else x
+        feats = self.preprocess(feats) 
+
+        if (T is not None) and (self.time_fc is not None):
+            t_emb = timestep_embedding(T, self.args.n_hidden)
+            t_emb = self.time_fc(t_emb)
+            if t_emb.ndim == 2:
+                t_emb = t_emb.unsqueeze(1)
+            feats = feats + t_emb
+
+        xg = self.fftproject_in(
+            feats.permute(0, 2, 1), x_in=x, iphi=self.iphi, code=None
+        )
+
+        xg = self._gfno_stem_forward(xg)
+
+        xp = self.fftproject_out(xg, x_out=x, iphi=self.iphi, code=None).permute(
+            0, 2, 1
+        ) 
+
+        out = self.point_q(xp) 
+        return out
+
+    def forward(self, x, fx=None, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/GNOT.py

@@ -0,0 +1,116 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from timm.layers import trunc_normal_
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.gnot_transformer_block import GNOTTransformerBlock
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+
+class Model(nn.Module):
+    """
+    GNOT (General Neural Operator Transformer) 模型。
+    """
+    def __init__(self, args, device, n_experts=3):
+        super(Model, self).__init__()
+        self.__name__ = "GNOT"
+        self.args = args
+        
+        # 1. Embedding & Preprocessing
+        # -----------------------------------------------------------
+        if args.unified_pos and args.geotype != "unstructured":
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            dim_x = args.ref ** len(args.shapelist)
+            dim_z = args.fun_dim + args.ref ** len(args.shapelist)
+        else:
+            dim_x = args.space_dim
+            dim_z = args.fun_dim + args.space_dim
+
+        self.preprocess_x = StandardMLP(
+            input_dim=dim_x,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+        
+        self.preprocess_z = StandardMLP(
+            input_dim=dim_z,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # 2. Transformer Blocks (MoE Style)
+        # -----------------------------------------------------------
+        self.blocks = nn.ModuleList([
+            GNOTTransformerBlock(
+                num_heads=args.n_heads,
+                hidden_dim=args.n_hidden,
+                dropout=args.dropout,
+                act=args.act,
+                mlp_ratio=args.mlp_ratio,
+                space_dim=args.space_dim,
+                n_experts=n_experts,
+            )
+            for _ in range(args.n_layers)
+        ])
+        
+        self.placeholder = nn.Parameter(
+            (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
+        )
+        
+        # 3. Projectors (Decoder)
+        self.fc1 = nn.Linear(args.n_hidden, args.n_hidden * 2)
+        self.fc2 = nn.Linear(args.n_hidden * 2, args.out_dim)
+        
+        self.initialize_weights()
+
+    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 forward(self, x, fx, T=None, geo=None):
+        pos = x
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 1)
+            
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess_z(fx)
+        else:
+            fx = self.preprocess_z(x)
+            
+        fx = fx + self.placeholder[None, None, :]
+        x = self.preprocess_x(x) # x here becomes embedding of geometric info
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden) # (B, C)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1) # (B, 1, C)
+            fx = fx + Time_emb
+
+        for block in self.blocks:
+            # GNOT block 需要三个参数: x(geo_emb), fx(phys_emb), pos(coords)
+            fx = block(x, fx, pos)
+            
+        fx = self.fc1(fx)
+        fx = F.gelu(fx)
+        fx = self.fc2(fx)
+        return fx

model/Galerkin_Transformer.py

@@ -0,0 +1,252 @@
+# import torch
+# import torch.nn as nn
+# import torch.nn.functional as F
+# import numpy as np
+# from timm.layers import trunc_normal_
+# from onescience.modules.layer.layers.Basic import MLP
+# from onescience.modules.attention.linearattention import LinearAttention
+# from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+# from einops import rearrange, repeat
+# from einops.layers.torch import Rearrange
+
+
+# class Galerkin_Transformer_block(nn.Module):
+#     """Transformer encoder block."""
+
+#     def __init__(
+#         self,
+#         num_heads: int,
+#         hidden_dim: int,
+#         dropout: float,
+#         act="gelu",
+#         mlp_ratio=4,
+#         last_layer=False,
+#         out_dim=1,
+#     ):
+#         super().__init__()
+#         self.last_layer = last_layer
+#         self.ln_1 = nn.LayerNorm(hidden_dim)
+#         self.ln_1a = nn.LayerNorm(hidden_dim)
+#         self.Attn = LinearAttention(
+#             hidden_dim,
+#             heads=num_heads,
+#             dim_head=hidden_dim // num_heads,
+#             dropout=dropout,
+#             attn_type="galerkin",
+#         )
+#         self.ln_2 = nn.LayerNorm(hidden_dim)
+#         self.mlp = MLP(
+#             hidden_dim,
+#             hidden_dim * mlp_ratio,
+#             hidden_dim,
+#             n_layers=0,
+#             res=False,
+#             act=act,
+#         )
+#         if self.last_layer:
+#             self.ln_3 = nn.LayerNorm(hidden_dim)
+#             self.mlp2 = nn.Linear(hidden_dim, out_dim)
+
+#     def forward(self, fx):
+#         fx = self.Attn(self.ln_1(fx), self.ln_1a(fx)) + fx
+#         fx = self.mlp(self.ln_2(fx)) + fx
+#         if self.last_layer:
+#             return self.mlp2(self.ln_3(fx))
+#         else:
+#             return fx
+
+
+# class Model(nn.Module):
+#     ## Galerkin_Transformer
+#     def __init__(self, args, device):
+#         super(Model, self).__init__()
+#         self.__name__ = "Galerkin_Transformer"
+#         self.args = args
+#         ## embedding
+#         if (
+#             args.unified_pos and args.geotype != "unstructured"
+#         ):  # only for structured mesh
+#             self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+#             self.preprocess = MLP(
+#                 args.fun_dim + args.ref ** len(args.shapelist),
+#                 args.n_hidden * 2,
+#                 args.n_hidden,
+#                 n_layers=0,
+#                 res=False,
+#                 act=args.act,
+#             )
+#         else:
+#             self.preprocess = MLP(
+#                 args.fun_dim + args.space_dim,
+#                 args.n_hidden * 2,
+#                 args.n_hidden,
+#                 n_layers=0,
+#                 res=False,
+#                 act=args.act,
+#             )
+#         if args.time_input:
+#             self.time_fc = nn.Sequential(
+#                 nn.Linear(args.n_hidden, args.n_hidden),
+#                 nn.SiLU(),
+#                 nn.Linear(args.n_hidden, args.n_hidden),
+#             )
+
+#         ## models
+#         self.blocks = nn.ModuleList(
+#             [
+#                 Galerkin_Transformer_block(
+#                     num_heads=args.n_heads,
+#                     hidden_dim=args.n_hidden,
+#                     dropout=args.dropout,
+#                     act=args.act,
+#                     mlp_ratio=args.mlp_ratio,
+#                     out_dim=args.out_dim,
+#                     last_layer=(_ == args.n_layers - 1),
+#                 )
+#                 for _ in range(args.n_layers)
+#             ]
+#         )
+#         self.placeholder = nn.Parameter(
+#             (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
+#         )
+#         self.initialize_weights()
+
+#     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 forward(self, x, fx, T=None, geo=None):
+#         if self.args.unified_pos:
+#             x = self.pos.repeat(x.shape[0], 1, 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, :]
+
+#         if T is not None:    
+#             Time_emb = timestep_embedding(T, self.args.n_hidden) # (B, C)
+#             Time_emb = self.time_fc(Time_emb) # (B, C)
+            
+#             if Time_emb.ndim == 2:
+#                 Time_emb = Time_emb.unsqueeze(1) # (B, 1, C)
+                
+#             fx = fx + Time_emb # Broadcasting: (B, N, C) + (B, 1, C) -> (B, N, C)
+
+#         for block in self.blocks:
+#             fx = block(fx)
+#         return fx
+
+
+
+import torch
+import torch.nn as nn
+from timm.layers import trunc_normal_
+
+# --- 引入模块工厂 ---
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.galerkin_transformer_block import Galerkin_Transformer_block
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+
+class Model(nn.Module):
+    """
+    Galerkin Transformer 模型。
+    
+    使用 Galerkin 线性注意力机制处理物理场数据的 Transformer 架构。
+    """
+    def __init__(self, args, device):
+        super(Model, self).__init__()
+        self.__name__ = "Galerkin_Transformer"
+        self.args = args
+        
+        # 1. Embedding & Preprocessing
+        # -----------------------------------------------------------
+        input_dim = args.fun_dim
+        if args.unified_pos and args.geotype != "unstructured":
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            input_dim += args.ref ** len(args.shapelist)
+        else:
+            input_dim += args.space_dim
+
+        # 对应原代码: MLP(input_dim, hidden*2, hidden, n_layers=0)
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2], # 中间层
+            activation=args.act,
+            use_bias=True
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # 2. Transformer Blocks
+        # -----------------------------------------------------------
+        # 使用工厂实例化 Galerkin_Transformer_block
+        self.blocks = nn.ModuleList([
+            Galerkin_Transformer_block(
+                num_heads=args.n_heads,
+                hidden_dim=args.n_hidden,
+                dropout=args.dropout,
+                act=args.act,
+                mlp_ratio=args.mlp_ratio,
+                out_dim=args.out_dim,
+                last_layer=(_ == args.n_layers - 1)
+            )
+            for _ in range(args.n_layers)
+        ])
+        
+        self.placeholder = nn.Parameter(
+            (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
+        )
+        self.initialize_weights()
+
+    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 forward(self, x, fx, T=None, geo=None):
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 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, :]
+
+        if T is not None:    
+            Time_emb = timestep_embedding(T, self.args.n_hidden) # (B, C)
+            Time_emb = self.time_fc(Time_emb) # (B, C)
+            
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1) # (B, 1, C)
+                
+            fx = fx + Time_emb # Broadcasting: (B, N, C) + (B, 1, C) -> (B, N, C)
+
+        for block in self.blocks:
+            fx = block(fx)
+        return fx

model/GraphSAGE.py

@@ -0,0 +1,87 @@
+import torch
+import torch.nn as nn
+import torch_geometric.nn as nng
+from onescience.modules.mlp.MLP import StandardMLP
+
+class Model(nn.Module):
+    """
+    GraphSAGE 模型。
+
+    使用 SAGEConv 进行邻居聚合，并结合 MLP 进行特征编码和解码。
+    """
+    def __init__(self, args, device):
+        super(Model, self).__init__()
+        self.__name__ = "GraphSAGE"
+
+        self.nb_hidden_layers = args.n_layers
+        self.size_hidden_layers = args.n_hidden
+        self.bn_bool = True
+        self.activation = nn.ReLU()
+
+        self.encoder = StandardMLP(
+            input_dim=args.fun_dim + args.space_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+        
+        self.decoder = StandardMLP(
+            input_dim=args.n_hidden,
+            output_dim=args.out_dim,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        # Graph Layers (Keep PyG implementation for consistency)
+        self.in_layer = nng.SAGEConv(
+            in_channels=args.n_hidden, out_channels=self.size_hidden_layers
+        )
+
+        self.hidden_layers = nn.ModuleList()
+        for n in range(self.nb_hidden_layers - 1):
+            self.hidden_layers.append(
+                nng.SAGEConv(
+                    in_channels=self.size_hidden_layers,
+                    out_channels=self.size_hidden_layers,
+                )
+            )
+
+        self.out_layer = nng.SAGEConv(
+            in_channels=self.size_hidden_layers, out_channels=self.size_hidden_layers
+        )
+
+        if self.bn_bool:
+            self.bn = nn.ModuleList()
+            for n in range(self.nb_hidden_layers):
+                self.bn.append(
+                    nn.BatchNorm1d(self.size_hidden_layers, track_running_stats=False)
+                )
+
+    def forward(self, x, fx, T=None, geo=None):
+        if x.dim() == 3:
+            x = x.squeeze(0)  # [1, N, C] → [N, C]
+        if fx is not None and fx.dim() == 3:
+            fx = fx.squeeze(0)  # [1, N, C] → [N, C]
+        if geo.dim() == 3:
+            edge_index = geo.squeeze(0)  # [1, 2, E] → [2, E]
+        else:
+            edge_index = geo
+
+        z = torch.cat((x, fx), dim=-1)
+        z = self.encoder(z)
+        z = self.in_layer(z, edge_index)
+        if self.bn_bool:
+            z = self.bn[0](z)
+        z = self.activation(z)
+
+        for n in range(self.nb_hidden_layers - 1):
+            z = self.hidden_layers[n](z, edge_index)
+            if self.bn_bool:
+                z = self.bn[n + 1](z)
+            z = self.activation(z)
+
+        z = self.out_layer(z, edge_index)
+        z = self.decoder(z)
+        return z.unsqueeze(0)

model/Graph_UNet.py

@@ -0,0 +1,200 @@
+
+import torch
+import torch.nn as nn
+import torch_geometric.nn as nng
+import random
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.sample.SpatialGraphDownsample import SpatialGraphDownsample
+from onescience.modules.sample.SpatialGraphUpsample import SpatialGraphUpsample
+
+class Model(nn.Module):
+    """
+    Graph U-Net 模型。
+    
+    基于图神经网络的 U-Net 结构，包含 Encoder-Decoder 和 Skip Connections。
+    使用 SpatialGraphDownsample 进行图池化，使用 SpatialGraphUpsample 进行反池化。
+    """
+    def __init__(
+        self,
+        args,
+        device,
+        pool="random",
+        scale=5,
+        list_r=[0.05, 0.2, 0.5, 1, 10],
+        pool_ratio=[0.5, 0.5, 0.5, 0.5, 0.5],
+        max_neighbors=64,
+        layer="SAGE",
+        head=2,
+    ):
+        super(Model, self).__init__()
+        self.__name__ = "Graph_UNet"
+        
+        # 参数绑定
+        self.L = scale
+        self.layer = layer
+        self.pool_type = pool
+        self.pool_ratio = pool_ratio
+        self.list_r = list_r
+        self.size_hidden = args.n_hidden
+        self.dim_enc = args.n_hidden
+        self.bn_bool = True
+        self.res = False
+        self.head = head
+        self.activation = nn.ReLU()
+        
+        self.encoder = StandardMLP(
+            input_dim=args.fun_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+        
+        self.decoder = StandardMLP(
+            input_dim=args.n_hidden,
+            output_dim=args.out_dim,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        # Down Path Layers
+        self.down_convs = nn.ModuleList()
+        self.down_samples = nn.ModuleList()
+        self.down_bns = nn.ModuleList()
+
+        # Level 0 (Initial)
+        self._add_conv_layer(self.down_convs, self.dim_enc, self.size_hidden)
+        if self.bn_bool:
+            self._add_bn_layer(self.down_bns, self.size_hidden)
+
+        # Level 1 to L-1
+        current_dim = self.size_hidden
+        for n in range(self.L - 1):
+            self.down_samples.append(
+                SpatialGraphDownsample(
+                    in_channels=current_dim,
+                    ratio=self.pool_ratio[n],
+                    r=self.list_r[n],
+                    max_num_neighbors=max_neighbors,
+                    pool_method=self.pool_type
+                )
+            )
+            
+            # Conv Layer
+            in_c = current_dim
+            out_c = 2 * current_dim if layer == "SAGE" else current_dim
+            self._add_conv_layer(self.down_convs, in_c, out_c)
+            current_dim = out_c
+            
+            if self.bn_bool:
+                self._add_bn_layer(self.down_bns, current_dim)
+
+        # Up Path Layers
+        self.up_convs = nn.ModuleList()
+        
+        # --- 3. Upsample Module ---
+        self.up_sampler = SpatialGraphUpsample()
+        
+        self.up_bns = nn.ModuleList()
+        
+        curr_h_init = args.n_hidden
+        
+        # Up Layer 0 (Top Layer)
+        if self.layer == "SAGE":
+            self.up_convs.append(nng.SAGEConv(3 * curr_h_init, self.dim_enc))
+            curr_h_init = 2 * curr_h_init
+        elif self.layer == "GAT":
+            self.up_convs.append(nng.GATConv(2 * self.head * curr_h_init, self.dim_enc, heads=2, concat=False))
+        
+        if self.bn_bool:
+             self.up_bns.append(nng.BatchNorm(self.dim_enc, track_running_stats=False))
+
+        # Up Layer 1 to L-1 (Middle Layers)
+        for n in range(1, self.L - 1):
+            if self.layer == "SAGE":
+                self.up_convs.append(nng.SAGEConv(3 * curr_h_init, curr_h_init))
+                bn_dim = curr_h_init
+                curr_h_init = 2 * curr_h_init
+            elif self.layer == "GAT":
+                self.up_convs.append(nng.GATConv(2 * self.head * curr_h_init, curr_h_init, heads=2, concat=True))
+                bn_dim = curr_h_init * 2 # GAT concat=True
+            
+            if self.bn_bool:
+                self.up_bns.append(nng.BatchNorm(bn_dim, track_running_stats=False))
+
+    def _add_conv_layer(self, module_list, in_c, out_c):
+        if self.layer == "SAGE":
+            module_list.append(nng.SAGEConv(in_c, out_c))
+        elif self.layer == "GAT":
+            module_list.append(nng.GATConv(in_c, out_c, heads=self.head, concat=True, add_self_loops=False))
+
+    def _add_bn_layer(self, module_list, in_c):
+        dim = in_c * self.head if self.layer == "GAT" else in_c
+        module_list.append(nng.BatchNorm(dim, track_running_stats=False))
+
+    def forward(self, x, fx, T=None, geo=None):
+        if geo is None: raise ValueError("Edge index required")
+        if fx.dim() == 3: fx = fx.squeeze(0)
+        if geo.dim() == 3: edge_index = geo.squeeze(0)
+        else: edge_index = geo
+        
+        # Encoder
+        z = self.encoder(fx)
+        if self.res: z_res = z.clone()
+
+        # Downsampling Path
+        skip_connections = [] 
+        pos_history = []     
+        edge_index_history = [edge_index.clone()]
+        
+        # Level 0 Conv
+        z = self.down_convs[0](z, edge_index)
+        if self.bn_bool: z = self.down_bns[0](z)
+        z = self.activation(z)
+        
+        skip_connections.append(z.clone())
+        
+        # Assuming x contains coords in first 2 columns as per original code logic
+        current_pos = x[:, :2] 
+        pos_history.append(current_pos.clone())
+
+        # Levels 1 to L-1
+        for n in range(self.L - 1):
+            z, current_pos, edge_index, _ = self.down_samples[n](z, current_pos, edge_index)
+            
+            pos_history.append(current_pos.clone())
+            edge_index_history.append(edge_index.clone())
+
+            z = self.down_convs[n+1](z, edge_index)
+            if self.bn_bool: z = self.down_bns[n+1](z)
+            z = self.activation(z)
+            
+            skip_connections.append(z.clone())
+        
+        # Up Path
+        for n in range(self.L - 1, 0, -1):
+            layer_idx = n - 1
+            
+            pos_low = pos_history[n]
+            pos_high = pos_history[n-1]
+            z_skip = skip_connections[n-1]
+            
+            target_edge_index = edge_index_history[n-1]
+
+            z = self.up_sampler(z, pos_low, pos_high)
+            
+            z = torch.cat([z, z_skip], dim=1)
+            
+            z = self.up_convs[layer_idx](z, target_edge_index)
+            
+            if self.bn_bool: 
+                z = self.up_bns[layer_idx](z)
+            
+            if n != 1:
+                z = self.activation(z)
+
+        # Decoder
+        if self.res: z = z + z_res
+        z = self.decoder(z)
+        return z.unsqueeze(0)

model/LSM.py

@@ -0,0 +1,262 @@
+
+import torch
+import math
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.Neural_Spectral_Block import (
+    NeuralSpectralBlock1D,
+    NeuralSpectralBlock2D,
+    NeuralSpectralBlock3D,
+)
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI
+from onescience.modules.layer.unet_layer import (
+    DoubleConv1D, Down1D, Up1D, OutConv1D,
+    DoubleConv2D, Down2D, Up2D, OutConv2D,
+    DoubleConv3D, Down3D, Up3D, OutConv3D,
+)
+
+ConvList = [None, DoubleConv1D, DoubleConv2D, DoubleConv3D]
+DownList = [None, Down1D, Down2D, Down3D]
+UpList = [None, Up1D, Up2D, Up3D]
+OutList = [None, OutConv1D, OutConv2D, OutConv3D]
+NeuralSpectralBlockList = [
+    None,
+    NeuralSpectralBlock1D,
+    NeuralSpectralBlock2D,
+    NeuralSpectralBlock3D,
+]
+
+class Model(nn.Module):
+    """
+    LSM (Latent Spectral Model) 模型。
+    结合了 U-Net 的多尺度特征提取与 Neural Spectral Block (Latent Transformer) 的全局谱处理能力。
+    """
+    def __init__(
+        self, args, device, bilinear=True, num_token=4, num_basis=12, s1=96, s2=96
+    ):
+        super(Model, self).__init__()
+        self.__name__ = "LSM"
+        self.args = args
+        
+        if args.task == "steady":
+            normtype = "bn"
+        else:
+            normtype = "in"  
+
+        # 1. Embedding & Preprocessing 
+        if args.unified_pos and args.geotype != "unstructured":
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            in_dim = args.fun_dim + args.ref ** len(args.shapelist)
+        else:
+            in_dim = args.fun_dim + args.space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=in_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+        else:
+            self.time_fc = None
+
+        # 2. Geometry Projection 
+        if self.args.geotype == "unstructured":
+            self.fftproject_in = GeoSpectralConv2d(
+                in_channels=args.n_hidden,
+                out_channels=args.n_hidden,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            self.fftproject_out = GeoSpectralConv2d(
+                in_channels=args.n_hidden,
+                out_channels=args.n_hidden,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            self.iphi = IPHI()
+            
+            patch_size = [(size + (16 - size % 16) % 16) // 16 for size in [s1, s2]]
+            self.padding = [(16 - size % 16) % 16 for size in [s1, s2]]
+        else:
+            patch_size = [
+                (size + (16 - size % 16) % 16) // 16 for size in args.shapelist
+            ]
+            self.padding = [(16 - size % 16) % 16 for size in args.shapelist]
+
+        # 3. Multiscale U-Net Modules
+        # -----------------------------------------------------------
+        dim = len(patch_size)
+        
+        self.inc = ConvList[dim](args.n_hidden, args.n_hidden, normtype=normtype)
+        self.down1 = DownList[dim](args.n_hidden, args.n_hidden * 2, normtype=normtype)
+        self.down2 = DownList[dim](args.n_hidden * 2, args.n_hidden * 4, normtype=normtype)
+        self.down3 = DownList[dim](args.n_hidden * 4, args.n_hidden * 8, normtype=normtype)
+        
+        factor = 2 if bilinear else 1
+        self.down4 = DownList[dim](args.n_hidden * 8, args.n_hidden * 16 // factor, normtype=normtype)
+        
+        self.up1 = UpList[dim](args.n_hidden * 16, args.n_hidden * 8 // factor, bilinear, normtype=normtype)
+        self.up2 = UpList[dim](args.n_hidden * 8, args.n_hidden * 4 // factor, bilinear, normtype=normtype)
+        self.up3 = UpList[dim](args.n_hidden * 4, args.n_hidden * 2 // factor, bilinear, normtype=normtype)
+        self.up4 = UpList[dim](args.n_hidden * 2, args.n_hidden, bilinear, normtype=normtype)
+        self.outc = OutList[dim](args.n_hidden, args.n_hidden)
+
+        # 4. Patchified Neural Spectral Blocks 
+        block_class = NeuralSpectralBlockList[dim]
+        
+        self.process1 = block_class(
+            width=args.n_hidden,
+            num_basis=num_basis,
+            patch_size=patch_size,
+            num_token=num_token,
+            n_heads=args.n_heads
+        )
+        self.process2 = block_class(
+            width=args.n_hidden * 2,
+            num_basis=num_basis,
+            patch_size=patch_size,
+            num_token=num_token,
+            n_heads=args.n_heads
+        )
+        self.process3 = block_class(
+            width=args.n_hidden * 4,
+            num_basis=num_basis,
+            patch_size=patch_size,
+            num_token=num_token,
+            n_heads=args.n_heads
+        )
+        self.process4 = block_class(
+            width=args.n_hidden * 8,
+            num_basis=num_basis,
+            patch_size=patch_size,
+            num_token=num_token,
+            n_heads=args.n_heads
+        )
+        self.process5 = block_class(
+            width=args.n_hidden * 16 // factor,
+            num_basis=num_basis,
+            patch_size=patch_size,
+            num_token=num_token,
+            n_heads=args.n_heads
+        )
+
+        # 5. Projectors
+        # -----------------------------------------------------------
+        self.fc1 = nn.Linear(args.n_hidden, args.n_hidden * 2)
+        self.fc2 = nn.Linear(args.n_hidden * 2, args.out_dim)
+
+    def structured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+        
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 1)
+            
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+            
+        x = fx.permute(0, 2, 1).reshape(B, self.args.n_hidden, *self.args.shapelist)
+        
+        if not all(item == 0 for item in self.padding):
+            pad_arg = []
+            for p in reversed(self.padding):
+                pad_arg.extend([0, p])
+            x = F.pad(x, pad_arg)
+
+        # LSM 核心处理流
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        
+        x = self.up1(self.process5(x5), self.process4(x4))
+        x = self.up2(x, self.process3(x3))
+        x = self.up3(x, self.process2(x2))
+        x = self.up4(x, self.process1(x1))
+        x = self.outc(x)
+
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 1:
+                x = x[..., : -self.padding[0]]
+            elif len(self.args.shapelist) == 2:
+                x = x[..., : -self.padding[0], : -self.padding[1]]
+            elif len(self.args.shapelist) == 3:
+                x = x[..., : -self.padding[0], : -self.padding[1], : -self.padding[2]]
+                
+        x = x.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def unstructured_geo(self, x, fx, T=None):
+        original_pos = x
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+
+        x = self.fftproject_in(
+            fx.permute(0, 2, 1), x_in=original_pos, iphi=self.iphi, code=None
+        )
+        
+        # LSM 核心处理流
+        x1 = self.inc(x)
+        x2 = self.down1(x1)
+        x3 = self.down2(x2)
+        x4 = self.down3(x3)
+        x5 = self.down4(x4)
+        
+        x = self.up1(self.process5(x5), self.process4(x4))
+        x = self.up2(x, self.process3(x3))
+        x = self.up3(x, self.process2(x2))
+        x = self.up4(x, self.process1(x1))
+        x = self.outc(x)
+        
+        x = self.fftproject_out(
+            x, x_out=original_pos, iphi=self.iphi, code=None
+        ).permute(0, 2, 1)
+        
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/MWT.py

@@ -0,0 +1,217 @@
+import torch
+import math
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+from timm.layers import trunc_normal_
+from onescience.modules.fourier.MultiWaveletTransform import (
+    MultiWaveletTransform1D,
+    MultiWaveletTransform2D,
+    MultiWaveletTransform3D,
+)
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI
+
+MultiWaveletTransformList = [
+    None,
+    MultiWaveletTransform1D,
+    MultiWaveletTransform2D,
+    MultiWaveletTransform3D,
+]
+
+class Model(nn.Module):
+    # this model requires H = W = Z and H, W, Z is the power of two
+    def __init__(
+        self, args, device, alpha=2, L=0, c=1, base="legendre", s1=128, s2=128
+    ):
+        super(Model, self).__init__()
+        self.__name__ = "MWT"
+        self.args = args
+        self.k = args.mwt_k
+        self.WMT_dim = c * self.k**2
+        if args.geotype == "structured_1D":
+            self.WMT_dim = c * self.k
+        self.c = c
+        self.s1 = s1
+        self.s2 = s2
+        
+        ## embedding
+        if (
+            args.unified_pos and args.geotype != "unstructured"
+        ):  # only for structured mesh
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            self.preprocess = StandardMLP(
+                input_dim=args.fun_dim + args.ref ** len(args.shapelist),
+                output_dim=self.WMT_dim,
+                hidden_dims=[args.n_hidden * 2],
+                activation=args.act,
+                use_bias=True
+            )
+        else:
+            self.preprocess = StandardMLP(
+                input_dim=args.fun_dim + args.space_dim,
+                output_dim=self.WMT_dim,
+                hidden_dims=[args.n_hidden * 2],
+                activation=args.act,
+                use_bias=True
+            )
+            
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(self.WMT_dim, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, self.WMT_dim),
+            )
+            
+        # geometry projection
+        if self.args.geotype == "unstructured":
+            self.fftproject_in = GeoSpectralConv2d(
+                in_channels=self.WMT_dim,
+                out_channels=self.WMT_dim,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            self.fftproject_out = GeoSpectralConv2d(
+                in_channels=self.WMT_dim,
+                out_channels=self.WMT_dim,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            self.iphi = IPHI()
+            self.augmented_resolution = [s1, s2]
+            self.padding = [(16 - size % 16) % 16 for size in [s1, s2]]
+        else:
+            target = 2 ** (math.ceil(np.log2(max(args.shapelist))))
+            self.padding = [(target - size) for size in args.shapelist]
+            self.augmented_resolution = [target for _ in range(len(self.padding))]
+
+        dim = len(self.padding)
+        transform_class = MultiWaveletTransformList[dim]
+
+        self.spectral_layers = nn.ModuleList(
+            [
+                transform_class(
+                    k=self.k,
+                    alpha=alpha,
+                    L=L,
+                    c=c,
+                    base=base
+                )
+                for _ in range(args.n_layers)
+            ]
+        )
+        
+        # projectors
+        self.fc1 = nn.Linear(self.WMT_dim, args.n_hidden)
+        self.fc2 = nn.Linear(args.n_hidden, args.out_dim)
+
+    def structured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 1)
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.WMT_dim)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+            
+        x = fx.permute(0, 2, 1).reshape(B, self.WMT_dim, *self.args.shapelist)
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 2:
+                x = F.pad(x, [0, self.padding[1], 0, self.padding[0]])
+            elif len(self.args.shapelist) == 3:
+                x = F.pad(
+                    x, [0, self.padding[2], 0, self.padding[1], 0, self.padding[0]]
+                )
+        x = (
+            x.reshape(B, self.WMT_dim, -1)
+            .permute(0, 2, 1)
+            .contiguous()
+            .reshape(
+                B,
+                *self.augmented_resolution,
+                self.c,
+                self.k**2 if self.args.geotype != "structured_1D" else self.k
+            )
+        )
+        for i in range(self.args.n_layers):
+            x = self.spectral_layers[i](x)
+            if i < self.args.n_layers - 1:
+                x = F.gelu(x)
+        x = (
+            x.reshape(B, -1, self.WMT_dim)
+            .permute(0, 2, 1)
+            .contiguous()
+            .reshape(B, self.WMT_dim, *self.augmented_resolution)
+        )
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 2:
+                x = x[..., : -self.padding[0], : -self.padding[1]]
+            elif len(self.args.shapelist) == 3:
+                x = x[..., : -self.padding[0], : -self.padding[1], : -self.padding[2]]
+        x = x.reshape(B, self.WMT_dim, -1).permute(0, 2, 1)
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def unstructured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+        original_pos = x
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.WMT_dim)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+
+        x = self.fftproject_in(
+            fx.permute(0, 2, 1), x_in=original_pos, iphi=self.iphi, code=None
+        )
+        x = (
+            x.reshape(B, self.WMT_dim, -1)
+            .permute(0, 2, 1)
+            .contiguous()
+            .reshape(B, *self.augmented_resolution, self.c, self.k**2)
+        )
+        for i in range(self.args.n_layers):
+            x = self.spectral_layers[i](x)
+            if i < self.args.n_layers - 1:
+                x = F.gelu(x)
+        x = (
+            x.reshape(B, -1, self.WMT_dim)
+            .permute(0, 2, 1)
+            .contiguous()
+            .reshape(B, self.WMT_dim, *self.augmented_resolution)
+        )
+        x = self.fftproject_out(
+            x, x_out=original_pos, iphi=self.iphi, code=None
+        ).permute(0, 2, 1)
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/MeshGraphNet.py

@@ -0,0 +1,252 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+
+try:
+    import dgl
+    from dgl import DGLGraph
+except ImportError:
+    pass 
+
+from dataclasses import dataclass
+from itertools import chain
+from typing import Callable, List, Tuple, Union
+
+# --- 引入模块工厂 ---
+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
+    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 Model(Module):
+    """
+    LSMMeshGraphNet 网络架构 (Refactored).
+    
+    使用网格图的 MLP、边更新和节点更新模块构建。
+    """
+
+    def __init__(
+        self,
+        args,
+        device,
+        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())
+        self.__name__ = "LSMMeshGraphNet"
+        
+        # 参数绑定
+        self.input_dim_nodes = args.fun_dim
+        self.input_dim_edges = 4 
+        self.output_dim = args.out_dim
+        
+        activation_fn = get_activation(mlp_activation_fn)
+
+        # 1. Edge Encoder
+        self.edge_encoder = MeshGraphMLP(
+            input_dim=self.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=self.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=self.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
+        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 block constructed from edge and node update modules.
+    """
+
+    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

model/ONO.py

@@ -0,0 +1,104 @@
+import torch
+import torch.nn as nn
+from timm.layers import trunc_normal_
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.orthogonal_neural_block import OrthogonalNeuralBlock
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+
+class Model(nn.Module):
+    """
+    Orthogonal Neural Operator (ONO) 模型。
+    """
+    def __init__(self, args, device):
+        super(Model, self).__init__()
+        self.__name__ = "ONO"
+        self.args = args
+        
+        # Embedding & Preprocessing 
+        if args.unified_pos and args.geotype != "unstructured":
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            dim_x = args.ref ** len(args.shapelist)
+            dim_z = args.fun_dim + args.ref ** len(args.shapelist)
+        else:
+            dim_x = args.fun_dim + args.space_dim
+            dim_z = args.fun_dim + args.space_dim
+
+        self.preprocess_x = StandardMLP(
+            input_dim=dim_x,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+        self.preprocess_z = StandardMLP(
+            input_dim=dim_z,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # ONO Blocks 
+        self.blocks = nn.ModuleList([
+            OrthogonalNeuralBlock(
+                num_heads=args.n_heads,
+                hidden_dim=args.n_hidden,
+                dropout=args.dropout,
+                act=args.act,
+                attn_type=args.attn_type,
+                mlp_ratio=args.mlp_ratio,
+                last_layer=(_ == args.n_layers - 1),
+                psi_dim=args.psi_dim,
+                out_dim=args.out_dim,
+            )
+            for _ in range(args.n_layers)
+        ])
+        
+        self.placeholder = nn.Parameter(
+            (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
+        )
+        self.initialize_weights()
+
+    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 forward(self, x, fx, T=None, geo=None):
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 1)
+            
+        if fx is not None:
+            x = torch.cat((x, fx), -1)
+            fx = self.preprocess_z(x)
+            x = self.preprocess_x(x)
+        else:
+            fx = self.preprocess_z(x)
+            x = self.preprocess_x(x)
+            
+        fx = fx + self.placeholder[None, None, :]
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+
+        for block in self.blocks:
+            x, fx = block(x, fx)
+        return fx

model/PointNet.py

@@ -0,0 +1,100 @@
+import torch
+import torch.nn as nn
+import torch_geometric.nn as nng
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+from onescience.modules.mlp.MLP import StandardMLP
+
+class Model(nn.Module):
+    """
+    PointNet 模型。
+    
+    用于处理点云数据，通过 MLP 提取局部特征，并使用全局最大池化提取全局特征。
+    """
+    def __init__(self, args, device):
+        super(Model, self).__init__()
+        self.__name__ = "PointNet"
+
+        # 1. Input Block
+        self.in_block = StandardMLP(
+            input_dim=args.n_hidden,
+            output_dim=args.n_hidden * 2,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        # 2. Max Pooling Block
+        self.max_block = StandardMLP(
+            input_dim=args.n_hidden * 2,
+            output_dim=args.n_hidden * 32,
+            hidden_dims=[args.n_hidden * 8],
+            activation=args.act,
+            use_bias=True
+        )
+
+        # 3. Output Block
+        self.out_block = StandardMLP(
+            input_dim=args.n_hidden * (2 + 32), # 34 * hidden
+            output_dim=args.n_hidden * 4,
+            hidden_dims=[args.n_hidden * 16],
+            activation=args.act,
+            use_bias=True
+        )
+
+        # 4. Encoder
+        self.encoder = StandardMLP(
+            input_dim=args.fun_dim + args.space_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        # 5. Decoder
+        self.decoder = StandardMLP(
+            input_dim=args.n_hidden,
+            output_dim=args.out_dim,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        self.fcfinal = nn.Linear(args.n_hidden * 4, args.n_hidden)
+
+    def forward(self, x, fx, T=None, geo=None):
+        if geo is None:
+            raise ValueError("Please provide edge index for Graph Neural Networks")
+        
+        # 兼容 batch_size = 1 输入
+        if x.dim() == 3:
+            x = x.squeeze(0)  # [1, N, C] → [N, C]
+        if fx is not None and fx.dim() == 3:
+            fx = fx.squeeze(0)
+
+        assert (
+            x.size(0) > 0 # Simple check
+        ), "Input cannot be empty"
+
+        # 构造 batch 索引
+        batch = torch.zeros(x.shape[0], dtype=torch.long, device=x.device)
+
+        # 编码 + 局部特征提取
+        z = torch.cat((x, fx), dim=-1).float()
+        z = self.encoder(z)
+        z = self.in_block(z)
+
+        # 全局特征（max pooling）
+        global_coef = self.max_block(z)
+        global_coef = nng.global_max_pool(global_coef, batch=batch)
+
+        # 重复 global coef 到每个点
+        nb_points = torch.tensor([batch.shape[0]], device=z.device)
+        global_coef = global_coef.repeat_interleave(nb_points, dim=0)
+
+        # 拼接全局 + 局部特征
+        z = torch.cat([z, global_coef], dim=1)
+        z = self.out_block(z)
+        z = self.fcfinal(z)
+        z = self.decoder(z)
+
+        return z.unsqueeze(0)

model/RegDGCNN.py

@@ -0,0 +1,361 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import os
+import copy
+import math
+import numpy as np
+from torch_geometric.data import Data, DataLoader
+from torch_geometric.nn import (
+    GCNConv,
+    GATConv,
+    global_mean_pool,
+    global_max_pool,
+    JumpingKnowledge,
+)
+from torch.nn import Sequential, Linear, ReLU, BatchNorm1d, Dropout
+from torch_geometric.nn import BatchNorm
+
+
+def knn(x, k):
+    """
+    Computes the k-nearest neighbors for each point in x.
+
+    Args:
+        x (torch.Tensor): The input tensor of shape (batch_size, num_dims, num_points).
+        k (int): The number of nearest neighbors to find.
+
+    Returns:
+        torch.Tensor: Indices of the k-nearest neighbors for each point, shape (batch_size, num_points, k).
+    """
+    # Calculate pairwise distance, shape (batch_size, num_points, num_points)
+    inner = -2 * torch.matmul(x.transpose(2, 1), x)
+    xx = torch.sum(x**2, dim=1, keepdim=True)
+    pairwise_distance = -xx - inner - xx.transpose(2, 1)
+
+    # Retrieve the indices of the k nearest neighbors
+    idx = pairwise_distance.topk(k=k, dim=-1)[1]
+    return idx
+
+
+def get_graph_feature(x, k=20, idx=None):
+    """
+    Constructs local graph features for each point by finding its k-nearest neighbors and
+    concatenating the relative position vectors.
+
+    Args:
+        x (torch.Tensor): The input tensor of shape (batch_size, num_dims, num_points).
+        k (int): The number of neighbors to consider for graph construction.
+        idx (torch.Tensor, optional): Precomputed k-nearest neighbor indices.
+
+    Returns:
+        torch.Tensor: The constructed graph features of shape (batch_size, 2*num_dims, num_points, k).
+    """
+    batch_size = x.size(0)
+    num_points = x.size(2)
+    x = x.view(batch_size, -1, num_points)
+
+    # Compute k-nearest neighbors if not provided
+    if idx is None:
+        idx = knn(x, k=k)
+
+    # Prepare indices for gathering
+    device = x.device
+    idx_base = torch.arange(0, batch_size, device=device).view(-1, 1, 1) * num_points
+    idx = idx + idx_base
+    idx = idx.view(-1)
+
+    _, num_dims, _ = x.size()
+    x = x.transpose(2, 1).contiguous()
+
+    # Gather neighbors for each point to construct local regions
+    feature = x.view(batch_size * num_points, -1)[idx, :]
+    feature = feature.view(batch_size, num_points, k, num_dims)
+
+    # Expand x to match the dimensions for broadcasting subtraction
+    x = x.view(batch_size, num_points, 1, num_dims).repeat(1, 1, k, 1)
+
+    # Concatenate the original point features with the relative positions to form the graph features
+    feature = torch.cat((feature - x, x), dim=3).permute(0, 3, 1, 2).contiguous()
+
+    return feature
+
+
+# class Model(nn.Module):
+#     """
+#     Deep Graph Convolutional Neural Network for Regression Tasks (RegDGCNN) for processing 3D point cloud data.
+
+#     This network architecture extracts hierarchical features from point clouds using graph-based convolutions,
+#     enabling effective learning of spatial structures.
+#     """
+
+#     def __init__(self, args, device):
+#         """
+#         Initializes the RegDGCNN model with specified configurations.
+
+#         Args:
+#             args (dict): Configuration parameters including 'k' for the number of neighbors, 'emb_dims' for embedding
+#             dimensions, and 'dropout' rate.
+#             output_channels (int): Number of output channels (e.g., for drag prediction, this is 1).
+#         """
+#         super(Model, self).__init__()
+#         self.__name__ = 'RegDGCNN'
+#         self.args = args
+#         self.k = 20  # Number of nearest neighbors
+
+#         # Batch normalization layers to stabilize and accelerate training
+#         self.bn1 = nn.BatchNorm2d(args.n_hidden)
+#         self.bn2 = nn.BatchNorm2d(args.n_hidden * 2)
+#         self.bn3 = nn.BatchNorm2d(args.n_hidden * 2)
+#         self.bn4 = nn.BatchNorm2d(args.n_hidden * 4)
+#         self.bn5 = nn.BatchNorm1d(args.emb_dims)
+
+#         # EdgeConv layers: Convolutional layers leveraging local neighborhood information
+#         self.conv1 = nn.Sequential(nn.Conv2d(2*(args.fun_dim + args.space_dim), args.n_hidden, kernel_size=1, bias=False),
+#                                    self.bn1,
+#                                    nn.LeakyReLU(negative_slope=0.2))
+#         self.conv2 = nn.Sequential(nn.Conv2d(args.n_hidden * 2, args.n_hidden * 2, kernel_size=1, bias=False),
+#                                    self.bn2,
+#                                    nn.LeakyReLU(negative_slope=0.2))
+#         self.conv3 = nn.Sequential(nn.Conv2d(args.n_hidden * 4, args.n_hidden * 2, kernel_size=1, bias=False),
+#                                    self.bn3,
+#                                    nn.LeakyReLU(negative_slope=0.2))
+#         self.conv4 = nn.Sequential(nn.Conv2d(args.n_hidden * 4, args.n_hidden * 4, kernel_size=1, bias=False),
+#                                    self.bn4,
+#                                    nn.LeakyReLU(negative_slope=0.2))
+#         self.conv5 = nn.Sequential(nn.Conv1d(args.n_hidden * 9, args.emb_dims, kernel_size=1, bias=False),
+#                                    self.bn5,
+#                                    nn.LeakyReLU(negative_slope=0.2))
+
+#         # Fully connected layers to interpret the extracted features and make predictions
+#         self.linear1 = nn.Linear(args.emb_dims*2, 128, bias=False)
+#         self.bn6 = nn.LayerNorm(128)
+#         self.dp1 = nn.Dropout(p=args.dropout)
+
+#         self.linear2 = nn.Linear(128, 64)
+#         self.bn7 = nn.LayerNorm(64)
+#         self.dp2 = nn.Dropout(p=args.dropout)
+
+#         self.linear3 = nn.Linear(64, 32)
+#         self.bn8 = nn.LayerNorm(32)
+#         self.dp3 = nn.Dropout(p=args.dropout)
+
+#         self.linear4 = nn.Linear(32, 16)
+#         self.bn9 = nn.LayerNorm(16)
+#         self.dp4 = nn.Dropout(p=args.dropout)
+
+#         self.linear5 = nn.Linear(16, args.out_dim)  # The final output layer
+
+#     def _make_norm_layer(self, channels):
+#         """创建归一化层,支持batch size=1"""
+#         return nn.GroupNorm(num_groups=min(32, channels), num_channels=channels)
+
+#     def forward(self, x, fx, T=None, geo=None):
+#         """
+#         Forward pass of the model to process input data and predict outputs.
+
+#         Args:
+#             x (torch.Tensor): Input tensor representing a batch of point clouds.
+
+#         Returns:
+#             torch.Tensor: Model predictions for the input batch.
+#         """
+
+#         batch_size = x.size(0)
+
+#         # 兼容 batch_size = 1 输入：去除 batch 维度
+#         x = torch.cat([x, fx], dim=-1)
+#         x = x.permute(0, 2, 1)
+
+#         print(f"x:{x.shape}")
+#         # Extract graph features and apply EdgeConv blocks
+#         x = get_graph_feature(x, k=self.k)  # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k)
+#         print(f"x:{x.shape}")
+#         x = self.conv1(x)  # (batch_size, 3*2, num_points, k) -> (batch_size, 256, num_points, k)
+
+#         # Global max pooling
+#         x1 = x.max(dim=-1, keepdim=False)[0]  # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points)
+
+#         # Repeat the process for subsequent EdgeConv blocks
+#         x = get_graph_feature(x1, k=self.k)   # (batch_size, 256, num_points) -> (batch_size, 256*2, num_points, k)
+#         x = self.conv2(x)                     # (batch_size, 256*2, num_points, k) -> (batch_size, 512, num_points, k)
+#         x2 = x.max(dim=-1, keepdim=False)[0]  # (batch_size, 512, num_points, k) -> (batch_size, 512, num_points)
+
+#         x = get_graph_feature(x2, k=self.k)   # (batch_size, 512, num_points) -> (batch_size, 512*2, num_points, k)
+#         x = self.conv3(x)                     # (batch_size, 512*2, num_points, k) -> (batch_size, 512, num_points, k)
+#         x3 = x.max(dim=-1, keepdim=False)[0]  # (batch_size, 512, num_points, k) -> (batch_size, 512, num_points)
+
+#         x = get_graph_feature(x3, k=self.k)   # (batch_size, 512, num_points) -> (batch_size, 512*2, num_points, k)
+#         x = self.conv4(x)                     # (batch_size, 512*2, num_points, k) -> (batch_size, 1024, num_points, k)
+#         x4 = x.max(dim=-1, keepdim=False)[0]  # (batch_size, 1024, num_points, k) -> (batch_size, 1024, num_points)
+
+#         # Concatenate features from all EdgeConv blocks
+#         x = torch.cat((x1, x2, x3, x4), dim=1)  # (batch_size, 256+512+512+1024, num_points)
+
+#         # Apply the final convolutional block
+#         x = self.conv5(x)  # (batch_size, 256+512+512+1024, num_points) -> (batch_size, emb_dims, num_points)
+#         # Combine global max and average pooling features
+#         # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims)
+#         x1 = F.adaptive_max_pool1d(x, 1).view(batch_size, -1)
+#         # (batch_size, emb_dims, num_points) -> (batch_size, emb_dims)
+#         x2 = F.adaptive_avg_pool1d(x, 1).view(batch_size, -1)
+#         x = torch.cat((x1, x2), 1)   # (batch_size, emb_dims*2)
+#         print(f"x:{x.shape}")
+#         # Process features through fully connected layers with dropout and batch normalization
+#         x = F.leaky_relu(self.bn6(self.linear1(x)), negative_slope=0.2)  # (batch_size, emb_dims*2) -> (batch_size, 128)
+#         x = self.dp1(x)
+#         x = F.leaky_relu(self.bn7(self.linear2(x)), negative_slope=0.2)  # (batch_size, 128) -> (batch_size, 64)
+#         x = self.dp2(x)
+#         x = F.leaky_relu(self.bn8(self.linear3(x)), negative_slope=0.2)  # (batch_size, 64) -> (batch_size, 32)
+#         x = self.dp3(x)
+#         x = F.leaky_relu(self.bn9(self.linear4(x)), negative_slope=0.2)  # (batch_size, 32) -> (batch_size, 16)
+#         x = self.dp4(x)
+
+#         # Final linear layer to produce the output
+#         x = self.linear5(x)                                              # (batch_size, 16) -> (batch_size, 1)
+#         print(f"x:{x.shape}")
+#         x = x.permute(0, 2, 1)
+#         return x
+
+
+class Model(nn.Module):
+    """
+    Deep Graph Convolutional Neural Network for Regression Tasks (RegDGCNN) for processing 3D point cloud data.
+
+    This network architecture extracts hierarchical features from point clouds using graph-based convolutions,
+    enabling effective learning of spatial structures.
+    """
+
+    def __init__(self, args, device):
+        """
+        Initializes the RegDGCNN model with specified configurations.
+
+        Args:
+            args (dict): Configuration parameters including 'k' for the number of neighbors, 'emb_dims' for embedding
+            dimensions, and 'dropout' rate.
+            output_channels (int): Number of output channels (e.g., for drag prediction, this is 1).
+        """
+        super(Model, self).__init__()
+        self.__name__ = "RegDGCNN"
+        self.args = args
+        self.k = 40  # Number of nearest neighbors
+
+        # Batch normalization layers to stabilize and accelerate training
+        self.bn1 = nn.BatchNorm2d(args.n_hidden)
+        self.bn2 = nn.BatchNorm2d(args.n_hidden * 2)
+        self.bn3 = nn.BatchNorm2d(args.n_hidden * 2)
+        self.bn4 = nn.BatchNorm2d(args.n_hidden * 4)
+        self.bn5 = nn.BatchNorm1d(args.emb_dims)
+
+        # EdgeConv layers: Convolutional layers leveraging local neighborhood information
+        self.conv1 = nn.Sequential(
+            nn.Conv2d(
+                2 * (args.fun_dim + args.space_dim),
+                args.n_hidden,
+                kernel_size=1,
+                bias=False,
+            ),
+            self.bn1,
+            nn.LeakyReLU(negative_slope=0.2),
+        )
+        self.conv2 = nn.Sequential(
+            nn.Conv2d(args.n_hidden * 2, args.n_hidden * 2, kernel_size=1, bias=False),
+            self.bn2,
+            nn.LeakyReLU(negative_slope=0.2),
+        )
+        self.conv3 = nn.Sequential(
+            nn.Conv2d(args.n_hidden * 4, args.n_hidden * 2, kernel_size=1, bias=False),
+            self.bn3,
+            nn.LeakyReLU(negative_slope=0.2),
+        )
+        self.conv4 = nn.Sequential(
+            nn.Conv2d(args.n_hidden * 4, args.n_hidden * 4, kernel_size=1, bias=False),
+            self.bn4,
+            nn.LeakyReLU(negative_slope=0.2),
+        )
+        self.conv5 = nn.Sequential(
+            nn.Conv1d(args.n_hidden * 9, args.emb_dims, kernel_size=1, bias=False),
+            self.bn5,
+            nn.LeakyReLU(negative_slope=0.2),
+        )
+
+        # Fully connected layers to interpret the extracted features and make predictions
+        self.point_pred = nn.Sequential(
+            nn.Linear(args.emb_dims, 64), nn.ReLU(), nn.Linear(64, args.out_dim)
+        )
+
+    def _make_norm_layer(self, channels):
+        return nn.GroupNorm(num_groups=min(32, channels), num_channels=channels)
+
+    def forward(self, x, fx, T=None, geo=None):
+        """
+        Forward pass of the model to process input data and predict outputs.
+
+        Args:
+            x (torch.Tensor): Input tensor representing a batch of point clouds.
+
+        Returns:
+            torch.Tensor: Model predictions for the input batch.
+        """
+
+        batch_size = x.size(0)
+
+        # 兼容 batch_size = 1 输入：去除 batch 维度
+        x = torch.cat([x, fx], dim=-1)
+        x = x.permute(0, 2, 1)
+
+        # Extract graph features and apply EdgeConv blocks
+        x = get_graph_feature(
+            x, k=self.k
+        )  # (batch_size, 3, num_points) -> (batch_size, 3*2, num_points, k)
+        x = self.conv1(
+            x
+        )  # (batch_size, 3*2, num_points, k) -> (batch_size, 256, num_points, k)
+        # Global max pooling
+        x1 = x.max(dim=-1, keepdim=False)[
+            0
+        ]  # (batch_size, 64, num_points, k) -> (batch_size, 64, num_points)
+
+        # Repeat the process for subsequent EdgeConv blocks
+        x = get_graph_feature(
+            x1, k=self.k
+        )  # (batch_size, 256, num_points) -> (batch_size, 256*2, num_points, k)
+        x = self.conv2(
+            x
+        )  # (batch_size, 256*2, num_points, k) -> (batch_size, 512, num_points, k)
+        x2 = x.max(dim=-1, keepdim=False)[
+            0
+        ]  # (batch_size, 512, num_points, k) -> (batch_size, 512, num_points)
+
+        x = get_graph_feature(
+            x2, k=self.k
+        )  # (batch_size, 512, num_points) -> (batch_size, 512*2, num_points, k)
+        x = self.conv3(
+            x
+        )  # (batch_size, 512*2, num_points, k) -> (batch_size, 512, num_points, k)
+        x3 = x.max(dim=-1, keepdim=False)[
+            0
+        ]  # (batch_size, 512, num_points, k) -> (batch_size, 512, num_points)
+
+        x = get_graph_feature(
+            x3, k=self.k
+        )  # (batch_size, 512, num_points) -> (batch_size, 512*2, num_points, k)
+        x = self.conv4(
+            x
+        )  # (batch_size, 512*2, num_points, k) -> (batch_size, 1024, num_points, k)
+        x4 = x.max(dim=-1, keepdim=False)[
+            0
+        ]  # (batch_size, 1024, num_points, k) -> (batch_size, 1024, num_points)
+
+        # Concatenate features from all EdgeConv blocks
+        x = torch.cat(
+            (x1, x2, x3, x4), dim=1
+        )  # (batch_size, 256+512+512+1024, num_points)
+
+        # Apply the final convolutional block
+        x = self.conv5(
+            x
+        )  # (batch_size, 256+512+512+1024, num_points) -> (batch_size, emb_dims, num_points)
+        x = x.permute(0, 2, 1)  # (batch_size, num_points, emb_dims)
+        out = self.point_pred(x)  # 点级预测层，输出 (batch_size, num_points, out_dim)
+        return out

model/Swin_Transformer.py

@@ -0,0 +1,211 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint as checkpoint
+from timm.layers import trunc_normal_
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.SwinTransformerBlock import SwinTransformerBlock
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+
+class BasicLayer(nn.Module):
+    """
+    Swin Transformer Layer (Stage).
+    封装了多个 SwinTransformerBlock。
+    """
+    def __init__(
+        self,
+        dim,
+        input_resolution,
+        depth,
+        num_heads,
+        window_size,
+        mlp_ratio=4.0,
+        qkv_bias=True,
+        qk_scale=None,
+        drop=0.0,
+        attn_drop=0.0,
+        drop_path=0.0,
+        norm_layer=nn.LayerNorm,
+        downsample=None,
+        use_checkpoint=False,
+        fused_window_process=False,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.input_resolution = input_resolution
+        self.depth = depth
+        self.use_checkpoint = use_checkpoint
+
+        # build blocks
+        self.blocks = nn.ModuleList([
+            SwinTransformerBlock(
+                dim=dim,
+                input_resolution=input_resolution,
+                num_heads=num_heads,
+                window_size=window_size,
+                shift_size=0 if (i % 2 == 0) else window_size // 2,
+                mlp_ratio=mlp_ratio,
+                qkv_bias=qkv_bias,
+                qk_scale=qk_scale,
+                drop=drop,
+                attn_drop=attn_drop,
+                drop_path=(drop_path[i] if isinstance(drop_path, list) else drop_path),
+                norm_layer=norm_layer,
+                fused_window_process=fused_window_process,
+            )
+            for i in range(depth)
+        ])
+
+        # patch merging layer
+        if downsample is not None:
+            self.downsample = downsample(
+                input_resolution, dim=dim, norm_layer=norm_layer
+            )
+        else:
+            self.downsample = None
+
+    def forward(self, x):
+        for blk in self.blocks:
+            if self.use_checkpoint:
+                x = checkpoint.checkpoint(blk, x)
+            else:
+                x = blk(x)
+        if self.downsample is not None:
+            x = self.downsample(x)
+        return x
+
+
+class Model(nn.Module):
+    """
+    Swin Transformer 主模型。
+    """
+    def __init__(self, args, device, window_size=4):
+        super(Model, self).__init__()
+        self.__name__ = "SwinTransformer"
+        self.args = args
+        
+        if args.geotype != "structured_2D":
+            raise ValueError("Swin Transformer only supports Structured 2D geometry")
+
+        # 1. Embedding & Preprocessing
+        if args.unified_pos: 
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            in_dim = args.fun_dim + args.ref ** len(args.shapelist)
+        else:
+            in_dim = args.fun_dim + args.space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=in_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+        
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+            
+        self.placeholder = nn.Parameter(
+            (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
+        )
+        
+        self.padding = [
+            (window_size - size % window_size) % window_size for size in args.shapelist
+        ]
+        self.augmented_resolution = [
+            (self.padding[i] + args.shapelist[i]) for i in range(len(self.padding))
+        ]
+        
+        # 3. Swin Layers
+        self.blocks = nn.ModuleList([
+            BasicLayer(
+                dim=args.n_hidden,
+                input_resolution=self.augmented_resolution,
+                depth=2,
+                num_heads=args.n_heads,
+                window_size=window_size,
+            )
+            for _ in range(args.n_layers)
+        ])
+        
+        # 4. Projectors
+        self.fc1 = nn.Linear(args.n_hidden, args.n_hidden * 2)
+        self.fc2 = nn.Linear(args.n_hidden * 2, args.out_dim)
+        
+        self.initialize_weights()
+
+    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 structured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 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, :]
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+            
+        ## aug shape
+        fx = fx.permute(0, 2, 1).reshape(B, self.args.n_hidden, *self.args.shapelist)
+        
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 2:
+                fx = F.pad(fx, [0, self.padding[1], 0, self.padding[0]])
+            elif len(self.args.shapelist) == 3:
+                fx = F.pad(fx, [0, self.padding[2], 0, self.padding[1], 0, self.padding[0]])
+                
+        fx = fx.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
+        
+        ## swin transformer
+        for block in self.blocks:
+            fx = block(fx)
+            
+        ## back to original shape
+        fx = fx.permute(0, 2, 1).reshape(
+            B, self.args.n_hidden, *self.augmented_resolution
+        )
+        
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 2:
+                fx = fx[..., : -self.padding[0], : -self.padding[1]]
+            elif len(self.args.shapelist) == 3:
+                fx = fx[..., : -self.padding[0], : -self.padding[1], : -self.padding[2]]
+                
+        fx = fx.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
+        
+        ## projection
+        fx = self.fc1(fx)
+        fx = F.gelu(fx)
+        fx = self.fc2(fx)
+        return fx
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "structured_2D":
+            return self.structured_geo(x, fx, T)
+        else:
+            raise ValueError("Swin Transformer only supports Structured 2D geometry")

model/Transformer.py

@@ -0,0 +1,149 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from timm.layers import trunc_normal_
+from onescience.modules.attention.flashattention import FlashAttention
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+from einops import rearrange, repeat
+
+class Transformer_block(nn.Module):
+    """
+    Transformer encoder block.
+    包含多头自注意力和前馈神经网络，最后一层可选是否直接输出预测维度。
+    """
+
+    def __init__(
+        self,
+        num_heads: int,
+        hidden_dim: int,
+        dropout: float,
+        act="gelu",
+        mlp_ratio=4,
+        last_layer=False,
+        out_dim=1,
+    ):
+        super().__init__()
+        self.last_layer = last_layer
+        self.ln_1 = nn.LayerNorm(hidden_dim)
+
+        self.Attn = FlashAttention(
+            dim=hidden_dim,
+            heads=num_heads,
+            dim_head=hidden_dim // num_heads,
+            dropout=dropout,
+        )
+        
+        self.ln_2 = nn.LayerNorm(hidden_dim)
+        
+        self.mlp = StandardMLP(
+            input_dim=hidden_dim,
+            output_dim=hidden_dim,
+            hidden_dims=[hidden_dim * mlp_ratio],
+            activation=act,
+            use_bias=True
+        )
+        
+        if self.last_layer:
+            self.ln_3 = nn.LayerNorm(hidden_dim)
+            self.mlp2 = nn.Linear(hidden_dim, out_dim)
+
+    def forward(self, fx):
+        fx = self.Attn(self.ln_1(fx)) + fx
+        fx = self.mlp(self.ln_2(fx)) + fx
+        if self.last_layer:
+            return self.mlp2(self.ln_3(fx))
+        else:
+            return fx
+
+
+class Model(nn.Module):
+    """
+    标准的 Transformer 模型架构。
+    通过堆叠 Transformer_block 处理物理场数据。
+    """
+    def __init__(self, args, device):
+        super(Model, self).__init__()
+        self.__name__ = "Transformer"
+        self.args = args
+        
+        ## embedding
+        if (
+            args.unified_pos and args.geotype != "unstructured"
+        ):  # only for structured mesh
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            input_dim = args.fun_dim + args.ref ** len(args.shapelist)
+        else:
+            input_dim = args.fun_dim + args.space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        self.blocks = nn.ModuleList(
+            [
+                Transformer_block(
+                    num_heads=args.n_heads,
+                    hidden_dim=args.n_hidden,
+                    dropout=args.dropout,
+                    act=args.act,
+                    mlp_ratio=args.mlp_ratio,
+                    out_dim=args.out_dim,
+                    last_layer=(_ == args.n_layers - 1),
+                )
+                for _ in range(args.n_layers)
+            ]
+        )
+        
+        self.placeholder = nn.Parameter(
+            (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
+        )
+        self.initialize_weights()
+
+    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 forward(self, x, fx, T=None, geo=None):
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 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, :]
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+
+        for block in self.blocks:
+            fx = block(fx)
+            
+        return fx

model/Transolver.py

@@ -0,0 +1,130 @@
+import torch
+import torch.nn as nn
+import numpy as np
+from timm.layers import trunc_normal_
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.transformer.Transolver_block import Transolver_block
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+
+class Model(nn.Module):
+    """
+    Transolver 模型。
+    通过物理启发的切片机制 (Slicing) 解决 PDE 和物理场预测问题。
+    """
+    def __init__(self, args, device):
+        super(Model, self).__init__()
+        self.__name__ = "Transolver"
+        self.args = args
+        
+        ## embedding
+        if (
+            args.unified_pos and args.geotype != "unstructured"
+        ):  
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            self.preprocess = StandardMLP(
+                input_dim=args.fun_dim + args.ref ** len(args.shapelist),
+                output_dim=args.n_hidden,
+                hidden_dims=[args.n_hidden * 2],
+                activation=args.act,
+                use_bias=True
+            )
+        else:
+            self.preprocess = StandardMLP(
+                input_dim=args.fun_dim + args.space_dim,
+                output_dim=args.n_hidden,
+                hidden_dims=[args.n_hidden * 2],
+                activation=args.act,
+                use_bias=True
+            )
+            
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        ## models
+        self.blocks = nn.ModuleList(
+            [
+                Transolver_block(
+                    num_heads=args.n_heads,
+                    hidden_dim=args.n_hidden,
+                    dropout=args.dropout,
+                    act=args.act,
+                    mlp_ratio=args.mlp_ratio,
+                    out_dim=args.out_dim,
+                    slice_num=args.slice_num,
+                    last_layer=(_ == args.n_layers - 1),
+                    geotype=args.geotype,
+                    shapelist=args.shapelist,
+                )
+                for _ in range(args.n_layers)
+            ]
+        )
+        
+        self.placeholder = nn.Parameter(
+            (1 / (args.n_hidden)) * torch.rand(args.n_hidden, dtype=torch.float)
+        )
+        self.initialize_weights()
+
+    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 structured_geo(self, x, fx, T=None):
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 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, :]
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+
+        for block in self.blocks:
+            fx = block(fx)
+        return fx
+
+    def unstructured_geo(self, x, fx, T=None):
+        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, :]
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+
+        for block in self.blocks:
+            fx = block(fx)
+        return fx
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/U_FNO.py

@@ -0,0 +1,241 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from onescience.modules.fourier.fno_layers import (
+    SpectralConv1d,
+    SpectralConv2d,
+    SpectralConv3d,
+)
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI
+
+SpectralConvList = [None, SpectralConv1d, SpectralConv2d, SpectralConv3d]
+
+from .U_Net import Model as U_Net
+
+class Model(nn.Module):
+    """
+    U-FNO 模型。
+    
+    结合了 U-Net (用于多尺度特征提取) 和 FNO (用于全局谱特征提取)。
+    U-Net 作为 FNO 层的并联分支，增强了局部特征捕捉能力。
+
+    """
+    def __init__(self, args, device, s1=96, s2=96):
+        super(Model, self).__init__()
+        self.__name__ = "U-FNO"
+        self.args = args
+        self.device = device
+
+        # 1. Embedding & Preprocessing
+        # -----------------------------------------------------------
+        input_dim = args.fun_dim
+        if args.unified_pos and args.geotype != "unstructured":
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            input_dim += args.ref ** len(args.shapelist)
+        else:
+            input_dim += args.space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # 2. Geometry Projection & Padding Logic
+        # -----------------------------------------------------------
+        if self.args.geotype == "unstructured":
+            self.fftproject_in = GeoSpectralConv2d(
+                in_channels=args.n_hidden,
+                out_channels=args.n_hidden,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            self.fftproject_out = GeoSpectralConv2d(
+                in_channels=args.n_hidden,
+                out_channels=args.n_hidden,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            self.iphi = IPHI()
+            
+            self.padding = [(16 - size % 16) % 16 for size in [s1, s2]]
+        else:
+            self.padding = [(16 - size % 16) % 16 for size in args.shapelist]
+
+        # 3. FNO Blocks
+        dim = len(self.padding)
+        
+        # 辅助函数：构建 FNO 参数 
+        def get_fno_layer(in_c, out_c):
+            kwargs = {
+                "in_channels": in_c,
+                "out_channels": out_c
+            }
+            # 动态添加 modes1, modes2, modes3
+            mode_names = ["modes1", "modes2", "modes3"]
+            for i in range(dim):
+                if i < len(mode_names):
+                    kwargs[mode_names[i]] = args.modes
+            return SpectralConvList[dim](**kwargs)
+
+        self.conv0 = get_fno_layer(args.n_hidden, args.n_hidden)
+        self.conv1 = get_fno_layer(args.n_hidden, args.n_hidden)
+        self.conv2 = get_fno_layer(args.n_hidden, args.n_hidden)
+        self.conv3 = get_fno_layer(args.n_hidden, args.n_hidden)
+
+        ConvClass = [None, nn.Conv1d, nn.Conv2d, nn.Conv3d][dim]
+        self.w0 = ConvClass(args.n_hidden, args.n_hidden, 1)
+        self.w1 = ConvClass(args.n_hidden, args.n_hidden, 1)
+        self.w2 = ConvClass(args.n_hidden, args.n_hidden, 1)
+        self.w3 = ConvClass(args.n_hidden, args.n_hidden, 1)
+
+        # 4. U-Net Branches (Parallel)
+        # -----------------------------------------------------------
+        self.u_net2 = U_Net(args, device)
+        self.u_net3 = U_Net(args, device)
+
+        # 5. Projectors
+        self.fc1 = nn.Linear(args.n_hidden, args.n_hidden)
+        self.fc2 = nn.Linear(args.n_hidden, args.out_dim)
+
+    def structured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+        
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 1)
+            
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden) # (B, C)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1) # (B, 1, C)
+            fx = fx + Time_emb
+
+        x = fx.permute(0, 2, 1).reshape(B, self.args.n_hidden, *self.args.shapelist)
+        
+        if not all(item == 0 for item in self.padding):
+            pad_arg = []
+            for p in reversed(self.padding):
+                pad_arg.extend([0, p])
+            x = F.pad(x, pad_arg)
+
+        # Layer 0
+        x1 = self.conv0(x)
+        x2 = self.w0(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        # Layer 1
+        x1 = self.conv1(x)
+        x2 = self.w1(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        # Layer 2 (with U-Net)
+        x1 = self.conv2(x)
+        x2 = self.w2(x)
+        x3 = self.u_net2.multiscale(x)
+        x = x1 + x2 + x3
+        x = F.gelu(x)
+
+        # Layer 3 (with U-Net)
+        x1 = self.conv3(x)
+        x2 = self.w3(x)
+        x3 = self.u_net3.multiscale(x)
+        x = x1 + x2 + x3
+
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 1:
+                x = x[..., :-self.padding[0]]
+            elif len(self.args.shapelist) == 2:
+                x = x[..., :-self.padding[0], :-self.padding[1]]
+            elif len(self.args.shapelist) == 3:
+                x = x[..., :-self.padding[0], :-self.padding[1], :-self.padding[2]]
+        
+        x = x.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def unstructured_geo(self, x, fx, T=None):
+        original_pos = x
+        
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+
+        x = self.fftproject_in(
+            fx.permute(0, 2, 1), x_in=original_pos, iphi=self.iphi, code=None
+        )
+
+        # Layer 0
+        x1 = self.conv0(x)
+        x2 = self.w0(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        # Layer 1
+        x1 = self.conv1(x)
+        x2 = self.w1(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        # Layer 2
+        x1 = self.conv2(x)
+        x2 = self.w2(x)
+        x3 = self.u_net2.multiscale(x)
+        x = x1 + x2 + x3
+        x = F.gelu(x)
+
+        # Layer 3
+        x1 = self.conv3(x)
+        x2 = self.w3(x)
+        x3 = self.u_net3.multiscale(x)
+        x = x1 + x2 + x3
+
+        x = self.fftproject_out(
+            x, x_out=original_pos, iphi=self.iphi, code=None
+        ).permute(0, 2, 1)
+        
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/U_NO.py

@@ -0,0 +1,321 @@
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+
+from onescience.modules.fourier.fno_layers import (
+    SpectralConv1d,
+    SpectralConv2d,
+    SpectralConv3d,
+)
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI
+
+# --- 引入 U-Net 基础组件 (替代 UNet_Blocks) ---
+from onescience.modules.layer.unet_layer import (
+    DoubleConv1D, Down1D, Up1D, OutConv1D,
+    DoubleConv2D, Down2D, Up2D, OutConv2D,
+    DoubleConv3D, Down3D, Up3D, OutConv3D,
+)
+
+ConvList = [None, DoubleConv1D, DoubleConv2D, DoubleConv3D]
+DownList = [None, Down1D, Down2D, Down3D]
+UpList = [None, Up1D, Up2D, Up3D]
+OutList = [None, OutConv1D, OutConv2D, OutConv3D]
+SpectralConvList = [None, SpectralConv1d, SpectralConv2d, SpectralConv3d]
+
+class Model(nn.Module):
+    """
+    U-NO (U-Net Neural Operator) 模型。
+    
+    结合了 U-Net 的多尺度结构和 FNO 的谱卷积能力。
+    在 U-Net 的每个 Encoder 和 Decoder 层级之间插入了 FNO Block (SpectralConv) 和 1x1 卷积残差。
+    """
+    def __init__(self, args, device, bilinear=True, s1=96, s2=96):
+        super(Model, self).__init__()
+        self.__name__ = "U_NO"
+        self.args = args
+
+        if args.task == "steady":
+            normtype = "bn"
+        else:
+            normtype = "in"
+
+        # 1. Embedding & Preprocessing
+        # -----------------------------------------------------------
+        input_dim = args.fun_dim
+        if args.unified_pos and args.geotype != "unstructured":
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            input_dim += args.ref ** len(args.shapelist)
+        else:
+            input_dim += args.space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            output_dim=args.n_hidden,
+            hidden_dims=[args.n_hidden * 2],
+            activation=args.act,
+            use_bias=True
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # 2. Geometry Projection (GeoFNO for unstructured)
+        # -----------------------------------------------------------
+        if self.args.geotype == "unstructured":
+            self.fftproject_in = GeoSpectralConv2d(
+                in_channels=args.n_hidden,
+                out_channels=args.n_hidden,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            self.fftproject_out = GeoSpectralConv2d(
+                in_channels=args.n_hidden,
+                out_channels=args.n_hidden,
+                modes1=args.modes,
+                modes2=args.modes,
+                s1=s1,
+                s2=s2
+            )
+            self.iphi = IPHI()
+            
+            patch_size = [(size + (16 - size % 16) % 16) // 16 for size in [s1, s2]]
+            self.padding = [(16 - size % 16) % 16 for size in [s1, s2]]
+            self.augmented_resolution = [s1, s2]
+        else:
+            patch_size = [
+                (size + (16 - size % 16) % 16) // 16 for size in args.shapelist
+            ]
+            self.padding = [(16 - size % 16) % 16 for size in args.shapelist]
+            self.augmented_resolution = [
+                shape + padding for shape, padding in zip(args.shapelist, self.padding)
+            ]
+
+        dim = len(patch_size) 
+
+        # 3. Multiscale U-Net Modules 
+        self.inc = ConvList[dim](args.n_hidden, args.n_hidden, normtype=normtype)
+        
+        self.down1 = DownList[dim](args.n_hidden, args.n_hidden * 2, normtype=normtype)
+        self.down2 = DownList[dim](args.n_hidden * 2, args.n_hidden * 4, normtype=normtype)
+        self.down3 = DownList[dim](args.n_hidden * 4, args.n_hidden * 8, normtype=normtype)
+        
+        factor = 2 if bilinear else 1
+        self.down4 = DownList[dim](args.n_hidden * 8, args.n_hidden * 16 // factor, normtype=normtype)
+        
+        self.up1 = UpList[dim](args.n_hidden * 16, args.n_hidden * 8 // factor, bilinear, normtype=normtype)
+        self.up2 = UpList[dim](args.n_hidden * 8, args.n_hidden * 4 // factor, bilinear, normtype=normtype)
+        self.up3 = UpList[dim](args.n_hidden * 4, args.n_hidden * 2 // factor, bilinear, normtype=normtype)
+        self.up4 = UpList[dim](args.n_hidden * 2, args.n_hidden, bilinear, normtype=normtype)
+        
+        self.outc = OutList[dim](args.n_hidden, args.n_hidden)
+
+        # 4. FNO Blocks
+        def get_fno_layer(in_c, out_c, res_list, divisor):
+            modes_list = [
+                max(1, min(args.modes, res // divisor))
+                for res in res_list
+            ]
+
+            kwargs = {
+                "in_channels": in_c,
+                "out_channels": out_c
+            }
+            mode_names = ["modes1", "modes2", "modes3"]
+            for i, m in enumerate(modes_list):
+                if i < len(mode_names):
+                    kwargs[mode_names[i]] = m
+            return SpectralConvList[dim](**kwargs)
+
+        # Down Path FNOs
+        self.process1_down = get_fno_layer(args.n_hidden, args.n_hidden, self.augmented_resolution, 2)
+        self.process2_down = get_fno_layer(args.n_hidden * 2, args.n_hidden * 2, self.augmented_resolution, 4)
+        self.process3_down = get_fno_layer(args.n_hidden * 4, args.n_hidden * 4, self.augmented_resolution, 8)
+        self.process4_down = get_fno_layer(args.n_hidden * 8, args.n_hidden * 8, self.augmented_resolution, 16)
+        self.process5_down = get_fno_layer(args.n_hidden * 16 // factor, args.n_hidden * 16 // factor, self.augmented_resolution, 32)
+
+        # Residual Weights (1x1 Conv)
+        self.w1_down = ConvList[dim](args.n_hidden, args.n_hidden, 1) # kernel_size=1
+        self.w2_down = ConvList[dim](args.n_hidden * 2, args.n_hidden * 2, 1)
+        self.w3_down = ConvList[dim](args.n_hidden * 4, args.n_hidden * 4, 1)
+        self.w4_down = ConvList[dim](args.n_hidden * 8, args.n_hidden * 8, 1)
+        self.w5_down = ConvList[dim](args.n_hidden * 16 // factor, args.n_hidden * 16 // factor, 1)
+
+        # Up Path FNOs
+        self.process1_up = get_fno_layer(args.n_hidden, args.n_hidden, self.augmented_resolution, 2)
+        self.process2_up = get_fno_layer(args.n_hidden * 2 // factor, args.n_hidden * 2 // factor, self.augmented_resolution, 4)
+        self.process3_up = get_fno_layer(args.n_hidden * 4 // factor, args.n_hidden * 4 // factor, self.augmented_resolution, 8)
+        self.process4_up = get_fno_layer(args.n_hidden * 8 // factor, args.n_hidden * 8 // factor, self.augmented_resolution, 16)
+        self.process5_up = get_fno_layer(args.n_hidden * 16 // factor, args.n_hidden * 16 // factor, self.augmented_resolution, 32)
+
+        self.w1_up = ConvList[dim](args.n_hidden, args.n_hidden, 1)
+        self.w2_up = ConvList[dim](args.n_hidden * 2 // factor, args.n_hidden * 2 // factor, 1)
+        self.w3_up = ConvList[dim](args.n_hidden * 4 // factor, args.n_hidden * 4 // factor, 1)
+        self.w4_up = ConvList[dim](args.n_hidden * 8 // factor, args.n_hidden * 8 // factor, 1)
+        self.w5_up = ConvList[dim](args.n_hidden * 16 // factor, args.n_hidden * 16 // factor, 1)
+
+        # 5. Projectors
+        self.fc1 = nn.Linear(args.n_hidden, args.n_hidden * 2)
+        self.fc2 = nn.Linear(args.n_hidden * 2, args.out_dim)
+
+    def structured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 1)
+        
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb # Broadcast
+
+        # Reshape to Grid (B, C, H, W...)
+        x = fx.permute(0, 2, 1).reshape(B, self.args.n_hidden, *self.args.shapelist)
+        
+        # Padding
+        if not all(item == 0 for item in self.padding):
+            pad_arg = []
+            for p in reversed(self.padding):
+                pad_arg.extend([0, p])
+            x = F.pad(x, pad_arg)
+
+        # === U-NO Body ===
+        # Level 1 Down
+        x1 = self.inc(x)
+        # FNO Processing + Residual
+        x1 = F.gelu(self.process1_down(x1) + self.w1_down(x1))
+        
+        # Level 2 Down
+        x2 = self.down1(x1)
+        x2 = F.gelu(self.process2_down(x2) + self.w2_down(x2))
+        
+        # Level 3 Down
+        x3 = self.down2(x2)
+        x3 = F.gelu(self.process3_down(x3) + self.w3_down(x3))
+        
+        # Level 4 Down
+        x4 = self.down3(x3)
+        x4 = F.gelu(self.process4_down(x4) + self.w4_down(x4))
+        
+        # Level 5 (Bottleneck)
+        x5 = self.down4(x4)
+        x5 = F.gelu(self.process5_down(x5) + self.w5_down(x5))
+        # Bottleneck Up Process
+        x5 = F.gelu(self.process5_up(x5) + self.w5_up(x5))
+        
+        # Level 4 Up
+        x = self.up1(x5, x4)
+        x = F.gelu(self.process4_up(x) + self.w4_up(x))
+        
+        # Level 3 Up
+        x = self.up2(x, x3)
+        x = F.gelu(self.process3_up(x) + self.w3_up(x))
+        
+        # Level 2 Up
+        x = self.up3(x, x2)
+        x = F.gelu(self.process2_up(x) + self.w2_up(x))
+        
+        # Level 1 Up
+        x = self.up4(x, x1)
+        x = F.gelu(self.process1_up(x) + self.w1_up(x))
+        
+        x = self.outc(x)
+
+        # Un-padding
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 1:
+                x = x[..., : -self.padding[0]]
+            elif len(self.args.shapelist) == 2:
+                x = x[..., : -self.padding[0], : -self.padding[1]]
+            elif len(self.args.shapelist) == 3:
+                x = x[..., : -self.padding[0], : -self.padding[1], : -self.padding[2]]
+        
+        # Output Projection
+        x = x.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def unstructured_geo(self, x, fx, T=None):
+        original_pos = x
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden)
+            Time_emb = self.time_fc(Time_emb)
+            if Time_emb.ndim == 2:
+                Time_emb = Time_emb.unsqueeze(1)
+            fx = fx + Time_emb
+
+        # Projection to Grid
+        x = self.fftproject_in(
+            fx.permute(0, 2, 1), x_in=original_pos, iphi=self.iphi, code=None
+        )
+        
+        # === U-NO Body===
+        x1 = self.inc(x)
+        x1 = F.gelu(self.process1_down(x1) + self.w1_down(x1))
+        
+        x2 = self.down1(x1)
+        x2 = F.gelu(self.process2_down(x2) + self.w2_down(x2))
+        
+        x3 = self.down2(x2)
+        x3 = F.gelu(self.process3_down(x3) + self.w3_down(x3))
+        
+        x4 = self.down3(x3)
+        x4 = F.gelu(self.process4_down(x4) + self.w4_down(x4))
+        
+        x5 = self.down4(x4)
+        x5 = F.gelu(self.process5_down(x5) + self.w5_down(x5))
+        x5 = F.gelu(self.process5_up(x5) + self.w5_up(x5))
+        
+        x = self.up1(x5, x4)
+        x = F.gelu(self.process4_up(x) + self.w4_up(x))
+        
+        x = self.up2(x, x3)
+        x = F.gelu(self.process3_up(x) + self.w3_up(x))
+        
+        x = self.up3(x, x2)
+        x = F.gelu(self.process2_up(x) + self.w2_up(x))
+        
+        x = self.up4(x, x1)
+        x = F.gelu(self.process1_up(x) + self.w1_up(x))
+        
+        x = self.outc(x)
+        
+        # Projection Back to Points
+        x = self.fftproject_out(
+            x, x_out=original_pos, iphi=self.iphi, code=None
+        ).permute(0, 2, 1)
+        
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/U_Net.py

@@ -0,0 +1,219 @@
+import torch
+import math
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+from onescience.modules.decoder.unet_decoder import (
+    UNetDecoder1D,
+    UNetDecoder2D,
+    UNetDecoder3D,
+)
+from onescience.modules.encoder.unet_encoder import (
+    UNetEncoder1D,
+    UNetEncoder2D,
+    UNetEncoder3D,
+)
+from onescience.modules.fourier.geo_spectral import GeoSpectralConv2d, IPHI
+from onescience.modules.head.unet_head import UNetHead1D, UNetHead2D, UNetHead3D
+from onescience.modules.mlp.MLP import StandardMLP
+from onescience.modules.embedding import timestep_embedding, unified_pos_embedding
+
+EncoderList = [None, UNetEncoder1D, UNetEncoder2D, UNetEncoder3D]
+DecoderList = [None, UNetDecoder1D, UNetDecoder2D, UNetDecoder3D]
+HeadList = [None, UNetHead1D, UNetHead2D, UNetHead3D]
+
+class Model(nn.Module):
+    """
+    多尺度物理场 U-Net 模型。
+
+    该模型支持结构化网格（1D/2D/3D）和非结构化网格（通过 GeoFNO 的几何投影）的物理场预测。
+    利用编码器和解码器实现多尺度特征提取与融合。
+
+    Args:
+        args: 包含模型配置的参数命名空间 (如 task, geotype, n_hidden 等)。
+        device: 运行设备。
+        bilinear (bool, optional): U-Net 上采样是否使用双线性插值。默认值: True。
+        s1 (int, optional): 非结构化网格投影的潜在空间高度。默认值: 96。
+        s2 (int, optional): 非结构化网格投影的潜在空间宽度。默认值: 96。
+
+    形状:
+        输入 x: 坐标张量 (B, N, space_dim)。
+        输入 fx: 物理场特征张量 (B, N, fun_dim)。
+        输入 T: 可选的时间步张量。
+        输出: (B, N, out_dim)
+    """
+    def __init__(self, args, device, bilinear=True, s1=96, s2=96):
+        super(Model, self).__init__()
+        self.__name__ = "U-Net"
+        self.args = args
+        self.bilinear = bilinear
+        
+        normtype = "bn" if args.task == "steady" else "in"
+
+        # ==========================================
+        # 1. 位置编码与预处理 (Embedding)
+        # ==========================================
+        if args.unified_pos and args.geotype != "unstructured":  # structured mesh
+            self.pos = unified_pos_embedding(args.shapelist, args.ref, device=device)
+            input_dim = args.fun_dim + args.ref ** len(args.shapelist)
+        else:
+            input_dim = args.fun_dim + args.space_dim
+
+        self.preprocess = StandardMLP(
+            input_dim=input_dim,
+            hidden_dims=[args.n_hidden * 2],
+            output_dim=args.n_hidden,
+            activation=args.act,
+        )
+
+        if args.time_input:
+            self.time_fc = nn.Sequential(
+                nn.Linear(args.n_hidden, args.n_hidden),
+                nn.SiLU(),
+                nn.Linear(args.n_hidden, args.n_hidden),
+            )
+
+        # ==========================================
+        # 2. 几何投影 (Geometry Projection)
+        # ==========================================
+        if self.args.geotype == "unstructured":
+            self.fftproject_in = GeoSpectralConv2d(
+                in_channels=args.n_hidden, 
+                out_channels=args.n_hidden, 
+                modes1=args.modes, 
+                modes2=args.modes, 
+                s1=s1, 
+                s2=s2
+            )
+            self.fftproject_out = GeoSpectralConv2d(
+                in_channels=args.n_hidden, 
+                out_channels=args.n_hidden, 
+                modes1=args.modes, 
+                modes2=args.modes, 
+                s1=s1, 
+                s2=s2
+            )
+            self.iphi = IPHI()
+            
+            patch_size = [(size + (16 - size % 16) % 16) // 16 for size in [s1, s2]]
+            self.padding = [(16 - size % 16) % 16 for size in [s1, s2]]
+        else:
+            patch_size = [(size + (16 - size % 16) % 16) // 16 for size in args.shapelist]
+            self.padding = [(16 - size % 16) % 16 for size in args.shapelist]
+
+        # ==========================================
+        # 3. U-Net 核心 (Multiscale modules)
+        # ==========================================
+        dim = len(patch_size)
+        num_stages = 4 
+        
+        self.encoder = EncoderList[dim](
+            in_channels=args.n_hidden,
+            base_channels=args.n_hidden,
+            num_stages=num_stages,
+            bilinear=bilinear,
+            normtype=normtype
+        )
+        
+        self.decoder = DecoderList[dim](
+            base_channels=args.n_hidden,
+            num_stages=num_stages,
+            bilinear=bilinear,
+            normtype=normtype
+        )
+        
+        self.outc = HeadList[dim](
+            in_channels=args.n_hidden,
+            out_channels=args.n_hidden
+        )
+
+        # 最终投影
+        self.fc1 = nn.Linear(args.n_hidden, args.n_hidden)
+        self.fc2 = nn.Linear(args.n_hidden, args.out_dim)
+
+    def multiscale(self, x):
+        """完全解耦的 U-Net ���向计算"""
+        # 提取多尺度特征
+        features = self.encoder(x)
+        # 融合上采样
+        x = self.decoder(features)
+        # 预测头输出
+        return self.outc(x)
+
+    def structured_geo(self, x, fx, T=None):
+        B, N, _ = x.shape
+        if self.args.unified_pos:
+            x = self.pos.repeat(x.shape[0], 1, 1)
+            
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden).repeat(1, x.shape[1], 1)
+            Time_emb = self.time_fc(Time_emb)
+            fx = fx + Time_emb
+            
+        x = fx.permute(0, 2, 1).reshape(B, self.args.n_hidden, *self.args.shapelist)
+        
+        # Padding
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 2:
+                x = F.pad(x, [0, self.padding[1], 0, self.padding[0]])
+            elif len(self.args.shapelist) == 3:
+                x = F.pad(x, [0, self.padding[2], 0, self.padding[1], 0, self.padding[0]])
+                
+        # 极简调用多尺度 U-Net
+        x = self.multiscale(x)
+        
+        # Unpadding
+        if not all(item == 0 for item in self.padding):
+            if len(self.args.shapelist) == 2:
+                x = x[..., : -self.padding[0], : -self.padding[1]]
+            elif len(self.args.shapelist) == 3:
+                x = x[..., : -self.padding[0], : -self.padding[1], : -self.padding[2]]
+                
+        x = x.reshape(B, self.args.n_hidden, -1).permute(0, 2, 1)
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def unstructured_geo(self, x, fx, T=None):
+        original_pos = x
+        if fx is not None:
+            fx = torch.cat((x, fx), -1)
+            fx = self.preprocess(fx)
+        else:
+            fx = self.preprocess(x)
+
+        if T is not None:
+            Time_emb = timestep_embedding(T, self.args.n_hidden).repeat(1, x.shape[1], 1)
+            Time_emb = self.time_fc(Time_emb)
+            fx = fx + Time_emb
+
+        # GeoFNO: 坐标映射与入场变换
+        x = self.fftproject_in(
+            fx.permute(0, 2, 1), x_in=original_pos, iphi=self.iphi, code=None
+        )
+        
+        # 极简调用多尺度 U-Net
+        x = self.multiscale(x)
+        
+        # GeoFNO: 出场逆变换
+        x = self.fftproject_out(
+            x, x_out=original_pos, iphi=self.iphi, code=None
+        ).permute(0, 2, 1)
+        
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x
+
+    def forward(self, x, fx, T=None, geo=None):
+        if self.args.geotype == "unstructured":
+            return self.unstructured_geo(x, fx, T)
+        else:
+            return self.structured_geo(x, fx, T)

model/model_factory.py

@@ -0,0 +1,48 @@
+from . import Transolver
+from . import LSM
+from . import FNO
+from . import U_Net
+from . import Transformer
+from . import Factformer
+from . import Swin_Transformer
+from . import Galerkin_Transformer
+from . import GNOT
+from . import U_NO
+from . import U_FNO
+from . import F_FNO
+from . import ONO
+from . import MWT
+from . import GraphSAGE
+from . import Graph_UNet
+from . import PointNet
+from . import DeepONet
+from . import MeshGraphNet
+from . import RegDGCNN
+from . import GFNO
+
+
+def get_model(args, device):
+    model_dict = {
+        "PointNet": PointNet,
+        "Graph_UNet": Graph_UNet,
+        "GraphSAGE": GraphSAGE,
+        "MWT": MWT,
+        "ONO": ONO,
+        "F_FNO": F_FNO,
+        "U_FNO": U_FNO,
+        "U_NO": U_NO,
+        "GNOT": GNOT,
+        "Galerkin_Transformer": Galerkin_Transformer,
+        "Swin_Transformer": Swin_Transformer,
+        "Factformer": Factformer,
+        "Transformer": Transformer,
+        "U_Net": U_Net,
+        "FNO": FNO,
+        "Transolver": Transolver,
+        "LSM": LSM,
+        "DeepONet": DeepONet,
+        "MeshGraphNet": MeshGraphNet,
+        "RegDGCNN": RegDGCNN,
+        "GFNO": GFNO,
+    }
+    return model_dict[args.model].Model(args, device)

scripts/fake_data.py

@@ -0,0 +1,50 @@
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT / "model"))
+
+
+def load_config():
+    with open(PROJECT_ROOT / "config" / "config.yaml", "r", encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+
+def main():
+    os.chdir(PROJECT_ROOT)
+    cfg = load_config()
+    data_dir = PROJECT_ROOT / cfg["paths"]["data_dir"]
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    n_total = cfg["data"]["ntrain"] + cfg["data"]["ntest"] + 1
+    height, width = 221, 51
+    rng = np.random.default_rng(42)
+
+    y_axis = np.linspace(-1.0, 1.0, height, dtype=np.float32)
+    x_axis = np.linspace(0.0, 1.0, width, dtype=np.float32)
+    grid_y, grid_x = np.meshgrid(y_axis, x_axis, indexing="ij")
+
+    input_x = np.broadcast_to(grid_x, (n_total, height, width)).copy()
+    input_y = np.broadcast_to(grid_y, (n_total, height, width)).copy()
+    q = np.zeros((n_total, 5, height, width), dtype=np.float32)
+
+    for i in range(n_total):
+        phase = 0.15 * i
+        field = np.sin(np.pi * (grid_x + phase)) * np.cos(np.pi * grid_y)
+        q[i, 4] = field + 0.01 * rng.standard_normal((height, width))
+
+    np.save(data_dir / "NACA_Cylinder_X.npy", input_x.astype(np.float32))
+    np.save(data_dir / "NACA_Cylinder_Y.npy", input_y.astype(np.float32))
+    np.save(data_dir / "NACA_Cylinder_Q.npy", q.astype(np.float32))
+
+    print(f"Fake airfoil data written to {data_dir}")
+    print(f"Samples={n_total}, X/Y shape={input_x.shape}, Q shape={q.shape}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,99 @@
+import json
+import os
+from pathlib import Path
+
+import train as train_runtime
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_model(args, device):
+    model = train_runtime.get_model(args, device).to(device)
+    checkpoint_path = PROJECT_ROOT / "checkpoints" / f"{args.save_name}.pt"
+    state = train_runtime.torch.load(
+        checkpoint_path, map_location=device, weights_only=False
+    )
+    model_state = state["model_state"] if isinstance(state, dict) and "model_state" in state else state
+    model.load_state_dict(model_state)
+    model.eval()
+    return model
+
+
+def evaluate(model, test_loader, dataset, args, device):
+    loss_func = train_runtime.L2Loss(size_average=False)
+    totals = {
+        "rel_err": 0.0,
+        "abs_err": 0.0,
+        "mse": 0.0,
+        "mae": 0.0,
+        "maxae": 0.0,
+        "r2": 0.0,
+    }
+    count = 0
+    with train_runtime.torch.no_grad():
+        for pos, fx, y in test_loader:
+            x, fx, y = pos.to(device), fx.to(device), y.to(device)
+            if args.fun_dim == 0:
+                fx = None
+            out = model(x, fx)
+            if args.normalize:
+                out = dataset.y_normalizer.decode(out)
+
+            batch_size = y.shape[0]
+            totals["rel_err"] += loss_func.rel(out, y).item()
+            totals["abs_err"] += loss_func.abs(out, y).item()
+            totals["mse"] += loss_func.MSE(out, y).item()
+            totals["mae"] += loss_func.MAE(out, y).item()
+            totals["maxae"] += loss_func.MaxAE(out, y).item()
+            totals["r2"] += loss_func.R2Score(out, y).item()
+            count += batch_size
+
+    return {name: value / args.ntest for name, value in totals.items()}
+
+
+def main():
+    os.chdir(PROJECT_ROOT)
+    cfg = train_runtime.load_config()
+    if cfg.get("runtime", {}).get("device", "cpu") == "cpu":
+        os.environ["CUDA_VISIBLE_DEVICES"] = ""
+
+    weight_path = PROJECT_ROOT / cfg["paths"]["weight_path"]
+    if not weight_path.exists():
+        raise FileNotFoundError(
+            f"Checkpoint not found: {weight_path}. Run python scripts/train.py first "
+            "or update paths.weight_path/save_name in config/config.yaml."
+        )
+
+    args = train_runtime.build_args(cfg)
+    args.eval = cfg["inference"]["eval"]
+    train_runtime.load_onescience()
+    train_runtime.DistributedManager.initialize()
+    dist = train_runtime.DistributedManager()
+    dataset, _, test_loader, args.shapelist = train_runtime.get_data(args, dist)
+    device = train_runtime.get_device(args, dist)
+    if hasattr(dataset, "x_normalizer"):
+        dataset.x_normalizer = dataset.x_normalizer.to(device)
+    if hasattr(dataset, "y_normalizer"):
+        dataset.y_normalizer = dataset.y_normalizer.to(device)
+
+    model = load_model(args, device)
+    metrics = evaluate(model, test_loader, dataset, args, device)
+
+    result_dir = PROJECT_ROOT / cfg["paths"]["result_dir"] / args.save_name
+    result_dir.mkdir(parents=True, exist_ok=True)
+    with open(result_dir / "metrics.json", "w", encoding="utf-8") as f:
+        json.dump(metrics, f, indent=2, ensure_ascii=False)
+
+    print("\n===== 测试结果 =====")
+    print(f"平均相对误差: {metrics['rel_err']:.6e}")
+    print(f"平均绝对误差: {metrics['abs_err']:.6e}")
+    print(f"平均MSE: {metrics['mse']:.6e}")
+    print(f"平均MAE: {metrics['mae']:.6e}")
+    print(f"平均MaxAE: {metrics['maxae']:.6e}")
+    print(f"平均R2分数: {metrics['r2']:.6f}")
+    print("====================")
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,44 @@
+import json
+from pathlib import Path
+
+import torch
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+
+
+def load_config():
+    with open(PROJECT_ROOT / "config" / "config.yaml", "r", encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+
+def main():
+    cfg = load_config()
+    checkpoint = PROJECT_ROOT / cfg["paths"]["weight_path"]
+    result_dir = PROJECT_ROOT / cfg["paths"]["result_dir"] / cfg["train"]["save_name"]
+    metrics_path = result_dir / "metrics.json"
+
+    summary = {
+        "checkpoint": str(checkpoint),
+        "checkpoint_exists": checkpoint.exists(),
+        "result_dir": str(result_dir),
+        "result_dir_exists": result_dir.exists(),
+        "metrics": str(metrics_path),
+        "metrics_exists": metrics_path.exists(),
+    }
+    if checkpoint.exists():
+        state = torch.load(checkpoint, map_location="cpu", weights_only=False)
+        if isinstance(state, dict):
+            summary["epoch"] = state.get("epoch")
+            summary["best_epoch"] = state.get("best_epoch")
+            summary["best_test_loss"] = state.get("best_test_loss")
+            summary["has_model_state"] = "model_state" in state
+    if metrics_path.exists():
+        with open(metrics_path, "r", encoding="utf-8") as f:
+            summary["metrics_values"] = json.load(f)
+    print(json.dumps(summary, indent=2, ensure_ascii=False))
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,274 @@
+import os
+from pathlib import Path
+from types import SimpleNamespace
+
+import yaml
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+torch = None
+get_data = None
+DistributedManager = None
+replace_function = None
+get_model = None
+DerivLoss = None
+L2Loss = None
+
+
+def load_config():
+    with open(PROJECT_ROOT / "config" / "config.yaml", "r", encoding="utf-8") as f:
+        return yaml.safe_load(f)
+
+
+def load_onescience():
+    global torch
+    global get_data
+    global DistributedManager
+    global replace_function
+    global get_model
+    global DerivLoss
+    global L2Loss
+
+    if torch is not None:
+        return
+
+    import torch as torch_module
+    from onescience.datapipes.cfd_benchmark.data_factory import get_data as get_data_fn
+    from onescience.distributed.manager import DistributedManager as DistributedManager_cls
+    from onescience.memory.checkpoint import replace_function as replace_function_fn
+    from onescience.models.cfd_benchmark.model_factory import get_model as get_model_fn
+    from onescience.utils.cfd_benchmark.loss import DerivLoss as DerivLoss_cls
+    from onescience.utils.cfd_benchmark.loss import L2Loss as L2Loss_cls
+
+    torch = torch_module
+    get_data = get_data_fn
+    DistributedManager = DistributedManager_cls
+    replace_function = replace_function_fn
+    get_model = get_model_fn
+    DerivLoss = DerivLoss_cls
+    L2Loss = L2Loss_cls
+
+
+def build_args(cfg):
+    data = cfg["data"]
+    model = cfg["model"]
+    train = cfg["train"]
+    return SimpleNamespace(
+        lr=train["lr"],
+        epochs=train["epochs"],
+        weight_decay=train["weight_decay"],
+        pct_start=train["pct_start"],
+        batch_size=data["batch_size"],
+        gpu=0,
+        max_grad_norm=train["max_grad_norm"],
+        derivloss=train["derivloss"],
+        optimizer=train["optimizer"],
+        scheduler=train["scheduler"],
+        step_size=100,
+        gamma=0.5,
+        find_unused_parameters=False,
+        use_checkpoint=train["use_checkpoint"],
+        checkpoint_layers=train["checkpoint_layers"],
+        resume=False,
+        data_path=cfg["paths"]["data_dir"],
+        loader=data["loader"],
+        config_name="config_name",
+        train_ratio=0.8,
+        ntrain=data["ntrain"],
+        ntest=data["ntest"],
+        normalize=data["normalize"],
+        norm_type=data["norm_type"],
+        geotype=data["geotype"],
+        time_input=False,
+        space_dim=data["space_dim"],
+        fun_dim=data["fun_dim"],
+        out_dim=data["out_dim"],
+        shapelist=None,
+        downsamplex=data["downsamplex"],
+        downsampley=data["downsampley"],
+        downsamplez=1,
+        radius=0.2,
+        task=data["task"],
+        T_in=10,
+        T_out=10,
+        model=model["name"],
+        n_hidden=model["n_hidden"],
+        n_layers=model["n_layers"],
+        n_heads=model["n_heads"],
+        act=model["act"],
+        mlp_ratio=model["mlp_ratio"],
+        dropout=model["dropout"],
+        unified_pos=model["unified_pos"],
+        ref=model["ref"],
+        slice_num=model["slice_num"],
+        modes=12,
+        psi_dim=8,
+        attn_type="nystrom",
+        mwt_k=3,
+        branch_depth=5,
+        trunk_depth=6,
+        hidden_channels=[],
+        kernel_size=5,
+        emb_dims=128,
+        eval=0,
+        save_name=train["save_name"],
+        vis_num=cfg["inference"]["vis_num"],
+        vis_bound=None,
+    )
+
+
+def get_device(args, dist):
+    if dist.world_size == 1 and torch.cuda.is_available():
+        return torch.device(f"cuda:{args.gpu}")
+    if dist.world_size > 1 and torch.cuda.is_available():
+        return dist.device
+    return torch.device("cpu")
+
+
+def make_optimizer(args, model):
+    if args.optimizer == "AdamW":
+        return torch.optim.AdamW(
+            model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+        )
+    if args.optimizer == "Adam":
+        return torch.optim.Adam(
+            model.parameters(), lr=args.lr, weight_decay=args.weight_decay
+        )
+    raise ValueError("Optimizer only AdamW or Adam")
+
+
+def make_scheduler(args, optimizer, train_loader):
+    if args.scheduler == "OneCycleLR":
+        return torch.optim.lr_scheduler.OneCycleLR(
+            optimizer,
+            max_lr=args.lr,
+            epochs=args.epochs,
+            steps_per_epoch=len(train_loader),
+            pct_start=args.pct_start,
+        )
+    if args.scheduler == "CosineAnnealingLR":
+        return torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=args.epochs)
+    if args.scheduler == "StepLR":
+        return torch.optim.lr_scheduler.StepLR(
+            optimizer, step_size=args.step_size, gamma=args.gamma
+        )
+    raise ValueError("Scheduler only OneCycleLR, CosineAnnealingLR or StepLR")
+
+
+def evaluate(model, test_loader, dataset, args, device):
+    myloss = L2Loss(size_average=False)
+    model.eval()
+    rel_err = 0.0
+    with torch.no_grad():
+        for pos, fx, y in test_loader:
+            x, fx, y = pos.to(device), fx.to(device), y.to(device)
+            if args.fun_dim == 0:
+                fx = None
+            out = model(x, fx)
+            if args.normalize:
+                out = dataset.y_normalizer.decode(out)
+            rel_err += myloss(out, y).item()
+    return rel_err / args.ntest
+
+
+def train(args):
+    load_onescience()
+    DistributedManager.initialize()
+    dist = DistributedManager()
+    dataset, train_loader, test_loader, args.shapelist = get_data(args, dist)
+    device = get_device(args, dist)
+    model = get_model(args, device).to(device)
+    if hasattr(dataset, "x_normalizer"):
+        dataset.x_normalizer = dataset.x_normalizer.to(device)
+    if hasattr(dataset, "y_normalizer"):
+        dataset.y_normalizer = dataset.y_normalizer.to(device)
+
+    optimizer = make_optimizer(args, model)
+    scheduler = make_scheduler(args, optimizer, train_loader)
+    checkpoint_dir = PROJECT_ROOT / "checkpoints"
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+    checkpoint_path = checkpoint_dir / f"{args.save_name}.pt"
+    checkpoint_layers = (
+        [layer.strip() for layer in args.checkpoint_layers.split(",") if layer.strip()]
+        if args.use_checkpoint and args.checkpoint_layers
+        else []
+    )
+    myloss = L2Loss(size_average=False)
+    regloss = DerivLoss(size_average=False, shapelist=args.shapelist) if args.derivloss else None
+    best_test_loss = float("inf")
+    best_epoch = 0
+
+    print(args)
+    print(model)
+    print(f"Use device: {device}")
+
+    for epoch in range(args.epochs):
+        model.train()
+        train_loss = 0.0
+        for pos, fx, y in train_loader:
+            x, fx, y = pos.to(device), fx.to(device), y.to(device)
+            if args.fun_dim == 0:
+                fx = None
+            with replace_function(
+                module=model,
+                replace_layers_list=checkpoint_layers,
+                ddp_flag=(dist.world_size > 1),
+            ):
+                out = model(x, fx)
+            if args.normalize:
+                out = dataset.y_normalizer.decode(out)
+                y = dataset.y_normalizer.decode(y)
+            loss = myloss(out, y)
+            if regloss is not None:
+                loss = loss + 0.1 * regloss(out, y)
+
+            train_loss += loss.item()
+            optimizer.zero_grad()
+            loss.backward()
+            if args.max_grad_norm is not None:
+                torch.nn.utils.clip_grad_norm_(model.parameters(), args.max_grad_norm)
+            optimizer.step()
+            if args.scheduler == "OneCycleLR":
+                scheduler.step()
+
+        if args.scheduler in {"CosineAnnealingLR", "StepLR"}:
+            scheduler.step()
+
+        train_loss = train_loss / args.ntrain
+        rel_err = evaluate(model, test_loader, dataset, args, device)
+        if rel_err < best_test_loss:
+            best_test_loss = rel_err
+            best_epoch = epoch
+            torch.save(
+                {
+                    "epoch": epoch,
+                    "model_state": model.state_dict(),
+                    "optimizer_state": optimizer.state_dict(),
+                    "scheduler_state": scheduler.state_dict(),
+                    "best_test_loss": best_test_loss,
+                    "best_epoch": best_epoch,
+                    "args": vars(args),
+                },
+                checkpoint_path,
+            )
+        if epoch % 10 == 0:
+            print("Epoch {} Train loss : {:.5f}".format(epoch, train_loss))
+            print("rel_err:{}".format(rel_err))
+
+    print(
+        "Training completed. Best model saved at epoch {} with rel_err: {:.5f}".format(
+            best_epoch, best_test_loss
+        )
+    )
+
+
+def main():
+    os.chdir(PROJECT_ROOT)
+    cfg = load_config()
+    if cfg.get("runtime", {}).get("device", "cpu") == "cpu":
+        os.environ["CUDA_VISIBLE_DEVICES"] = ""
+    train(build_args(cfg))
+
+
+if __name__ == "__main__":
+    main()