Upload folder using huggingface_hub

README.md

@@ -0,0 +1,127 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- BENO
+- elliptic PDEs
+- boundary-embedded neural operator
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">BENO</span>
+  </strong>
+</p>
+
+# Model Overview
+
+BENO is a boundary-embedded neural operator developed by researchers at Peking University for solving elliptic PDEs subject to complex boundary conditions. It enables rapid prediction of steady-state physical fields across varying geometries and nonhomogeneous boundary conditions.
+
+Paper: [BENO: Boundary-embedded Neural Operators for Elliptic PDEs](https://openreview.net/forum?id=ZZTkLDRmkg)
+
+# Model Description
+BENO employs a boundary-embedded neural operator architecture that combines a dual-branch graph neural network with a Transformer. Trained on the BENO dataset, it predicts steady-state physical fields governed by elliptic PDEs such as the Poisson and Laplace equations on complex geometries with nonhomogeneous boundary conditions.
+
+# Use Cases
+
+| Use Case | Description |
+| ---------- | --------------------------------- |
+| Elliptic PDE solving | Rapidly approximate steady-state boundary-value problems such as the Poisson and Laplace equations |
+| Complex boundary modeling | Capture the effects of free-form boundaries, irregular domains, and nonhomogeneous boundary values on the solution field |
+| Steady-state field prediction | Predict equilibrium physical fields determined jointly by source terms and boundary conditions |
+| Numerical solver acceleration | Replace or augment conventional FEM, FDM, and FVM workflows to improve inference efficiency |
+
+# 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/BENO --local_dir ./BENO
+cd BENO
+```
+
+### Set Up the Runtime Environment
+
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# Installation with uv is also supported
+pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+**GPU Environment**
+```bash
+# Activate Conda first
+conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
+conda activate onescience311
+# Installation with uv is also supported
+pip install onescience[cfd-gpu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+### Training Data
+The OneScience community provides the `beno` dataset for training. Download it with the command below and verify that the data path in `config/config.yaml` is configured correctly:
+
+```bash
+modelscope download --dataset OneScience/beno --local_dir ./data
+```
+The complete dataset is also available from the [official download link](https://drive.google.com/file/d/11PbUrzJ-b18VhFGY_uICSciCkeGrsaTZ/view).
+
+### Training
+
+Single GPU:
+
+```bash
+python scripts/train.py
+```
+
+Multiple GPUs:
+```
+torchrun --standalone --nnodes=<num_nodes> --nproc_per_node=<num_GPUs> scripts/train.py
+```
+
+### Model Weights
+This repository will provide weights trained on the BENO dataset in the `weights/` directory. The weights will be uploaded soon.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+# Official OneScience Resources
+
+| Platform | OneScience Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citations and License
+
+- Original BENO paper: [BENO: Boundary-embedded Neural Operators for Elliptic PDEs](https://proceedings.iclr.cc/paper_files/paper/2024/file/218ca0d92e6ed8f9db00621e103dc70c-Paper-Conference.pdf)
+- This repository retains source attribution and has been adapted for automated execution through OneScience and ModelScope. Before public redistribution, verify the applicable licensing requirements of the upstream project.

conf/config.yaml

@@ -0,0 +1,66 @@
+root:
+  datapipe:
+    source:
+      data_dir: "${ONESCIENCE_BENO_DATA_DIR}/Dirichlet/"
+      cache_dir: "./cache_data"
+      file_prefix: "N8_smoke"
+    verbose: false
+
+    data:
+      name: "BENO_Elliptic"
+      ntrain: 2
+      ntest: 1
+      resolution: 8
+      ns: 4
+      domain_bounds: [[0, 1], [0, 1]]
+
+    dataloader:
+      batch_size: 1
+      num_workers: 0
+      pin_memory: false
+
+  model:
+    name: "HeteroGNS"
+    nnode_in_features: 10
+    nnode_out_features: 1
+    nedge_in_features: 7
+    nmlp_layers: 2
+    act: "relu"
+    boundary_dim: 16
+    trans_layer: 1
+    latent_dim: 16
+    nmessage_passing_steps: 1
+    mlp_hidden_dim: 16
+
+  training:
+    device: "cpu"
+    output_dir: "./weight"
+    epochs: 1
+    save_period: 1
+    checkpoint_name: "model_epoch_{epoch}.pt"
+    seed: 2025
+
+    optimizer:
+      name: "Adam"
+      lr: 1.0e-3
+      weight_decay: 5.0e-4
+
+    scheduler:
+      name: "CosineAnnealingWarmRestarts"
+      T_0: 2
+      T_mult: 1
+
+  inference:
+    device: "cpu"
+    checkpoint_dir: "./weight"
+    checkpoint_path: ""
+    result_dir: "./result/output"
+    picture_dir: "./result/picture"
+    metrics_path: "./result/metrics.json"
+    num_visualize: 1
+
+  fake_data:
+    root_dir: "./data"
+    seed: 2025
+    boundary_points: 128
+    noise_std: 0.001

configuration.json

@@ -0,0 +1,14 @@
+{
+  "framework": "Pytorch",
+  "task": "other",
+  "model": {
+    "type": "BENO",
+    "entrypoint": "model.HeteroGNS"
+  },
+  "pipeline": {
+    "fake_data": "scripts/fake_data.py",
+    "train": "scripts/train.py",
+    "inference": "scripts/inference.py",
+    "result": "scripts/result.py"
+  }
+}

model/BE_MPNN.py

@@ -0,0 +1,316 @@
+from typing import List
+import torch
+import torch.nn as nn
+from torch_geometric.nn import MessagePassing
+from .transformer import *
+
+def build_mlp(
+        input_size: int,
+        hidden_layer_sizes: List[int],
+        output_size: int = None,
+        output_activation: nn.Module = nn.Identity,
+        activation: nn.Module = nn.ReLU) -> nn.Module:
+
+  # Size of each layer
+  layer_sizes = [input_size] + hidden_layer_sizes
+  if output_size:
+    layer_sizes.append(output_size)
+
+  # Number of layers
+  nlayers = len(layer_sizes) - 1
+
+  # Create a list of activation functions and
+  # set the last element to output activation function
+  act = [activation for i in range(nlayers)]
+  act[-1] = output_activation
+
+  # Create a torch sequential container
+  mlp = nn.Sequential()
+  for i in range(nlayers):
+    mlp.add_module("NN-" + str(i), nn.Linear(layer_sizes[i],
+                                             layer_sizes[i + 1]))
+    mlp.add_module("Act-" + str(i), act[i]())
+
+  return mlp
+
+
+class Encoder(nn.Module):
+
+
+  def __init__(
+          self,
+          nnode_in_features: int,
+          nnode_out_features: int, #latent dim=128
+          nedge_in_features: int,
+          nedge_out_features: int,
+          nmlp_layers: int,
+          mlp_hidden_dim: int,
+          activation:nn.Module,):
+ 
+    super(Encoder, self).__init__()
+    # Encode node features as an MLP
+    self.node_fn = nn.Sequential(*[build_mlp(nnode_in_features,
+                                             [mlp_hidden_dim
+                                              for _ in range(nmlp_layers)],
+                                             nnode_out_features,activation=activation),
+                                   nn.LayerNorm(nnode_out_features)])
+    # Encode edge features as an MLP
+    self.edge_fn = nn.Sequential(*[build_mlp(nedge_in_features,
+                                             [mlp_hidden_dim
+                                              for _ in range(nmlp_layers)],
+                                             nedge_out_features,activation=activation),
+                                   nn.LayerNorm(nedge_out_features)])
+    
+    self.node_fn_inbd = nn.Sequential(*[build_mlp(nnode_in_features,
+                                             [mlp_hidden_dim
+                                              for _ in range(nmlp_layers)],
+                                             nnode_out_features,activation=activation),
+                                   nn.LayerNorm(nnode_out_features)])
+    
+    self.edge_fn_inbd = nn.Sequential(*[build_mlp(nedge_in_features,
+                                             [mlp_hidden_dim
+                                              for _ in range(nmlp_layers)],
+                                             nedge_out_features,activation=activation),
+                                   nn.LayerNorm(nedge_out_features)])
+
+  def forward(
+          self,
+          x: torch.tensor,
+          edge_features: torch.tensor,
+          x_inbd:torch.tensor,
+          edge_inbd_features: torch.tensor):
+
+    return self.node_fn(x), self.edge_fn(edge_features), self.node_fn_inbd(x_inbd), self.edge_fn_inbd(edge_inbd_features)
+
+
+class InteractionNetwork(MessagePassing):
+  def __init__(
+      self,
+      nnode_in: int,
+      nnode_out: int,
+      nedge_in: int,
+      nedge_out: int,
+      nmlp_layers: int,
+      mlp_hidden_dim: int,
+      boundary_dim: int,
+      trans_layer:int,
+      activation: nn.Module = nn.ReLU
+  ):
+
+    # Aggregate features from neighbors
+    super(InteractionNetwork, self).__init__(aggr='mean')
+    # Node MLP
+    self.node_fn = nn.Sequential(*[build_mlp(nnode_in + nedge_out + boundary_dim,
+                                            [mlp_hidden_dim
+                                             for _ in range(nmlp_layers)],  
+                                             nnode_out,activation=activation),
+                                   nn.LayerNorm(nnode_out)])
+
+    self.edge_fn = nn.Sequential(*[build_mlp(nnode_in + nnode_in + nedge_in,
+                                             [mlp_hidden_dim
+                                              for _ in range(nmlp_layers)],
+                                             nedge_out,activation=activation),
+                                   nn.LayerNorm(nedge_out)])
+
+    #boundary
+    self.boundary_fn = Transformer(enc_in = 3, d_model = boundary_dim, n_heads = 2, enc_layers = trans_layer)
+
+  def forward(self,
+              x: torch.tensor,
+              edge_index: torch.tensor,
+              edge_features: torch.tensor,
+              boundary: torch.tensor,
+              
+              ):
+
+    boundary = boundary.unsqueeze(0).float()
+   
+    boundary = self.boundary_fn(boundary)
+   
+    x_residual = x.clone()
+    edge_features_residual = edge_features.clone()
+
+    x, edge_features = self.propagate(
+        edge_index=edge_index, x=x, edge_features=edge_features, boundary=boundary)
+
+
+    return x + x_residual, edge_features + edge_features_residual 
+
+  def message(self,
+              x_i: torch.tensor,
+              x_j: torch.tensor,
+              edge_features: torch.tensor,
+              boundary: torch.tensor
+              ) -> torch.tensor:
+
+    # Concat edge features with a final shape of [nedges, latent_dim*3]
+
+    edge_features = torch.cat([x_i, x_j, edge_features], dim=-1)
+
+    edge_features = self.edge_fn(edge_features)
+    return edge_features
+
+  def update(self,
+             x_updated: torch.tensor,
+             x: torch.tensor,
+             edge_features: torch.tensor,
+             boundary: torch.tensor
+             ):
+ 
+    # Concat node features with a final shape of
+    boundary_all = boundary.repeat(x.shape[0], 1)
+    x_updated = torch.cat([x_updated, x,boundary_all], dim=-1)
+    x_updated = self.node_fn(x_updated)
+
+    return x_updated, edge_features
+
+class Processor(MessagePassing):
+
+
+  def __init__(
+      self,
+      nnode_in: int,
+      nnode_out: int,
+      nedge_in: int,
+      nedge_out: int,
+      nmessage_passing_steps: int,
+      nmlp_layers: int,
+      mlp_hidden_dim: int,
+      boundary_dim: int,
+      trans_layer:int,
+      activation:nn.Module
+  ):
+
+    super(Processor, self).__init__(aggr='mean')
+    # Create a stack of M Graph Networks GNs.
+    self.gnn_stacks = nn.ModuleList([
+        InteractionNetwork(
+            nnode_in=nnode_in,
+            nnode_out=nnode_out,
+            nedge_in=nedge_in,
+            nedge_out=nedge_out,
+            nmlp_layers=nmlp_layers,
+            mlp_hidden_dim=mlp_hidden_dim,
+            boundary_dim=boundary_dim,
+            trans_layer=trans_layer,
+            activation=activation
+        ) for _ in range(nmessage_passing_steps)])
+    
+    self.gnn_stacks_inbd = nn.ModuleList([
+        InteractionNetwork(
+            nnode_in=nnode_in,
+            nnode_out=nnode_out,
+            nedge_in=nedge_in,
+            nedge_out=nedge_out,
+            nmlp_layers=nmlp_layers,
+            mlp_hidden_dim=mlp_hidden_dim,
+            boundary_dim=boundary_dim,
+            trans_layer=trans_layer,
+            activation=activation,
+        ) for _ in range(nmessage_passing_steps)])
+
+
+  def forward(self,
+              x: torch.tensor,
+              edge_index: torch.tensor,
+              edge_features: torch.tensor,
+              boundary:torch.tensor,
+              x_inbd: torch.tensor,
+              edge_inbd_index: torch.tensor,
+              edge_inbd_features: torch.tensor,
+              boundary_inbd:torch.tensor):  
+
+    for gnn in self.gnn_stacks:
+      x, edge_features= gnn(x, edge_index, edge_features,boundary )  
+      
+    for gnn in self.gnn_stacks_inbd:
+      x_inbd, edge_inbd_features = gnn(x_inbd, edge_inbd_index, edge_inbd_features,boundary_inbd)  
+    
+    return x, edge_features, x_inbd, edge_inbd_features
+
+class Decoder(nn.Module):
+
+
+  def __init__(
+          self,
+          nnode_in: int,
+          nnode_out: int,
+          nmlp_layers: int,
+          mlp_hidden_dim: int,
+          activation:nn.Module):
+
+    super(Decoder, self).__init__()
+    self.node_fn = build_mlp(
+        nnode_in, [mlp_hidden_dim for _ in range(nmlp_layers)], nnode_out,activation=activation)
+    self.node_fn_inbd = build_mlp(
+        nnode_in, [mlp_hidden_dim for _ in range(nmlp_layers)], nnode_out,activation=activation)
+
+  def forward(self,
+              x: torch.tensor,
+              x_inbd: torch.tensor):
+    u1=self.node_fn(x)
+    u2=self.node_fn_inbd(x_inbd)
+    u=u1+u2
+
+    return u
+
+class HeteroGNS(nn.Module):
+  def __init__(
+      self,
+      nnode_in_features: int,
+      nnode_out_features: int,
+      nedge_in_features: int,
+      latent_dim: int = 128,
+      nmessage_passing_steps: int = 10,
+      nmlp_layers: int = 2,
+      mlp_hidden_dim: int = 128,
+      activation: nn.Module = nn.ELU,
+      boundary_dim: int=128,
+      trans_layer: int =3,
+  ):
+
+    super(HeteroGNS, self).__init__()
+    self._encoder = Encoder(
+        nnode_in_features=nnode_in_features,
+        nnode_out_features=latent_dim,
+        nedge_in_features=nedge_in_features,
+        nedge_out_features=latent_dim,
+        nmlp_layers=nmlp_layers,
+        mlp_hidden_dim=mlp_hidden_dim,
+        activation=activation,
+
+        
+    )
+    self._processor = Processor(
+        nnode_in=latent_dim,
+        nnode_out=latent_dim,
+        nedge_in=latent_dim,
+        nedge_out=latent_dim,
+        nmessage_passing_steps=nmessage_passing_steps,
+        nmlp_layers=nmlp_layers,
+        mlp_hidden_dim=mlp_hidden_dim,
+        boundary_dim = boundary_dim,
+        trans_layer = trans_layer,
+        activation=activation,
+    )
+    self._decoder = Decoder(
+        nnode_in=latent_dim,
+        nnode_out=nnode_out_features,
+        nmlp_layers=nmlp_layers,
+        mlp_hidden_dim=mlp_hidden_dim,
+        activation=activation,
+        
+    )
+
+  def forward(self,data):
+
+    x, edge_index, edge_features,boundary = data['G1'].x, data['G1'].edge_index, data['G1'].edge_features,data['G1'].boundary
+    x_inbd, edge_inbd_index, edge_inbd_features,boundary_inbd = data['G2'].x, data['G2'].edge_index,data['G2'].edge_features,data['G2'].boundary
+
+    x, edge_features,x_inbd,edge_inbd_features = self._encoder(x, edge_features,x_inbd,edge_inbd_features)
+
+    x, edge_features, x_inbd, edge_inbd_features = self._processor(x, edge_index, edge_features, boundary,x_inbd,edge_inbd_index, edge_inbd_features,boundary_inbd)
+
+    u = self._decoder(x,x_inbd)
+
+    return u

model/__init__.py

@@ -0,0 +1,3 @@
+from .BE_MPNN import HeteroGNS
+
+__all__ = ["HeteroGNS"]

model/transformer.py

@@ -0,0 +1,547 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+import math
+from math import sqrt
+
+
+class PositionalEmbedding(nn.Module):
+    def __init__(self, d_model, max_len=5000):
+        super(PositionalEmbedding, self).__init__()
+        # Compute the positional encodings once in log space.
+        pe = torch.zeros(max_len, d_model).float()
+        pe.require_grad = False
+
+        position = torch.arange(0, max_len).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
+
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+
+        pe = pe.unsqueeze(0)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        return self.pe[:, :x.size(1)]
+
+
+class TokenEmbedding(nn.Module):
+    def __init__(self, c_in, d_model):
+        super(TokenEmbedding, self).__init__()
+        padding = 1 if torch.__version__ >= '1.5.0' else 2
+        self.tokenConv = nn.Conv1d(in_channels=c_in, out_channels=d_model,
+                                   kernel_size=3, padding=padding, padding_mode='circular', bias=False)
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_in', nonlinearity='leaky_relu')
+
+    def forward(self, x):
+        x = self.tokenConv(x.permute(0, 2, 1)).transpose(1, 2)
+        return x
+
+
+class FixedEmbedding(nn.Module):
+    def __init__(self, c_in, d_model):
+        super(FixedEmbedding, self).__init__()
+
+        w = torch.zeros(c_in, d_model).float()
+        w.require_grad = False
+
+        position = torch.arange(0, c_in).float().unsqueeze(1)
+        div_term = (torch.arange(0, d_model, 2).float() * -(math.log(10000.0) / d_model)).exp()
+
+        w[:, 0::2] = torch.sin(position * div_term)
+        w[:, 1::2] = torch.cos(position * div_term)
+
+        self.emb = nn.Embedding(c_in, d_model)
+        self.emb.weight = nn.Parameter(w, requires_grad=False)
+
+    def forward(self, x):
+        return self.emb(x).detach()
+
+
+class TemporalEmbedding(nn.Module):
+    def __init__(self, d_model, embed_type='fixed', freq='h'):
+        super(TemporalEmbedding, self).__init__()
+
+        minute_size = 4
+        hour_size = 24
+        weekday_size = 7
+        day_size = 32
+        month_size = 13
+
+        Embed = FixedEmbedding if embed_type == 'fixed' else nn.Embedding
+        if freq == 't':
+            self.minute_embed = Embed(minute_size, d_model)
+        self.hour_embed = Embed(hour_size, d_model)
+        self.weekday_embed = Embed(weekday_size, d_model)
+        self.day_embed = Embed(day_size, d_model)
+        self.month_embed = Embed(month_size, d_model)
+
+    def forward(self, x):
+        x = x.long()
+
+        minute_x = self.minute_embed(x[:, :, 4]) if hasattr(self, 'minute_embed') else 0.
+        hour_x = self.hour_embed(x[:, :, 3])
+        weekday_x = self.weekday_embed(x[:, :, 2])
+        day_x = self.day_embed(x[:, :, 1])
+        month_x = self.month_embed(x[:, :, 0])
+
+        return hour_x + weekday_x + day_x + month_x + minute_x
+
+
+class TimeFeatureEmbedding(nn.Module):
+    def __init__(self, d_model, embed_type='timeF', freq='h'):
+        super(TimeFeatureEmbedding, self).__init__()
+
+        freq_map = {'h': 4, 't': 5, 's': 6, 'm': 1, 'a': 1, 'w': 2, 'd': 3, 'b': 3}
+        d_inp = freq_map[freq]
+        self.embed = nn.Linear(d_inp, d_model, bias=False)
+
+    def forward(self, x):
+        return self.embed(x)
+
+
+class DataEmbedding(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0):
+        super(DataEmbedding, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        x = self.value_embedding(x) + self.temporal_embedding(x_mark) + self.position_embedding(x)
+        return self.dropout(x)
+
+
+class DataEmbedding_wo_pos(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0):
+        super(DataEmbedding_wo_pos, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        x = self.value_embedding(x) + self.temporal_embedding(x_mark)
+        return self.dropout(x)
+
+
+class DataEmbedding_wo_pos_temp(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0):
+        super(DataEmbedding_wo_pos_temp, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        x = self.value_embedding(x)
+        return self.dropout(x)
+
+
+class DataEmbedding_wo_temp(nn.Module):
+    def __init__(self, c_in, d_model, embed_type='fixed', freq='h', dropout=0):
+        super(DataEmbedding_wo_temp, self).__init__()
+
+        self.value_embedding = TokenEmbedding(c_in=c_in, d_model=d_model)
+        self.position_embedding = PositionalEmbedding(d_model=d_model)
+        self.temporal_embedding = TemporalEmbedding(d_model=d_model, embed_type=embed_type,
+                                                    freq=freq) if embed_type != 'timeF' else TimeFeatureEmbedding(
+            d_model=d_model, embed_type=embed_type, freq=freq)
+        self.dropout = nn.Dropout(p=dropout)
+
+    def forward(self, x, x_mark):
+        x = self.value_embedding(x) + self.position_embedding(x)
+        return self.dropout(x)
+
+
+class TriangularCausalMask():
+    def __init__(self, B, L, device="cpu"):
+        mask_shape = [B, 1, L, L]
+        with torch.no_grad():
+            self._mask = torch.triu(torch.ones(mask_shape, dtype=torch.bool), diagonal=1).to(device)
+
+    @property
+    def mask(self):
+        return self._mask
+
+
+class ProbMask():
+    def __init__(self, B, H, L, index, scores, device="cpu"):
+        _mask = torch.ones(L, scores.shape[-1], dtype=torch.bool).to(device).triu(1)
+        _mask_ex = _mask[None, None, :].expand(B, H, L, scores.shape[-1])
+        indicator = _mask_ex[torch.arange(B)[:, None, None],
+                    torch.arange(H)[None, :, None],
+                    index, :].to(device)
+        self._mask = indicator.view(scores.shape).to(device)
+
+    @property
+    def mask(self):
+        return self._mask
+
+
+class FullAttention(nn.Module):
+    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0, output_attention=False):
+        super(FullAttention, self).__init__()
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.output_attention = output_attention
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, H, E = queries.shape
+        _, S, _, D = values.shape
+        scale = self.scale or 1. / sqrt(E)
+
+        scores = torch.einsum("blhe,bshe->bhls", queries, keys)
+
+        if self.mask_flag:
+            if attn_mask is None:
+                attn_mask = TriangularCausalMask(B, L, device=queries.device)
+
+            scores.masked_fill_(attn_mask.mask, -np.inf)
+
+        A = self.dropout(torch.softmax(scale * scores, dim=-1))
+        V = torch.einsum("bhls,bshd->blhd", A, values)
+
+        if self.output_attention:
+            return (V.contiguous(), A)
+        else:
+            return (V.contiguous(), None)
+
+
+class ProbAttention(nn.Module):
+    def __init__(self, mask_flag=True, factor=5, scale=None, attention_dropout=0, output_attention=False):
+        super(ProbAttention, self).__init__()
+        self.factor = factor
+        self.scale = scale
+        self.mask_flag = mask_flag
+        self.output_attention = output_attention
+        self.dropout = nn.Dropout(attention_dropout)
+
+    def _prob_QK(self, Q, K, sample_k, n_top):  # n_top: c*ln(L_q)
+        # Q [B, H, L, D]
+        B, H, L_K, E = K.shape
+        _, _, L_Q, _ = Q.shape
+
+        # calculate the sampled Q_K
+        K_expand = K.unsqueeze(-3).expand(B, H, L_Q, L_K, E)
+        index_sample = torch.randint(L_K, (L_Q, sample_k))  # real U = U_part(factor*ln(L_k))*L_q
+        K_sample = K_expand[:, :, torch.arange(L_Q).unsqueeze(1), index_sample, :]
+        Q_K_sample = torch.matmul(Q.unsqueeze(-2), K_sample.transpose(-2, -1)).squeeze()
+
+        # find the Top_k query with sparisty measurement
+        M = Q_K_sample.max(-1)[0] - torch.div(Q_K_sample.sum(-1), L_K)
+        M_top = M.topk(n_top, sorted=False)[1]
+
+        # use the reduced Q to calculate Q_K
+        Q_reduce = Q[torch.arange(B)[:, None, None],
+                   torch.arange(H)[None, :, None],
+                   M_top, :]  # factor*ln(L_q)
+        Q_K = torch.matmul(Q_reduce, K.transpose(-2, -1))  # factor*ln(L_q)*L_k
+
+        return Q_K, M_top
+
+    def _get_initial_context(self, V, L_Q):
+        B, H, L_V, D = V.shape
+        if not self.mask_flag:
+            # V_sum = V.sum(dim=-2)
+            V_sum = V.mean(dim=-2)
+            contex = V_sum.unsqueeze(-2).expand(B, H, L_Q, V_sum.shape[-1]).clone()
+        else:  # use mask
+            assert (L_Q == L_V)  # requires that L_Q == L_V, i.e. for self-attention only
+            contex = V.cumsum(dim=-2)
+        return contex
+
+    def _update_context(self, context_in, V, scores, index, L_Q, attn_mask):
+        B, H, L_V, D = V.shape
+
+        if self.mask_flag:
+            attn_mask = ProbMask(B, H, L_Q, index, scores, device=V.device)
+            scores.masked_fill_(attn_mask.mask, -np.inf)
+
+        attn = torch.softmax(scores, dim=-1)  # nn.Softmax(dim=-1)(scores)
+
+        context_in[torch.arange(B)[:, None, None],
+        torch.arange(H)[None, :, None],
+        index, :] = torch.matmul(attn, V).type_as(context_in)
+        if self.output_attention:
+            attns = (torch.ones([B, H, L_V, L_V]) / L_V).type_as(attn).to(attn.device)
+            attns[torch.arange(B)[:, None, None], torch.arange(H)[None, :, None], index, :] = attn
+            return (context_in, attns)
+        else:
+            return (context_in, None)
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L_Q, H, D = queries.shape
+        _, L_K, _, _ = keys.shape
+
+        queries = queries.transpose(2, 1)
+        keys = keys.transpose(2, 1)
+        values = values.transpose(2, 1)
+
+        U_part = self.factor * np.ceil(np.log(L_K)).astype('int').item()  # c*ln(L_k)
+        u = self.factor * np.ceil(np.log(L_Q)).astype('int').item()  # c*ln(L_q)
+
+        U_part = U_part if U_part < L_K else L_K
+        u = u if u < L_Q else L_Q
+
+        scores_top, index = self._prob_QK(queries, keys, sample_k=U_part, n_top=u)
+
+        # add scale factor
+        scale = self.scale or 1. / sqrt(D)
+        if scale is not None:
+            scores_top = scores_top * scale
+        # get the context
+        context = self._get_initial_context(values, L_Q)
+        # update the context with selected top_k queries
+        context, attn = self._update_context(context, values, scores_top, index, L_Q, attn_mask)
+
+        return context.contiguous(), attn
+
+
+class AttentionLayer(nn.Module):
+    def __init__(self, attention, d_model, n_heads, d_keys=None,
+                 d_values=None):
+        super(AttentionLayer, self).__init__()
+
+        d_keys = d_keys or (d_model // n_heads)
+        d_values = d_values or (d_model // n_heads)
+
+        self.inner_attention = attention
+        self.query_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.key_projection = nn.Linear(d_model, d_keys * n_heads)
+        self.value_projection = nn.Linear(d_model, d_values * n_heads)
+        self.out_projection = nn.Linear(d_values * n_heads, d_model)
+        self.n_heads = n_heads
+
+    def forward(self, queries, keys, values, attn_mask):
+        B, L, _ = queries.shape
+        _, S, _ = keys.shape
+        H = self.n_heads
+
+        queries = self.query_projection(queries).view(B, L, H, -1)
+        keys = self.key_projection(keys).view(B, S, H, -1)
+        values = self.value_projection(values).view(B, S, H, -1)
+
+        out, attn = self.inner_attention(
+            queries,
+            keys,
+            values,
+            attn_mask
+        )
+        out = out.view(B, L, -1)
+
+        return self.out_projection(out), attn
+
+
+class ConvLayer(nn.Module):
+    def __init__(self, c_in):
+        super(ConvLayer, self).__init__()
+        self.downConv = nn.Conv1d(in_channels=c_in,
+                                  out_channels=c_in,
+                                  kernel_size=3,
+                                  padding=2,
+                                  padding_mode='circular')
+        self.norm = nn.BatchNorm1d(c_in)
+        self.activation = nn.ELU()
+        self.maxPool = nn.MaxPool1d(kernel_size=3, stride=2, padding=1)
+
+    def forward(self, x):
+        x = self.downConv(x.permute(0, 2, 1))
+        x = self.norm(x)
+        x = self.activation(x)
+        x = self.maxPool(x)
+        x = x.transpose(1, 2)
+        return x
+
+
+class EncoderLayer(nn.Module):
+    def __init__(self, attention, d_model, d_ff=None, dropout=0, activation="relu"):
+        super(EncoderLayer, self).__init__()
+        d_ff = d_ff or 4 * d_model
+        self.attention = attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, attn_mask=None):
+        new_x, attn = self.attention(
+            x, x, x,
+            attn_mask=attn_mask
+        )
+        x = x + self.dropout(new_x)
+
+        y = x = self.norm1(x)
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+        y = self.dropout(self.conv2(y).transpose(-1, 1))
+
+        return self.norm2(x + y), attn
+
+
+class Encoder(nn.Module):
+    def __init__(self, attn_layers, conv_layers=None, norm_layer=None):
+        super(Encoder, self).__init__()
+        self.attn_layers = nn.ModuleList(attn_layers)
+        self.conv_layers = nn.ModuleList(conv_layers) if conv_layers is not None else None
+        self.norm = norm_layer
+
+    def forward(self, x, attn_mask=None):
+        # x [B, L, D]
+        attns = []
+        if self.conv_layers is not None:
+            for attn_layer, conv_layer in zip(self.attn_layers, self.conv_layers):
+                x, attn = attn_layer(x, attn_mask=attn_mask)
+                x = conv_layer(x)
+                attns.append(attn)
+            x, attn = self.attn_layers[-1](x)
+            attns.append(attn)
+        else:
+            for attn_layer in self.attn_layers:
+                x, attn = attn_layer(x, attn_mask=attn_mask)
+                attns.append(attn)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        return x, attns
+
+
+class DecoderLayer(nn.Module):
+    def __init__(self, self_attention, cross_attention, d_model, d_ff=None,
+                 dropout=0, activation="relu"):
+        super(DecoderLayer, self).__init__()
+        d_ff = d_ff or 4 * d_model
+        self.self_attention = self_attention
+        self.cross_attention = cross_attention
+        self.conv1 = nn.Conv1d(in_channels=d_model, out_channels=d_ff, kernel_size=1)
+        self.conv2 = nn.Conv1d(in_channels=d_ff, out_channels=d_model, kernel_size=1)
+        self.norm1 = nn.LayerNorm(d_model)
+        self.norm2 = nn.LayerNorm(d_model)
+        self.norm3 = nn.LayerNorm(d_model)
+        self.dropout = nn.Dropout(dropout)
+        self.activation = F.relu if activation == "relu" else F.gelu
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None):
+        x = x + self.dropout(self.self_attention(
+            x, x, x,
+            attn_mask=x_mask
+        )[0])
+        x = self.norm1(x)
+
+        x = x + self.dropout(self.cross_attention(
+            x, cross, cross,
+            attn_mask=cross_mask
+        )[0])
+
+        y = x = self.norm2(x)
+        y = self.dropout(self.activation(self.conv1(y.transpose(-1, 1))))
+        y = self.dropout(self.conv2(y).transpose(-1, 1))
+
+        return self.norm3(x + y)
+
+
+class Decoder(nn.Module):
+    def __init__(self, layers, norm_layer=None, projection=None):
+        super(Decoder, self).__init__()
+        self.layers = nn.ModuleList(layers)
+        self.norm = norm_layer
+        self.projection = projection
+
+    def forward(self, x, cross, x_mask=None, cross_mask=None):
+        for layer in self.layers:
+            x = layer(x, cross, x_mask=x_mask, cross_mask=cross_mask)
+
+        if self.norm is not None:
+            x = self.norm(x)
+
+        if self.projection is not None:
+            x = self.projection(x)
+        return x
+
+class Transformer(nn.Module):
+    """
+    Vanilla Transformer with O(L^2) complexity
+    """
+    def __init__(self, enc_in = 2, d_model = 16, n_heads = 2, hidden_size=48, dropout=0, d_ff = 16, 
+                 factor = 1, output_attention = False, enc_layers = 2, activation = 'gelu'):
+        super(Transformer, self).__init__()
+        self.output_attention = output_attention
+
+        # Embedding
+        self.enc_embedding = DataEmbedding_wo_temp(enc_in,d_model)
+        # self.dec_embedding = DataEmbedding_wo_temp(configs.dec_in, configs.d_model)
+
+        # Encoder
+        self.encoder = Encoder(
+            [
+                EncoderLayer(
+                    AttentionLayer(
+                        FullAttention(False, factor, attention_dropout=dropout,
+                                      output_attention=output_attention), d_model, n_heads),
+                    d_model,
+                    d_ff,
+                    dropout = dropout,
+                    activation = activation
+                ) for l in range(enc_layers)
+            ],
+            norm_layer=torch.nn.LayerNorm(d_model)
+        )
+        self.GMP = torch.nn.AdaptiveMaxPool1d(1, return_indices=False)
+
+        # # Decoder
+        # self.decoder = Decoder(
+        #     [
+        #         DecoderLayer(
+        #             AttentionLayer(
+        #                 FullAttention(True, configs.factor, attention_dropout=configs.dropout, output_attention=False),
+        #                 configs.d_model, configs.n_heads),
+        #             AttentionLayer(
+        #                 FullAttention(False, configs.factor, attention_dropout=configs.dropout, output_attention=False),
+        #                 configs.d_model, configs.n_heads),
+        #             configs.d_model,
+        #             configs.d_ff,
+        #             dropout=configs.dropout,
+        #             activation=configs.activation,
+        #         )
+        #         for l in range(configs.d_layers)
+        #     ],
+        #     norm_layer=torch.nn.LayerNorm(configs.d_model),
+        #     projection=nn.Linear(configs.d_model, configs.c_out, bias=True)
+        # )
+
+    def forward(self, x_enc, x_mark_enc=None, x_dec=None, x_mark_dec=None,
+                enc_self_mask=None, dec_self_mask=None, dec_enc_mask=None):
+
+        enc_out = self.enc_embedding(x_enc, x_mark_enc)
+        # print(enc_out.shape) 
+        out, attns = self.encoder(enc_out, attn_mask=enc_self_mask)
+        # out [bs, num_node, embedding_dim]
+        out = self.GMP(out.permute(0,2,1)).squeeze(-1)
+        # dec_out = self.dec_embedding(x_dec, x_mark_dec)
+        # dec_out = self.decoder(dec_out, enc_out, x_mask=dec_self_mask, cross_mask=dec_enc_mask)
+
+        # if self.output_attention:
+        #     return dec_out[:, -self.pred_len:, :], attns
+        # else:
+        return out  # [B, L, D]
+
+if __name__ == '__main__':
+    net = Transformer()
+    a = torch.randn((1,78,2))
+    print(net(a).shape)

scripts/fake_data.py

@@ -0,0 +1,123 @@
+import math
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from onescience.utils.YParams import YParams
+
+
+def set_default_data_env():
+    os.environ.setdefault("ONESCIENCE_BENO_DATA_DIR", str(PROJECT_ROOT / "data"))
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def load_config():
+    set_default_data_env()
+    cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
+    cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
+    cfg.datapipe.source.cache_dir = str(resolve_path(cfg.datapipe.source.cache_dir))
+    cfg.fake_data.root_dir = str(resolve_path(cfg.fake_data.root_dir))
+    return cfg
+
+
+def square_boundary(points, boundary_points):
+    bottom = [(x, 0.0) for x in points]
+    top = [(x, 1.0) for x in points]
+    left = [(0.0, y) for y in points[1:-1]]
+    right = [(1.0, y) for y in points[1:-1]]
+    boundary = np.array(bottom + top + left + right, dtype=np.float64)
+    if boundary.shape[0] == 0:
+        raise ValueError("resolution is too small to build boundary points")
+    repeats = int(math.ceil(boundary_points / boundary.shape[0]))
+    return np.tile(boundary, (repeats, 1))[:boundary_points]
+
+
+def make_beno_arrays(cfg):
+    data_cfg = cfg.datapipe.data
+    fake_cfg = cfg.fake_data
+    total_samples = int(data_cfg.ntrain) + int(data_cfg.ntest)
+    resolution = int(data_cfg.resolution)
+    boundary_points = int(fake_cfg.boundary_points)
+    rng = np.random.default_rng(int(fake_cfg.seed))
+
+    axis = np.linspace(0.0, 1.0, resolution, dtype=np.float64)
+    yy, xx = np.meshgrid(axis, axis, indexing="ij")
+    coords = np.stack([xx.reshape(-1), yy.reshape(-1)], axis=-1)
+
+    cell_state = np.zeros((resolution, resolution), dtype=np.float64)
+    cell_state[0, :] = 1.0
+    cell_state[-1, :] = 1.0
+    cell_state[:, 0] = 1.0
+    cell_state[:, -1] = 1.0
+    cell_state = cell_state.reshape(-1)
+
+    boundary_xy = square_boundary(axis, boundary_points)
+
+    rhs = np.zeros((total_samples, resolution * resolution, 4), dtype=np.float64)
+    sol = np.zeros((total_samples, resolution * resolution, 1), dtype=np.float64)
+    bc = np.zeros((total_samples, boundary_points, 4), dtype=np.float64)
+
+    for sample_idx in range(total_samples):
+        phase = 0.25 * sample_idx
+        amplitude = 1.0 + 0.1 * sample_idx
+        forcing = (
+            amplitude * np.sin(np.pi * xx + phase) * np.sin(np.pi * yy)
+            + 0.15 * np.cos(2.0 * np.pi * xx) * np.sin(np.pi * yy + phase)
+        )
+        solution = 0.25 * forcing + 0.05 * (xx + yy)
+        if fake_cfg.noise_std:
+            solution += float(fake_cfg.noise_std) * rng.standard_normal(solution.shape)
+
+        boundary_value = (
+            0.05 * sample_idx
+            + 0.1 * boundary_xy[:, 0]
+            - 0.08 * boundary_xy[:, 1]
+        )
+
+        rhs[sample_idx, :, 0:2] = coords
+        rhs[sample_idx, :, 2] = forcing.reshape(-1)
+        rhs[sample_idx, :, 3] = cell_state
+        sol[sample_idx, :, 0] = solution.reshape(-1)
+        bc[sample_idx, :, 0:2] = boundary_xy
+        bc[sample_idx, :, 2] = boundary_value
+
+    return rhs, sol, bc
+
+
+def clear_matching_cache(cfg):
+    cache_dir = Path(cfg.datapipe.source.cache_dir)
+    prefix = cfg.datapipe.source.file_prefix
+    for split, count in (("train", cfg.datapipe.data.ntrain), ("test", cfg.datapipe.data.ntest)):
+        cache_file = cache_dir / f"cached_{prefix}_{split}_{count}.pt"
+        if cache_file.exists():
+            cache_file.unlink()
+
+
+def main():
+    cfg = load_config()
+    data_dir = Path(cfg.datapipe.source.data_dir)
+    data_dir.mkdir(parents=True, exist_ok=True)
+
+    rhs, sol, bc = make_beno_arrays(cfg)
+    prefix = cfg.datapipe.source.file_prefix
+    np.save(data_dir / f"RHS_{prefix}_all.npy", rhs)
+    np.save(data_dir / f"SOL_{prefix}_all.npy", sol)
+    np.save(data_dir / f"BC_{prefix}_all.npy", bc)
+    clear_matching_cache(cfg)
+
+    print(f"Fake BENO data written to {data_dir}")
+    print(f"RHS shape={rhs.shape}, SOL shape={sol.shape}, BC shape={bc.shape}")
+
+
+if __name__ == "__main__":
+    main()

scripts/inference.py

@@ -0,0 +1,204 @@
+import json
+import logging
+import os
+import sys
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model import HeteroGNS
+from onescience.datapipes.cfd import BENODatapipe
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+from onescience.utils.beno.util import record_data, to_np_array
+from onescience.utils.beno.utilities import LpLoss, plot_data
+
+
+logging.basicConfig(level=logging.INFO)
+LOGGER = logging.getLogger("BENO_Inference")
+
+
+def set_default_data_env():
+    os.environ.setdefault("ONESCIENCE_BENO_DATA_DIR", str(PROJECT_ROOT / "data"))
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def load_config():
+    set_default_data_env()
+    cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
+    cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
+    cfg.datapipe.source.cache_dir = str(resolve_path(cfg.datapipe.source.cache_dir))
+    cfg.inference.checkpoint_dir = str(resolve_path(cfg.inference.checkpoint_dir))
+    cfg.inference.result_dir = str(resolve_path(cfg.inference.result_dir))
+    cfg.inference.picture_dir = str(resolve_path(cfg.inference.picture_dir))
+    cfg.inference.metrics_path = str(resolve_path(cfg.inference.metrics_path))
+    if cfg.inference.checkpoint_path:
+        cfg.inference.checkpoint_path = str(resolve_path(cfg.inference.checkpoint_path))
+    return cfg
+
+
+def activation_from_name(name):
+    name = str(name).lower()
+    if name == "relu":
+        return nn.ReLU
+    if name == "elu":
+        return nn.ELU
+    if name == "leakyrelu":
+        return nn.LeakyReLU
+    return nn.SiLU
+
+
+def build_model(model_cfg):
+    return HeteroGNS(
+        nnode_in_features=model_cfg.nnode_in_features,
+        nnode_out_features=model_cfg.nnode_out_features,
+        nedge_in_features=model_cfg.nedge_in_features,
+        latent_dim=model_cfg.get("latent_dim", model_cfg.get("width", 128)),
+        nmessage_passing_steps=model_cfg.get("nmessage_passing_steps", 10),
+        nmlp_layers=model_cfg.nmlp_layers,
+        mlp_hidden_dim=model_cfg.get("mlp_hidden_dim", model_cfg.get("width", 128)),
+        activation=activation_from_name(model_cfg.act),
+        boundary_dim=model_cfg.boundary_dim,
+        trans_layer=model_cfg.trans_layer,
+    )
+
+
+def select_device(device_name, dist):
+    device_name = str(device_name).lower()
+    if device_name == "cpu":
+        return torch.device("cpu")
+    if device_name == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError("inference.device is cuda, but CUDA is not available.")
+        return torch.device(f"cuda:{dist.local_rank}")
+    return dist.device
+
+
+def find_checkpoint(cfg):
+    if cfg.inference.checkpoint_path:
+        path = Path(cfg.inference.checkpoint_path)
+        return path if path.exists() else None
+    checkpoints = sorted(Path(cfg.inference.checkpoint_dir).glob("model_epoch_*.pt"))
+    return checkpoints[-1] if checkpoints else None
+
+
+def main():
+    cfg = load_config()
+    DistributedManager.initialize()
+    dist = DistributedManager()
+    device = select_device(cfg.inference.get("device", "auto"), dist)
+
+    datapipe = BENODatapipe(cfg, distributed=(dist.world_size > 1))
+    test_loader, _ = datapipe.test_dataloader()
+    u_normalizer = datapipe.u_normalizer.to(device)
+    a_normalizer = datapipe.a_normalizer.to(device)
+    resolution = int(cfg.datapipe.data.resolution)
+
+    model = build_model(cfg.model).to(device)
+    checkpoint = find_checkpoint(cfg)
+    if checkpoint is not None:
+        LOGGER.info("Loading checkpoint: %s", checkpoint)
+        state_dict = torch.load(checkpoint, map_location=device)
+        model.load_state_dict(state_dict)
+    else:
+        LOGGER.warning("No checkpoint found. Running with random weights.")
+    model.eval()
+
+    result_dir = Path(cfg.inference.result_dir)
+    picture_dir = Path(cfg.inference.picture_dir)
+    result_dir.mkdir(parents=True, exist_ok=True)
+    picture_dir.mkdir(parents=True, exist_ok=True)
+
+    myloss = LpLoss(size_average=False)
+    analysis_record = {}
+    out_all = np.array([])
+    label_all = np.array([])
+    a_ori_all = np.array([])
+    mask_all = np.array([])
+    grid_all = np.array([])
+
+    with torch.no_grad():
+        for sample_id, data in enumerate(test_loader):
+            data = data.to(device)
+            out_indomain = model(data)
+
+            full_out = torch.zeros((resolution * resolution, 1), device=device)
+            full_label = torch.zeros((resolution * resolution), device=device)
+            full_input_a = torch.zeros((resolution * resolution, 10), device=device)
+            full_grid = torch.zeros((resolution * resolution, 2), device=device)
+            indices = data["G1"].sample_idx
+
+            full_out[indices] = out_indomain
+            full_label[indices] = data["G1+2"].y
+            full_input_a[indices, :] = data["G1"].x
+            full_grid[indices, :] = data["G1"].x[:, :2]
+
+            pred_decoded = u_normalizer.decode(full_out.view(1, -1))
+            a_decoded = a_normalizer.decode(full_input_a[:, 2].view(1, -1))
+            label_reshaped = full_label.view(1, -1)
+
+            cell_state_full = torch.zeros((1, resolution * resolution), device=device)
+            cell_state_full[0, :] = data["G1"].cell_state
+
+            l2_item = myloss(pred_decoded, label_reshaped).item()
+            mae_item = nn.L1Loss()(pred_decoded, label_reshaped).item()
+            record_data(analysis_record, [l2_item, mae_item], ["L2", "MAE"])
+
+            np.savez(
+                result_dir / f"sample_{sample_id:04d}.npz",
+                predict=to_np_array(pred_decoded).reshape(resolution, resolution),
+                label=to_np_array(label_reshaped).reshape(resolution, resolution),
+                forcing=to_np_array(a_decoded).reshape(resolution, resolution),
+                mask=to_np_array(cell_state_full).reshape(resolution, resolution),
+            )
+
+            out_all = np.append(out_all, to_np_array(pred_decoded))
+            label_all = np.append(label_all, to_np_array(label_reshaped))
+            a_ori_all = np.append(a_ori_all, to_np_array(a_decoded))
+            mask_all = np.append(mask_all, to_np_array(cell_state_full))
+            grid_all = np.append(grid_all, to_np_array(full_grid.unsqueeze(0)))
+
+    metrics = {
+        "mean_l2": float(np.mean(analysis_record["L2"])),
+        "std_l2": float(np.std(analysis_record["L2"])),
+        "mean_mae": float(np.mean(analysis_record["MAE"])),
+        "num_samples": int(len(analysis_record["L2"])),
+    }
+    metrics_path = Path(cfg.inference.metrics_path)
+    metrics_path.parent.mkdir(parents=True, exist_ok=True)
+    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
+    LOGGER.info("Metrics: %s", metrics)
+
+    plot_samples = min(int(cfg.inference.num_visualize), metrics["num_samples"])
+    if plot_samples >= 2:
+        plot_path = picture_dir / "forcing_solution_comparison.png"
+        plot_data(
+            predict_term=out_all,
+            true_term=label_all,
+            forcing_term=a_ori_all,
+            forcing_mask=mask_all,
+            grid_info=grid_all,
+            resolution=resolution,
+            num_samples=plot_samples,
+            interpolation="bilinear",
+            save_path=str(plot_path),
+        )
+        LOGGER.info("Plot saved to %s", plot_path)
+    elif plot_samples == 1:
+        LOGGER.info("Skipping plot_data for one sample because the runtime helper expects a 2D axes grid.")
+
+    dist.cleanup()
+
+
+if __name__ == "__main__":
+    main()

scripts/result.py

@@ -0,0 +1,19 @@
+import json
+from pathlib import Path
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+METRICS_PATH = PROJECT_ROOT / "result" / "metrics.json"
+
+
+def main():
+    if not METRICS_PATH.exists():
+        raise FileNotFoundError(f"Missing metrics file: {METRICS_PATH}. Run scripts/inference.py first.")
+    metrics = json.loads(METRICS_PATH.read_text(encoding="utf-8"))
+    print("BENO inference metrics")
+    for key, value in metrics.items():
+        print(f"{key}: {value}")
+
+
+if __name__ == "__main__":
+    main()

scripts/train.py

@@ -0,0 +1,210 @@
+import os
+import random
+import sys
+from pathlib import Path
+from timeit import default_timer
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.nn.parallel import DistributedDataParallel as DDP
+from tqdm import tqdm
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+sys.path.insert(0, str(PROJECT_ROOT))
+
+from model import HeteroGNS
+from onescience.datapipes.cfd import BENODatapipe
+from onescience.distributed.manager import DistributedManager
+from onescience.utils.YParams import YParams
+from onescience.utils.beno.utilities import LpLoss
+
+
+def set_default_data_env():
+    os.environ.setdefault("ONESCIENCE_BENO_DATA_DIR", str(PROJECT_ROOT / "data"))
+
+
+def resolve_path(path_value):
+    path = Path(path_value)
+    return path if path.is_absolute() else PROJECT_ROOT / path
+
+
+def load_config():
+    set_default_data_env()
+    cfg = YParams(str(PROJECT_ROOT / "conf" / "config.yaml"), "root")
+    cfg.datapipe.source.data_dir = str(resolve_path(cfg.datapipe.source.data_dir))
+    cfg.datapipe.source.cache_dir = str(resolve_path(cfg.datapipe.source.cache_dir))
+    cfg.training.output_dir = str(resolve_path(cfg.training.output_dir))
+    return cfg
+
+
+def seed_everything(seed):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def activation_from_name(name):
+    name = str(name).lower()
+    if name == "relu":
+        return nn.ReLU
+    if name == "elu":
+        return nn.ELU
+    if name == "leakyrelu":
+        return nn.LeakyReLU
+    return nn.SiLU
+
+
+def build_model(model_cfg):
+    return HeteroGNS(
+        nnode_in_features=model_cfg.nnode_in_features,
+        nnode_out_features=model_cfg.nnode_out_features,
+        nedge_in_features=model_cfg.nedge_in_features,
+        latent_dim=model_cfg.get("latent_dim", model_cfg.get("width", 128)),
+        nmessage_passing_steps=model_cfg.get("nmessage_passing_steps", 10),
+        nmlp_layers=model_cfg.nmlp_layers,
+        mlp_hidden_dim=model_cfg.get("mlp_hidden_dim", model_cfg.get("width", 128)),
+        activation=activation_from_name(model_cfg.act),
+        boundary_dim=model_cfg.boundary_dim,
+        trans_layer=model_cfg.trans_layer,
+    )
+
+
+def select_device(device_name, dist):
+    device_name = str(device_name).lower()
+    if device_name == "cpu":
+        return torch.device("cpu")
+    if device_name == "cuda":
+        if not torch.cuda.is_available():
+            raise RuntimeError("training.device is cuda, but CUDA is not available.")
+        return torch.device(f"cuda:{dist.local_rank}")
+    return dist.device
+
+
+def reduce_scalar(value, device, dist):
+    tensor = torch.tensor(value, device=device, dtype=torch.float32)
+    if dist.world_size > 1:
+        torch.distributed.all_reduce(tensor)
+        tensor /= dist.world_size
+    return tensor.item()
+
+
+def evaluate(model, test_loader, device, u_normalizer, myloss, dist):
+    model.eval()
+    total_l2 = 0.0
+    with torch.no_grad():
+        for batch in test_loader:
+            batch = batch.to(device)
+            out = model(batch)
+            pred = u_normalizer.decode(
+                out.view(batch.num_graphs, -1),
+                sample_idx=batch["G1"].sample_idx.view(batch.num_graphs, -1),
+            )
+            total_l2 += myloss(pred, batch["G1+2"].y.view(batch.num_graphs, -1)).item()
+    total_l2 = reduce_scalar(total_l2, device, dist)
+    return total_l2
+
+
+def main():
+    cfg = load_config()
+    seed_everything(int(cfg.training.seed))
+
+    DistributedManager.initialize()
+    dist = DistributedManager()
+    device = select_device(cfg.training.get("device", "auto"), dist)
+
+    output_dir = Path(cfg.training.output_dir)
+    if dist.rank == 0:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        print(f"Config: {PROJECT_ROOT / 'conf' / 'config.yaml'}")
+        print(f"Data: {cfg.datapipe.source.data_dir}")
+        print(f"Checkpoint directory: {output_dir}")
+
+    datapipe = BENODatapipe(cfg, distributed=(dist.world_size > 1))
+    train_loader, train_sampler = datapipe.train_dataloader()
+    test_loader, _ = datapipe.test_dataloader()
+    if len(train_loader) == 0:
+        raise RuntimeError("Training loader is empty. Check ntrain and batch_size.")
+
+    u_normalizer = datapipe.u_normalizer.to(device)
+    model = build_model(cfg.model).to(device)
+    if dist.world_size > 1:
+        device_ids = [dist.local_rank] if device.type == "cuda" else None
+        model = DDP(model, device_ids=device_ids)
+
+    optimizer = torch.optim.Adam(
+        model.parameters(),
+        lr=cfg.training.optimizer.lr,
+        weight_decay=cfg.training.optimizer.weight_decay,
+    )
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(
+        optimizer,
+        T_0=cfg.training.scheduler.T_0,
+        T_mult=cfg.training.scheduler.T_mult,
+    )
+    myloss = LpLoss(size_average=False)
+
+    for epoch in range(int(cfg.training.epochs)):
+        if train_sampler:
+            train_sampler.set_epoch(epoch)
+
+        model.train()
+        train_mse = 0.0
+        train_l2 = 0.0
+        batches = 0
+        start = default_timer()
+        iterator = tqdm(train_loader, desc=f"Epoch {epoch}", disable=(dist.rank != 0))
+
+        for batch in iterator:
+            batch = batch.to(device)
+            optimizer.zero_grad(set_to_none=True)
+            out = model(batch)
+            loss = F.mse_loss(out.view(-1, 1), batch["G1+2"].y.view(-1, 1))
+            loss.backward()
+            optimizer.step()
+
+            with torch.no_grad():
+                pred_denorm = u_normalizer.decode(
+                    out.view(batch.num_graphs, -1),
+                    sample_idx=batch["G1"].sample_idx.view(batch.num_graphs, -1),
+                )
+                target_denorm = u_normalizer.decode(
+                    batch["G1+2"].y.view(batch.num_graphs, -1),
+                    sample_idx=batch["G1"].sample_idx.view(batch.num_graphs, -1),
+                )
+                l2 = myloss(pred_denorm, target_denorm)
+
+            train_mse += loss.item()
+            train_l2 += l2.item()
+            batches += 1
+            if dist.rank == 0:
+                iterator.set_postfix({"mse": f"{loss.item():.2e}", "l2": f"{l2.item():.2e}"})
+
+        scheduler.step()
+        train_mse = reduce_scalar(train_mse / batches, device, dist)
+        train_l2 = reduce_scalar(train_l2 / datapipe.train_dataset.ntrain, device, dist)
+        test_l2 = evaluate(model, test_loader, device, u_normalizer, myloss, dist)
+        test_l2 /= datapipe.test_dataset.ntest
+
+        if dist.rank == 0:
+            print(
+                f"Epoch {epoch:03d} | Train MSE: {train_mse:.6f} | "
+                f"Train L2: {train_l2:.6f} | Test L2: {test_l2:.6f} | "
+                f"Time: {default_timer() - start:.1f}s"
+            )
+            if (epoch + 1) % int(cfg.training.save_period) == 0:
+                model_to_save = model.module if hasattr(model, "module") else model
+                checkpoint_name = cfg.training.checkpoint_name.format(epoch=epoch)
+                ckpt_path = output_dir / checkpoint_name
+                torch.save(model_to_save.state_dict(), ckpt_path)
+                print(f"Saved checkpoint to {ckpt_path}")
+
+    dist.cleanup()
+
+
+if __name__ == "__main__":
+    main()