Upload folder using huggingface_hub

README.md

@@ -0,0 +1,126 @@
+---
+license: apache-2.0
+language:
+- en
+tags:
+- OneScience
+- PDE
+- CFD
+- neural-operator
+- PDEBench
+frameworks: PyTorch
+
+datasets:
+ - OneScience/pdenneval
+---
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">PDENNEval</span>
+  </strong>
+</p>
+
+# Model Overview
+PDENNEval is a comprehensive benchmark developed by researchers at Sun Yat-sen University for evaluating neural-network methods for solving partial differential equations. It enables systematic comparison of function-learning and operator-learning approaches across a broad range of PDE tasks.
+
+Paper: [PDENNEval: A Comprehensive Evaluation of Neural Network Methods for Solving PDEs](https://www.ijcai.org/proceedings/2024/0573.pdf)
+
+# Model Description
+PDENNEval provides a unified framework for evaluating neural-network PDE solvers. It contains 16 datasets generated with high-accuracy conventional scientific computing methods, covering 19 PDE problem classes from fluid dynamics, materials science, finance, electromagnetics, and other domains. The benchmark evaluates the solution accuracy, computational efficiency, and robustness of 12 neural-network methods.
+
+## Use Cases
+
+| Use Case | Description |
+| --- | --- |
+| PDE neural-operator validation | Complete a minimal FNO training and inference workflow using PDEBench-style HDF5 data |
+| Neural PDE method development | Extend the model collection in `model/`, which already includes FNO, DeepONet, PINO-FNO, UNO, MPNN, and UNet definitions |
+| ModelScope/OneCode execution | Download the standalone model package, install its dependencies, and run the provided scripts |
+| Rapid pipeline checks | Use `fake_data.py` and `smoke_models.py` to verify the data, models, and 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/PDENNEval --local_dir ./PDENNEval
+cd PDENNEval
+```
+
+### 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 `pdenneval` dataset for training. Download it with the command below and verify that the data path in `conf/config.yaml` is configured correctly:
+
+```bash
+modelscope download --dataset OneScience/pdenneval  --local_dir ./data
+```
+
+### Training
+
+```bash
+python scripts/train.py
+```
+
+### Model Weights
+This repository will provide weights trained on the `pdenneval` dataset in the `weights/` directory. The weights will be uploaded soon.
+
+### Inference
+
+```bash
+python scripts/inference.py
+```
+
+By default, inference loads `weight/best_model.pt` and writes `.npz` results to `result/output/`.
+
+### Evaluation and Visualization
+
+```bash
+python scripts/result.py
+```
+
+# Official OneScience Resources
+
+| Platform | OneScience Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citations and License
+
+- Original PDENNEval paper: [PDENNEval: A Comprehensive Evaluation of Neural Network Methods for Solving PDEs](https://www.ijcai.org/proceedings/2024/0573.pdf).
+- This repository retains the relevant source and attribution notices. Follow all applicable license requirements when using, modifying, or distributing its contents.

conf/config.yaml

@@ -0,0 +1,52 @@
+fno_config:
+  datapipe:
+    source:
+      data_dir: "./data"
+      file_name: "2D_DarcyFlow_beta0.1_Train.hdf5"
+    verbose: true
+    data:
+      single_file: true
+      reduced_resolution: 1
+      reduced_resolution_t: 1
+      reduced_batch: 1
+      initial_step: 1
+      test_ratio: 0.1
+      pde_name: "2D_Darcy_Flow"
+    dataloader:
+      batch_size: 2
+      num_workers: 0
+      pin_memory: false
+
+  model:
+    name: "FNO"
+    num_channels: 1
+    modes: 4
+    width: 8
+
+  training:
+    output_dir: "./weight"
+    if_training: true
+    continue_training: false
+    model_path: null
+    save_period: 1
+    seed: 0
+    training_type: "single"
+    t_train: 2
+    epochs: 1
+    optimizer:
+      name: "Adam"
+      lr: 1.0e-3
+      weight_decay: 1.0e-4
+    scheduler:
+      name: "StepLR"
+      step_size: 100
+      gamma: 0.5
+
+  inference:
+    checkpoint: "./weight/best_model.pt"
+    output_dir: "./result/output"
+    max_batches: 2
+
+  result:
+    prediction_dir: "./result/output"
+    output_dir: "./result"

configuration.json

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

model/__init__.py

@@ -0,0 +1,3 @@
+from .fno import FNO1d, FNO2d, FNO3d, FNO_maxwell
+
+__all__ = ["FNO1d", "FNO2d", "FNO3d", "FNO_maxwell"]

model/deeponet.py

@@ -0,0 +1,257 @@
+# coding=utf-8
+import torch
+import torch.nn as nn
+from onescience.utils.pdenneval.deeponet_utils import _get_act, _get_initializer
+
+
+class MLP(nn.Module):
+    """Fully-connected neural network."""
+
+    def __init__(self, layer_sizes, activation, kernel_initializer):
+        super().__init__()
+        self.activation = _get_act(activation)
+        initializer = _get_initializer(kernel_initializer)
+        initializer_zero = _get_initializer("zeros")
+
+        self.linears = torch.nn.ModuleList()
+        for i in range(1, len(layer_sizes)):
+            self.linears.append(
+                torch.nn.Linear(
+                    layer_sizes[i - 1], layer_sizes[i], dtype=torch.float32
+                )
+            )
+            initializer(self.linears[-1].weight)
+            initializer_zero(self.linears[-1].bias)
+
+    def forward(self, inputs):
+        x = inputs
+        for linear in self.linears[:-1]:
+            x = self.activation(linear(x))
+        x = self.linears[-1](x)
+        return x
+
+class Modified_MLP(nn.Module):
+    def __init__(self, layer_sizes, activation, kernel_initializer) -> None:
+        super().__init__()
+        self.activation = _get_act(activation)
+        initializer = _get_initializer(kernel_initializer)
+        initializer_zero = _get_initializer("zeros")
+        self.linears = torch.nn.ModuleList()
+        for i in range(1, len(layer_sizes)):
+            self.linears.append(
+                torch.nn.Linear(
+                    layer_sizes[i - 1], layer_sizes[i], dtype=torch.float32
+                )
+            )
+            initializer(self.linears[-1].weight)
+            initializer_zero(self.linears[-1].bias)
+        self.linear1=torch.nn.Linear(layer_sizes[0], layer_sizes[1], dtype=torch.float32)
+        self.linear2=torch.nn.Linear(layer_sizes[0], layer_sizes[1], dtype=torch.float32)
+        initializer(self.linear1.weight),initializer(self.linear2.weight)
+        initializer_zero(self.linear1.bias),initializer_zero(self.linear2.bias)
+    def forward(self, inputs):
+        U = self.activation(self.linear1(inputs))
+        V = self.activation(self.linear2(inputs))
+        for linear in self.linears[:-1]:
+            outputs=torch.sigmoid(linear(inputs))
+            inputs= outputs*U + (1-outputs)* V
+        outputs = self.linears[-1](inputs)
+        return outputs
+
+
+class DeepONet(nn.Module):
+    """Deep operator network.
+
+    Args:
+        layer_sizes_branch: A list of integers as the width of a fully connected network,
+            or `(dim, f)` where `dim` is the input dimension and `f` is a network
+            function. The width of the last layer in the branch and trunk net should be
+            equal.
+        layer_sizes_trunk (list): A list of integers as the width of a fully connected
+            network.
+        activation: If `activation` is a ``string``, then the same activation is used in
+            both trunk and branch nets. If `activation` is a ``dict``, then the trunk
+            net uses the activation `activation["trunk"]`, and the branch net uses
+            `activation["branch"]`.
+    """
+
+    def __init__(
+        self,
+        layer_sizes_branch,
+        layer_sizes_trunk,
+        activation,
+        kernel_initializer,
+    ):
+        super().__init__()
+        if isinstance(activation, dict):
+            activation_branch = _get_act(activation["branch"])
+            self.activation_trunk = _get_act(activation["trunk"])
+        else:
+            activation_branch = self.activation_trunk = _get_act(activation)
+        if callable(layer_sizes_branch[0]):
+            # User-defined network
+            self.branch = layer_sizes_branch[0]
+        else:
+            # Fully connected network
+            self.branch = MLP(layer_sizes_branch, activation_branch, kernel_initializer)
+        self.trunk = MLP(layer_sizes_trunk, self.activation_trunk, kernel_initializer)
+        self.b = torch.nn.parameter.Parameter(torch.tensor(0.0))
+
+    def forward(self, inputs):
+        x_func = inputs[0]
+        x_loc = inputs[1]
+        # Branch net to encode the input function
+        x_func = self.branch(x_func)
+        # Trunk net to encode the domain of the output function
+        x_loc = self.activation_trunk(self.trunk(x_loc))
+        # Dot product
+        if x_func.shape[-1] != x_loc.shape[-1]:
+            raise AssertionError(
+                "Output sizes of branch net and trunk net do not match."
+            )
+        x = torch.einsum("bi,bi->b", x_func, x_loc)
+        x = torch.unsqueeze(x, 1)
+        # Add bias
+        x += self.b
+        return x
+
+class DeepONetCartesianProd(nn.Module):
+    """Deep operator network for dataset in the format of Cartesian product.
+
+    Args:
+        layer_sizes_branch: A list of integers as the width of a fully connected network,
+            or `(dim, f)` where `dim` is the input dimension and `f` is a network
+            function. The width of the last layer in the branch and trunk net should be
+            equal.
+        layer_sizes_trunk (list): A list of integers as the width of a fully connected
+            network.
+        activation: If `activation` is a ``string``, then the same activation is used in
+            both trunk and branch nets. If `activation` is a ``dict``, then the trunk
+            net uses the activation `activation["trunk"]`, and the branch net uses
+            `activation["branch"]`.
+    """
+
+    def __init__(
+        self,
+        layer_sizes_branch,
+        layer_sizes_trunk,
+        activation,
+        kernel_initializer,
+        base_model = "MLP"    # or Modified_MLP 
+    ):
+        super().__init__()
+        if isinstance(activation, dict):
+            activation_branch = _get_act(activation["branch"])
+            self.activation_trunk = _get_act(activation["trunk"])
+        else:
+            activation_branch = self.activation_trunk = _get_act(activation)
+        base_model= MLP if base_model=="MLP" else Modified_MLP
+        if callable(layer_sizes_branch[0]):
+            # User-defined network
+            self.branch = layer_sizes_branch[0]
+        else:
+            self.branch = base_model(layer_sizes_branch, activation_branch, kernel_initializer)
+        self.trunk = base_model(layer_sizes_trunk, self.activation_trunk, kernel_initializer)
+        self.b = torch.nn.parameter.Parameter(torch.tensor(0.0))
+
+    def forward(self, inputs):
+        x_func = inputs[0]
+        x_loc = inputs[1]
+        # Branch net to encode the input function
+        x_func = self.branch(x_func)
+        # Trunk net to encode the domain of the output function
+        x_loc = self.activation_trunk(self.trunk(x_loc))
+        # Dot product
+        if x_func.shape[-1] != x_loc.shape[-1]:
+            raise AssertionError(
+                "Output sizes of branch net and trunk net do not match."
+            )
+        x = torch.einsum("bi,ni->bn", x_func, x_loc)
+        # Add bias
+        x += self.b
+        return x
+
+class DeepONetCartesianProd2D(DeepONetCartesianProd): 
+    #   For multiple outputs, we choose the second approach mentioned in "https://arxiv.org/abs/2111.05512", i.e. split 
+    # the output of both the branch and the trunk into n groups, and the k-th groups outputs the k-th solution.
+    def __init__(self,
+        size: int,
+        in_channel_branch: int,
+        query_dim: int ,
+        out_channel: int,
+        activation: str = "relu",
+        kernel_initializer: str = "Glorot normal",
+        base_model = "MLP"):
+        layer_sizes_branch = [in_channel_branch*size**2]+[128]*4+[128*out_channel]
+        layer_sizes_trunk= [query_dim]+[128]*4+[128*out_channel]
+        super().__init__(layer_sizes_branch,layer_sizes_trunk,activation,kernel_initializer,base_model)
+        self.out_channel = out_channel
+        self.query_dim=query_dim
+        self.b = torch.nn.parameter.Parameter(torch.zeros(out_channel,dtype=torch.float32))
+
+    def forward(self, inputs):
+        x_func = inputs[0]
+        x_loc = inputs[1]  
+        batchsize=x_func.shape[0]
+        x_func = x_func.reshape([batchsize,-1]) 
+        grid_shape = x_loc.shape[:-1]
+        x_loc = x_loc.reshape([-1,self.query_dim])  #(num_point, query_dim)
+        num_points=x_loc.shape[0]
+        # Branch net to encode the input function
+        x_func = self.branch(x_func.reshape([batchsize,-1]))
+        # Trunk net to encode the domain of the output function
+        x_loc = self.activation_trunk(self.trunk(x_loc))
+        # Dot product
+        if x_func.shape[-1] != x_loc.shape[-1]:
+            raise AssertionError(
+                "Output sizes of branch net and trunk net do not match."
+            )
+        x_func = x_func.reshape([batchsize,self.out_channel,-1])
+        x_loc = x_loc.reshape([num_points,self.out_channel,-1])
+        x = torch.einsum("bci,nci->bnc", x_func, x_loc)
+        # Add bias
+        x += self.b
+        return x.reshape([-1,*grid_shape,self.out_channel])
+
+class DeepONetCartesianProd1D(DeepONetCartesianProd):
+    #   For multiple outputs, we choose the second approach mentioned in "https://arxiv.org/abs/2111.05512", i.e. split 
+    # the output of both the branch and the trunk into n groups, and the k-th groups outputs the k-th solution.
+    def __init__(self,
+        size :int,
+        in_channel_branch: int,
+        query_dim: int ,
+        out_channel: int,
+        activation: str = "relu",
+        kernel_initializer: str = "Glorot normal",
+        base_model="MLP"):
+        layer_sizes_branch= [in_channel_branch*size]+[128]*4+[128*out_channel]
+        layer_sizes_trunk= [query_dim]+[128]*3+[128*out_channel]
+        super().__init__(layer_sizes_branch,layer_sizes_trunk,activation,kernel_initializer,base_model)
+        self.out_channel = out_channel
+        self.b = torch.nn.parameter.Parameter(torch.zeros(out_channel,dtype=torch.float32))
+        self.query_dim=query_dim
+
+    def forward(self, inputs):
+        x_func = inputs[0]
+        x_loc = inputs[1]
+        grid_shape = x_loc.shape[:-1]
+        x_loc = x_loc.reshape([-1,self.query_dim])  #(num_point, query_dim)
+        num_points=x_loc.shape[0]
+        batchsize=x_func.shape[0]
+        x_func = x_func.reshape([batchsize,-1])
+        # Branch net to encode the input function
+        
+        x_func = self.branch(x_func)
+        # Trunk net to encode the domain of the output function
+        x_loc = self.activation_trunk(self.trunk(x_loc))
+        # Dot product
+        if x_func.shape[-1] != x_loc.shape[-1]:
+            raise AssertionError(
+                "Output sizes of branch net and trunk net do not match."
+            )
+        x_func = x_func.reshape([batchsize,self.out_channel,-1])
+        x_loc = x_loc.reshape([num_points,self.out_channel,-1])
+        x = torch.einsum("bci,nci->bnc", x_func, x_loc)
+        # Add bias
+        x += self.b
+        return x.reshape([-1,*grid_shape,self.out_channel])

model/fno.py

@@ -0,0 +1,412 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import torch.nn.functional as F
+
+
+class SpectralConv1d(nn.Module):
+    def __init__(self, in_channels, out_channels, modes1):
+        super(SpectralConv1d, self).__init__()
+
+        """
+        1D Fourier layer. It does FFT, linear transform, and Inverse FFT.    
+        """
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.modes1 = modes1  #Number of Fourier modes to multiply, at most floor(N/2) + 1
+
+        self.scale = (1 / (in_channels*out_channels))
+        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat))
+
+    # Complex multiplication
+    def compl_mul1d(self, input, weights):
+        # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
+        return torch.einsum("bix,iox->box", input, weights)
+
+    def forward(self, x):
+        batchsize = x.shape[0]
+        #Compute Fourier coeffcients up to factor of e^(- something constant)
+        x_ft = torch.fft.rfft(x)
+
+        # Multiply relevant Fourier modes 
+        out_ft = torch.zeros(batchsize, self.out_channels,  x.size(-1)//2 + 1,  device=x.device, dtype=torch.cfloat)
+        out_ft[:, :, :self.modes1] = self.compl_mul1d(x_ft[:, :, :self.modes1], self.weights1)
+
+        #Return to physical space
+        x = torch.fft.irfft(out_ft, n=x.size(-1))
+        return x
+
+class FNO1d(nn.Module):
+    def __init__(self, num_channels, modes=16, width=64, initial_step=10):
+        super(FNO1d, self).__init__()
+
+        """
+        The overall network. It contains 4 layers of the Fourier layer.
+        1. Lift the input to the desire channel dimension by self.fc0 .
+        2. 4 layers of the integral operators u' = (W + K)(u).
+            W defined by self.w; K defined by self.conv .
+        3. Project from the channel space to the output space by self.fc1 and self.fc2 .
+        
+        input: the solution of the initial condition and location (a(x), x)
+        input shape: (batchsize, x=s, c=2)
+        output: the solution of a later timestep
+        output shape: (batchsize, x=s, c=1)
+        """
+
+        self.modes1 = modes
+        self.width = width
+        self.padding = 2 # pad the domain if input is non-periodic
+        self.fc0 = nn.Linear(initial_step*num_channels+1, self.width) # input channel is 2: (a(x), x)
+
+        self.conv0 = SpectralConv1d(self.width, self.width, self.modes1)
+        self.conv1 = SpectralConv1d(self.width, self.width, self.modes1)
+        self.conv2 = SpectralConv1d(self.width, self.width, self.modes1)
+        self.conv3 = SpectralConv1d(self.width, self.width, self.modes1)
+        self.w0 = nn.Conv1d(self.width, self.width, 1)
+        self.w1 = nn.Conv1d(self.width, self.width, 1)
+        self.w2 = nn.Conv1d(self.width, self.width, 1)
+        self.w3 = nn.Conv1d(self.width, self.width, 1)
+
+        self.fc1 = nn.Linear(self.width, 128)
+        self.fc2 = nn.Linear(128, num_channels)
+
+    def forward(self, x, grid):
+        # x dim = [b, x1, t*v]
+        x = torch.cat((x, grid), dim=-1)
+        x = self.fc0(x)
+        x = x.permute(0, 2, 1)
+        
+        x = F.pad(x, [0, self.padding]) # pad the domain if input is non-periodic
+
+        x1 = self.conv0(x)
+        x2 = self.w0(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv1(x)
+        x2 = self.w1(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv2(x)
+        x2 = self.w2(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv3(x)
+        x2 = self.w3(x)
+        x = x1 + x2
+
+        x = x[..., :-self.padding]
+        x = x.permute(0, 2, 1)
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x.unsqueeze(-2)
+
+
+class SpectralConv2d_fast(nn.Module):
+    def __init__(self, in_channels, out_channels, modes1, modes2):
+        super(SpectralConv2d_fast, self).__init__()
+
+        """
+        2D Fourier layer. It does FFT, linear transform, and Inverse FFT.    
+        """
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1
+        self.modes2 = modes2
+
+        self.scale = (1 / (in_channels * out_channels))
+        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))
+        self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))
+
+    # Complex multiplication
+    def compl_mul2d(self, input, weights):
+        # (batch, in_channel, x,y ), (in_channel, out_channel, x,y) -> (batch, out_channel, x,y)
+        return torch.einsum("bixy,ioxy->boxy", input, weights)
+
+    def forward(self, x):
+        batchsize = x.shape[0]
+        #Compute Fourier coeffcients up to factor of e^(- something constant)
+        x_ft = torch.fft.rfft2(x)
+
+        # Multiply relevant Fourier modes
+        out_ft = torch.zeros(batchsize, self.out_channels,  x.size(-2), x.size(-1)//2 + 1, dtype=torch.cfloat, device=x.device)
+        out_ft[:, :, :self.modes1, :self.modes2] = \
+            self.compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)
+        out_ft[:, :, -self.modes1:, :self.modes2] = \
+            self.compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)
+
+        #Return to physical space
+        x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1)))
+        return x
+
+class FNO2d(nn.Module):
+    def __init__(self, num_channels, modes1=12, modes2=12, width=20, initial_step=10):
+        super(FNO2d, self).__init__()
+
+        """
+        The overall network. It contains 4 layers of the Fourier layer.
+        1. Lift the input to the desire channel dimension by self.fc0 .
+        2. 4 layers of the integral operators u' = (W + K)(u).
+            W defined by self.w; K defined by self.conv .
+        3. Project from the channel space to the output space by self.fc1 and self.fc2 .
+        
+        input: the solution of the previous 10 timesteps + 2 locations (u(t-10, x, y), ..., u(t-1, x, y),  x, y)
+        input shape: (batchsize, x, y, c)
+        output: the solution of the next timestep
+        output shape: (batchsize, x, y, c)
+        """
+
+        self.modes1 = modes1
+        self.modes2 = modes2
+        self.width = width
+        self.padding = 2 # pad the domain if input is non-periodic
+        self.fc0 = nn.Linear(initial_step*num_channels+2, self.width)
+        # input channel is 12: the solution of the previous 10 timesteps + 2 locations (u(t-10, x, y), ..., u(t-1, x, y),  x, y)
+
+        self.conv0 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2)
+        self.conv1 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2)
+        self.conv2 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2)
+        self.conv3 = SpectralConv2d_fast(self.width, self.width, self.modes1, self.modes2)
+        self.w0 = nn.Conv2d(self.width, self.width, 1)
+        self.w1 = nn.Conv2d(self.width, self.width, 1)
+        self.w2 = nn.Conv2d(self.width, self.width, 1)
+        self.w3 = nn.Conv2d(self.width, self.width, 1)
+
+        self.fc1 = nn.Linear(self.width, 128)
+        self.fc2 = nn.Linear(128, num_channels)
+
+    def forward(self, x, grid):
+        # x dim = [b, x1, x2, t*v]
+        x = torch.cat((x, grid), dim=-1)
+        x = self.fc0(x)
+        x = x.permute(0, 3, 1, 2)
+        
+        # Pad tensor with boundary condition
+        x = F.pad(x, [0, self.padding, 0, self.padding])
+
+        x1 = self.conv0(x)
+        x2 = self.w0(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv1(x)
+        x2 = self.w1(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv2(x)
+        x2 = self.w2(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv3(x)
+        x2 = self.w3(x)
+        x = x1 + x2
+
+        x = x[..., :-self.padding, :-self.padding] # Unpad the tensor
+        x = x.permute(0, 2, 3, 1)
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        
+        return x.unsqueeze(-2)
+    
+
+class SpectralConv3d(nn.Module):
+    def __init__(self, in_channels, out_channels, modes1, modes2, modes3):
+        super(SpectralConv3d, self).__init__()
+
+        """
+        3D Fourier layer. It does FFT, linear transform, and Inverse FFT.    
+        """
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.modes1 = modes1 #Number of Fourier modes to multiply, at most floor(N/2) + 1
+        self.modes2 = modes2
+        self.modes3 = modes3
+
+        self.scale = (1 / (in_channels * out_channels))
+        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))
+        self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))
+        self.weights3 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))
+        self.weights4 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))
+
+    # Complex multiplication
+    def compl_mul3d(self, input, weights):
+        # (batch, in_channel, x,y,t ), (in_channel, out_channel, x,y,t) -> (batch, out_channel, x,y,t)
+        return torch.einsum("bixyz,ioxyz->boxyz", input, weights)
+
+    def forward(self, x):
+        batchsize = x.shape[0]
+        #Compute Fourier coeffcients up to factor of e^(- something constant)
+        x_ft = torch.fft.rfftn(x, dim=[-3,-2,-1])
+
+        # Multiply relevant Fourier modes
+        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-3), x.size(-2), x.size(-1)//2 + 1, dtype=torch.cfloat, device=x.device)
+        out_ft[:, :, :self.modes1, :self.modes2, :self.modes3] = \
+            self.compl_mul3d(x_ft[:, :, :self.modes1, :self.modes2, :self.modes3], self.weights1)
+        out_ft[:, :, -self.modes1:, :self.modes2, :self.modes3] = \
+            self.compl_mul3d(x_ft[:, :, -self.modes1:, :self.modes2, :self.modes3], self.weights2)
+        out_ft[:, :, :self.modes1, -self.modes2:, :self.modes3] = \
+            self.compl_mul3d(x_ft[:, :, :self.modes1, -self.modes2:, :self.modes3], self.weights3)
+        out_ft[:, :, -self.modes1:, -self.modes2:, :self.modes3] = \
+            self.compl_mul3d(x_ft[:, :, -self.modes1:, -self.modes2:, :self.modes3], self.weights4)
+
+        #Return to physical space
+        x = torch.fft.irfftn(out_ft, s=(x.size(-3), x.size(-2), x.size(-1)))
+        return x
+
+class FNO3d(nn.Module):
+    def __init__(self, num_channels, modes1=8, modes2=8, modes3=8, width=20, initial_step=10):
+        super(FNO3d, self).__init__()
+
+        """
+        The overall network. It contains 4 layers of the Fourier layer.
+        1. Lift the input to the desire channel dimension by self.fc0 .
+        2. 4 layers of the integral operators u' = (W + K)(u).
+            W defined by self.w; K defined by self.conv .
+        3. Project from the channel space to the output space by self.fc1 and self.fc2 .
+        
+        input: the solution of the first 10 timesteps + 3 locations (u(1, x, y), ..., u(10, x, y),  x, y, t). It's a constant function in time, except for the last index.
+        input shape: (batchsize, x=64, y=64, t=40, c=13)
+        output: the solution of the next 40 timesteps
+        output shape: (batchsize, x=64, y=64, t=40, c=1)
+        """
+
+        self.modes1 = modes1
+        self.modes2 = modes2
+        self.modes3 = modes3
+        self.width = width
+        self.padding = 6 # pad the domain if input is non-periodic
+        self.fc0 = nn.Linear(initial_step*num_channels+3, self.width)
+        # input channel is 12: the solution of the first 10 timesteps + 3 locations (u(1, x, y), ..., u(10, x, y),  x, y, t)
+
+        self.conv0 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)
+        self.conv1 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)
+        self.conv2 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)
+        self.conv3 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)
+        self.w0 = nn.Conv3d(self.width, self.width, 1)
+        self.w1 = nn.Conv3d(self.width, self.width, 1)
+        self.w2 = nn.Conv3d(self.width, self.width, 1)
+        self.w3 = nn.Conv3d(self.width, self.width, 1)
+        self.bn0 = torch.nn.BatchNorm3d(self.width)
+        self.bn1 = torch.nn.BatchNorm3d(self.width)
+        self.bn2 = torch.nn.BatchNorm3d(self.width)
+        self.bn3 = torch.nn.BatchNorm3d(self.width)
+
+        self.fc1 = nn.Linear(self.width, 128)
+        self.fc2 = nn.Linear(128, num_channels)
+
+    def forward(self, x, grid):
+        # x dim = [b, x1, x2, x3, t*v]
+        x = torch.cat((x, grid), dim=-1)
+        x = self.fc0(x)
+        x = x.permute(0, 4, 1, 2, 3)
+        
+        x = F.pad(x, [0, self.padding]) # pad the domain if input is non-periodic
+
+        x1 = self.conv0(x)
+        x2 = self.w0(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv1(x)
+        x2 = self.w1(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv2(x)
+        x2 = self.w2(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv3(x)
+        x2 = self.w3(x)
+        x = x1 + x2
+
+        x = x[..., :-self.padding]
+        x = x.permute(0, 2, 3, 4, 1) # pad the domain if input is non-periodic
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x.unsqueeze(-2)
+    
+class FNO_maxwell(nn.Module):
+    def __init__(self, num_channels, modes1=8, modes2=8, modes3=8, width=20, initial_step=10):
+        super(FNO_maxwell, self).__init__()
+
+        """
+        The overall network. It contains 4 layers of the Fourier layer.
+        1. Lift the input to the desire channel dimension by self.fc0 .
+        2. 4 layers of the integral operators u' = (W + K)(u).
+            W defined by self.w; K defined by self.conv .
+        3. Project from the channel space to the output space by self.fc1 and self.fc2 .
+        
+        input: the solution of the first 10 timesteps + 3 locations (u(1, x, y), ..., u(10, x, y),  x, y, t). It's a constant function in time, except for the last index.
+        input shape: (batchsize, x=64, y=64, t=40, c=13)
+        output: the solution of the next 40 timesteps
+        output shape: (batchsize, x=64, y=64, t=40, c=1)
+        """
+
+        self.modes1 = modes1
+        self.modes2 = modes2
+        self.modes3 = modes3
+        self.width = width
+        self.padding = 6 # pad the domain if input is non-periodic
+        self.fc0 = nn.Linear(initial_step*num_channels, self.width)
+        # input channel is 12: the solution of the first 10 timesteps + 3 locations (u(1, x, y), ..., u(10, x, y),  x, y, t)
+
+        self.conv0 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)
+        self.conv1 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)
+        self.conv2 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)
+        self.conv3 = SpectralConv3d(self.width, self.width, self.modes1, self.modes2, self.modes3)
+        self.w0 = nn.Conv3d(self.width, self.width, 1)
+        self.w1 = nn.Conv3d(self.width, self.width, 1)
+        self.w2 = nn.Conv3d(self.width, self.width, 1)
+        self.w3 = nn.Conv3d(self.width, self.width, 1)
+        self.bn0 = torch.nn.BatchNorm3d(self.width)
+        self.bn1 = torch.nn.BatchNorm3d(self.width)
+        self.bn2 = torch.nn.BatchNorm3d(self.width)
+        self.bn3 = torch.nn.BatchNorm3d(self.width)
+
+        self.fc1 = nn.Linear(self.width, 128)
+        self.fc2 = nn.Linear(128, num_channels)
+
+    def forward(self, x, grid):
+        # x dim = [b, x1, x2, x3, t*v]
+        x = self.fc0(x)
+        x = x.permute(0, 4, 1, 2, 3)
+        
+        x = F.pad(x, [0, self.padding]) # pad the domain if input is non-periodic
+
+        x1 = self.conv0(x)
+        x2 = self.w0(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv1(x)
+        x2 = self.w1(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv2(x)
+        x2 = self.w2(x)
+        x = x1 + x2
+        x = F.gelu(x)
+
+        x1 = self.conv3(x)
+        x2 = self.w3(x)
+        x = x1 + x2
+
+        x = x[..., :-self.padding]
+        x = x.permute(0, 2, 3, 4, 1) # pad the domain if input is non-periodic
+        x = self.fc1(x)
+        x = F.gelu(x)
+        x = self.fc2(x)
+        return x.unsqueeze(-2)

model/mpnn.py

@@ -0,0 +1,201 @@
+# coding=utf-8
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+from torch import nn
+from torch_geometric.data import Data
+from torch_geometric.nn import MessagePassing, InstanceNorm
+
+if TYPE_CHECKING:
+    from onescience.utils.pdenneval.mpnn_utils import PDE
+
+class Swish(nn.Module):
+    """Swish activation function
+    """
+    def __init__(self, beta=1):
+        super(Swish, self).__init__()
+        self.beta = beta
+
+    def forward(self, x):
+        return x * torch.sigmoid(self.beta*x)
+    
+
+class GNN_Layer(MessagePassing):
+    """Message passing layer
+    """
+    def __init__(self,
+                 in_features: int,
+                 out_features: int,
+                 hidden_features: int,
+                 time_window: int,
+                 spatial_dim: int,
+                 n_variables: int):
+        """Initialize message passing layers
+        Args:
+            in_features (int): number of node input features
+            out_features (int): number of node output features
+            hidden_features (int): number of hidden features
+            time_window (int): number of input/output timesteps (temporal bundling)
+            spatial_dim (int): number of dimension of spatial domain  
+            n_variables (int): number of equation specific parameters used in the solver
+        """
+        super(GNN_Layer, self).__init__(node_dim=-2, aggr='mean') # node_dim: The axis along which to propagate. (default: -2)
+        self.in_features = in_features
+        self.out_features = out_features
+        self.hidden_features = hidden_features
+        assert (spatial_dim == 1 or spatial_dim == 2 or spatial_dim == 3)
+
+        self.message_net_1 = nn.Sequential(nn.Linear(2 * in_features + time_window + spatial_dim + n_variables, hidden_features), 
+                                           Swish()
+                                           )
+        self.message_net_2 = nn.Sequential(nn.Linear(hidden_features, hidden_features), 
+                                           Swish()
+                                           )
+        self.update_net_1 = nn.Sequential(nn.Linear(in_features + hidden_features + n_variables, hidden_features), 
+                                          Swish()
+                                          )
+        self.update_net_2 = nn.Sequential(nn.Linear(hidden_features, out_features), 
+                                          Swish()
+                                          )
+        self.norm = InstanceNorm(hidden_features)
+
+    def forward(self, x, u, pos, variables, edge_index, batch):
+        """Propagate messages along edges
+        """
+        x = self.propagate(edge_index, x=x, u=u, pos=pos, variables=variables)
+        x = self.norm(x, batch)
+        return x
+
+    def message(self, x_i, x_j, u_i, u_j, pos_i, pos_j, variables_i):
+        """Message update following formula 8 of the paper
+        """
+        message = self.message_net_1(torch.cat((x_i, x_j, u_i - u_j, pos_i - pos_j, variables_i), dim=-1))
+        message = self.message_net_2(message)
+        return message
+
+    def update(self, message, x, variables):
+        """Node update following formula 9 of the paper
+        """
+        update = self.update_net_1(torch.cat((x, message, variables), dim=-1))
+        update = self.update_net_2(update)
+        if self.in_features == self.out_features:
+            return x + update
+        else:
+            return update
+
+
+
+class MPNN(nn.Module):
+    def __init__(self,
+                 pde: PDE,
+                 time_window: int = 25,
+                 hidden_features: int = 128,
+                 hidden_layers: int = 6,
+                 eq_variables: dict = {}):
+        """Initialize MPNN
+        It contains 6 MP-PDE layers with skip connections
+        The input graph to the forward pass has the shape [batch*n_nodes, time_window].
+        The output graph has the shape [batch*n_nodes, time_window].
+        Args:
+            pde (PDE): PDE to solve
+            time_window (int): number of input/output timesteps (temporal bundling)
+            hidden features (int): number of hidden features
+            hidden_layer (int): number of hidden layers
+            eq_variables (dict): dictionary of equation specific parameters
+        """
+        super(MPNN, self).__init__()
+        # arguments
+        self.pde = pde
+        self.out_features = time_window
+        self.hidden_features = hidden_features
+        self.hidden_layers = hidden_layers
+        self.time_window = time_window
+        self.eq_variables = eq_variables
+
+        # encoder
+        self.embedding_mlp = nn.Sequential(
+            nn.Linear(self.time_window+self.pde.spatial_dim+1+len(self.eq_variables), self.hidden_features),
+            Swish(),
+            nn.Linear(self.hidden_features, self.hidden_features),
+            Swish()
+            )
+
+        # processor
+        self.gnn_layers = torch.nn.ModuleList(modules=(GNN_Layer(
+            in_features=self.hidden_features,
+            hidden_features=self.hidden_features,
+            out_features=self.hidden_features,
+            time_window=self.time_window,
+            spatial_dim=self.pde.spatial_dim,
+            n_variables=len(self.eq_variables)+1 # (time is treated as equation variable)
+            ) for _ in range(self.hidden_layers)))
+        
+        # decoder
+        if(self.time_window==10): # NEW ADD
+            self.output_mlp = nn.Sequential(
+                nn.Conv1d(1, 8, 18, stride=5),
+                Swish(),
+                nn.Conv1d(8, 1, 14, stride=1)
+                )
+        if(self.time_window==20):
+            self.output_mlp = nn.Sequential(
+                nn.Conv1d(1, 8, 15, stride=4),
+                Swish(),
+                nn.Conv1d(8, 1, 10, stride=1)
+                )
+        if (self.time_window==25):
+            self.output_mlp = nn.Sequential(
+                nn.Conv1d(1, 8, 16, stride=3),
+                Swish(),
+                nn.Conv1d(8, 1, 14, stride=1)
+                )
+        if(self.time_window==50):
+            self.output_mlp = nn.Sequential(
+                nn.Conv1d(1, 8, 12, stride=2),
+                Swish(),
+                nn.Conv1d(8, 1, 10, stride=1)
+                )
+            
+    def forward(self, data: Data, v: int=0) -> torch.Tensor:
+        """Forward pass of MP-PDE solver class.
+        The input graph has the shape [batch*n_nodes, time_window, v].
+        The output tensor has the shape [batch*n_nodes, time_window].
+        Args:
+            data (Data): Pytorch Geometric data graph
+            v (int): physics variable to solve
+        Returns:
+            torch.Tensor: predictive solutions with shape [batch*n_nodes, time_window]
+        """
+        u = data.x[..., v] # [bs*nx, tw]
+        x_pos = data.x_pos # [bs*nx, spatial_dim]
+        t_pos = data.t_pos # [bs*nx]
+        edge_index = data.edge_index # [2, num_edges]
+        batch = data.batch # [bs*nx]
+        variables = data.variables # [bs*nx, num_variables]
+
+        # normalize temporal and spatial coordinate
+        t_pos = t_pos / self.pde.tmax
+        for d in range(self.pde.spatial_dim):
+            x_pos[:, d] = x_pos[:, d] / (self.pde.spatial_domain[d][1] - self.pde.spatial_domain[d][0])
+
+        # encode
+        node_input = torch.cat((u, x_pos, t_pos.unsqueeze(1), variables), dim=-1)
+        f = self.embedding_mlp(node_input) # [bs*nx, hidden_dim]
+
+        # process
+        for i in range(self.hidden_layers):
+            # f = self.gnn_layers[i](f, u, x_pos, variables, edge_index, batch)
+            f = self.gnn_layers[i](f, u, x_pos, torch.cat((t_pos.unsqueeze(1), variables), dim=-1), edge_index, batch)
+
+        # decode
+        dt = (self.pde.tmax - self.pde.tmin) / self.pde.resolution_t
+        dt = (torch.ones(1, self.time_window) * dt).to(f.device) # [1, tw]
+        dt = torch.cumsum(dt, dim=1) # [1, tw]
+        # [bs*nx, hidden_dim] -> [bs*nx, 1, hidden_dim]) -> conv1d -> [bs*nx, 1, tw] -> [bs*nx, tw] (diff)
+        diff = self.output_mlp(f[:, None]).squeeze(1)
+        # [bs*nx] -> [tw, bs*nx] -> [bs*nx, tw]
+        out = u[:, -1].repeat(self.time_window, 1).transpose(0, 1) + dt * diff
+        # print(out.shape)
+        return out

model/pino_fno.py

@@ -0,0 +1,420 @@
+# coding=utf-8
+############################
+'''
+The implement codes of FNO were taken and modified from: https://github.com/neuraloperator/physics_informed
+'''
+
+import numpy as np
+
+import torch
+import torch.nn as nn
+
+from onescience.utils.pdenneval.pino_utils import add_padding, remove_padding, add_padding2, remove_padding2, add_padding3, remove_padding3, _get_act
+
+
+@torch.jit.script
+def compl_mul1d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    # (batch, in_channel, x ), (in_channel, out_channel, x) -> (batch, out_channel, x)
+    res = torch.einsum("bix,iox->box", a, b)
+    return res
+
+
+@torch.jit.script
+def compl_mul2d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    # (batch, in_channel, x,y,t ), (in_channel, out_channel, x,y,t) -> (batch, out_channel, x,y,t)
+    res =  torch.einsum("bixy,ioxy->boxy", a, b)
+    return res
+
+
+@torch.jit.script
+def compl_mul3d(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    res = torch.einsum("bixyz,ioxyz->boxyz", a, b)
+    return res
+
+################################################################
+# 1d fourier layer
+################################################################
+
+
+class SpectralConv1d(nn.Module):
+    def __init__(self, in_channels, out_channels, modes1):
+        super(SpectralConv1d, self).__init__()
+
+        """
+        1D Fourier layer. It does FFT, linear transform, and Inverse FFT.
+        """
+
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        # Number of Fourier modes to multiply, at most floor(N/2) + 1
+        self.modes1 = modes1
+
+        self.scale = (1 / (in_channels*out_channels))
+        self.weights1 = nn.Parameter(
+            self.scale * torch.rand(in_channels, out_channels, self.modes1, dtype=torch.cfloat))
+
+    def forward(self, x):
+        batchsize = x.shape[0]
+        # Compute Fourier coeffcients up to factor of e^(- something constant)
+        x_ft = torch.fft.rfftn(x, dim=[2])
+
+        # Multiply relevant Fourier modes
+        out_ft = torch.zeros(batchsize, self.in_channels, x.size(-1)//2 + 1, device=x.device, dtype=torch.cfloat)
+        out_ft[:, :, :self.modes1] = compl_mul1d(x_ft[:, :, :self.modes1], self.weights1)
+
+        # Return to physical space
+        x = torch.fft.irfftn(out_ft, s=[x.size(-1)], dim=[2])
+        return x
+
+################################################################
+# 2d fourier layer
+################################################################
+
+
+class SpectralConv2d(nn.Module):
+    def __init__(self, in_channels, out_channels, modes1, modes2):
+        super(SpectralConv2d, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        # Number of Fourier modes to multiply, at most floor(N/2) + 1
+        self.modes1 = modes1
+        self.modes2 = modes2
+
+        self.scale = (1 / (in_channels * out_channels))
+        self.weights1 = nn.Parameter(
+            self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))
+        self.weights2 = nn.Parameter(
+            self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, dtype=torch.cfloat))
+
+    def forward(self, x):
+        batchsize = x.shape[0]
+        size1 = x.shape[-2]
+        size2 = x.shape[-1]
+        # Compute Fourier coeffcients up to factor of e^(- something constant)
+        x_ft = torch.fft.rfftn(x, dim=[2, 3])
+
+        # Multiply relevant Fourier modes
+        out_ft = torch.zeros(batchsize, self.out_channels, x.size(-2), x.size(-1) // 2 + 1, device=x.device,
+                                dtype=torch.cfloat)
+        out_ft[:, :, :self.modes1, :self.modes2] = \
+            compl_mul2d(x_ft[:, :, :self.modes1, :self.modes2], self.weights1)
+        out_ft[:, :, -self.modes1:, :self.modes2] = \
+            compl_mul2d(x_ft[:, :, -self.modes1:, :self.modes2], self.weights2)
+
+        # Return to physical space
+        x = torch.fft.irfftn(out_ft, s=(x.size(-2), x.size(-1)), dim=[2, 3])
+        return x
+
+
+class SpectralConv3d(nn.Module):
+    def __init__(self, in_channels, out_channels, modes1, modes2, modes3):
+        super(SpectralConv3d, self).__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.modes1 = modes1  #Number of Fourier modes to multiply, at most floor(N/2) + 1
+        self.modes2 = modes2
+        self.modes3 = modes3
+
+        self.scale = (1 / (in_channels * out_channels))
+        self.weights1 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))
+        self.weights2 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))
+        self.weights3 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))
+        self.weights4 = nn.Parameter(self.scale * torch.rand(in_channels, out_channels, self.modes1, self.modes2, self.modes3, dtype=torch.cfloat))
+
+    def forward(self, x):
+        batchsize = x.shape[0]
+        # Compute Fourier coeffcients up to factor of e^(- something constant)
+        x_ft = torch.fft.rfftn(x, dim=[2,3,4])
+
+        z_dim = min(x_ft.shape[4], self.modes3)
+
+        # Multiply relevant Fourier modes
+        out_ft = torch.zeros(batchsize, self.out_channels, x_ft.shape[2], x_ft.shape[3], self.modes3, device=x.device, dtype=torch.cfloat)
+
+        # if x_ft.shape[4] > self.modes3, truncate; if x_ft.shape[4] < self.modes3, add zero padding
+        coeff = torch.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, device=x.device, dtype=torch.cfloat)
+        coeff[..., :z_dim] = x_ft[:, :, :self.modes1, :self.modes2, :z_dim]
+        out_ft[:, :, :self.modes1, :self.modes2, :] = compl_mul3d(coeff, self.weights1)
+
+        coeff = torch.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, device=x.device, dtype=torch.cfloat)
+        coeff[..., :z_dim] = x_ft[:, :, -self.modes1:, :self.modes2, :z_dim]
+        out_ft[:, :, -self.modes1:, :self.modes2, :] = compl_mul3d(coeff, self.weights2)
+
+        coeff = torch.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, device=x.device, dtype=torch.cfloat)
+        coeff[..., :z_dim] = x_ft[:, :, :self.modes1, -self.modes2:, :z_dim]
+        out_ft[:, :, :self.modes1, -self.modes2:, :] = compl_mul3d(coeff, self.weights3)
+
+        coeff = torch.zeros(batchsize, self.in_channels, self.modes1, self.modes2, self.modes3, device=x.device, dtype=torch.cfloat)
+        coeff[..., :z_dim] = x_ft[:, :, -self.modes1:, -self.modes2:, :z_dim]
+        out_ft[:, :, -self.modes1:, -self.modes2:, :] = compl_mul3d(coeff, self.weights4)
+
+        #Return to physical space
+        x = torch.fft.irfftn(out_ft, s=(x.size(2), x.size(3), x.size(4)), dim=[2,3,4])
+        return x
+
+
+class FourierBlock(nn.Module):
+    def __init__(self, in_channels, out_channels, modes1, modes2, modes3, act='tanh'):
+        super(FourierBlock, self).__init__()
+        self.in_channel = in_channels
+        self.out_channel = out_channels
+        self.speconv = SpectralConv3d(in_channels, out_channels, modes1, modes2, modes3)
+        self.linear = nn.Conv1d(in_channels, out_channels, 1)
+        if act in ['tanh','gelu','none']:
+            self.act=_get_act(act)
+        else:
+            raise ValueError(f'{act} is not supported')
+
+    def forward(self, x):
+        '''
+        input x: (batchsize, channel width, x_grid, y_grid, t_grid)
+        '''
+        x1 = self.speconv(x)
+        x2 = self.linear(x.view(x.shape[0], self.in_channel, -1))
+        out = x1 + x2.view(x.shape[0], self.out_channel, x.shape[2], x.shape[3], x.shape[4])
+        if self.act is not None:
+            out = self.act(out)
+        return out
+
+
+class FNO1d(nn.Module):
+    def __init__(self,
+                 modes, width=32, 
+                 layers=None,
+                 fc_dim=128,
+                 in_dim=2, out_dim=1,
+                 act='relu',
+                 pad_ratio=[0.,0.1]):
+        super(FNO1d, self).__init__()
+
+        """
+        The overall network. It contains several layers of the Fourier layer.
+        1. Lift the input to the desire channel dimension by self.fc0 .
+        2. 4 layers of the integral operators u' = (W + K)(u).
+            W defined by self.w; K defined by self.conv .
+        3. Project from the channel space to the output space by self.fc1 and self.fc2 .
+
+        input: the solution of the initial condition and location (a(x), x)
+        input shape: (batchsize, x=s, c=2)
+        output: the solution of a later timestep
+        output shape: (batchsize, x=s, c=1)
+        """
+
+        self.modes1 = modes
+        self.width = width
+        self.pad_ratio = pad_ratio
+        if layers is None:
+            layers = [width] * 4
+
+        self.fc0 = nn.Linear(in_dim, layers[0])  # input channel is 2: (a(x), x)
+
+        self.sp_convs = nn.ModuleList([SpectralConv1d(
+            in_size, out_size, num_modes) for in_size, out_size, num_modes in zip(layers, layers[1:], self.modes1)])
+
+        self.ws = nn.ModuleList([nn.Conv1d(in_size, out_size, 1)
+                                 for in_size, out_size in zip(layers, layers[1:])])
+
+        self.fc1 = nn.Linear(layers[-1], fc_dim)
+        self.fc2 = nn.Linear(fc_dim, out_dim)
+        self.act = _get_act(act)
+
+    def forward(self, x):
+        length = len(self.ws)
+        size_1= x.shape[1]
+        if max(self.pad_ratio) > 0:
+            num_pad = [round(size_1 * i) for i in self.pad_ratio]
+        else:
+            num_pad = [0., 0.]
+        x = self.fc0(x)
+        x = x.permute(0, 2, 1)
+        
+        x = add_padding(x,num_pad)
+        for i, (speconv, w) in enumerate(zip(self.sp_convs, self.ws)):
+            x1 = speconv(x)
+            x2 = w(x)
+            x = x1 + x2
+            if i != length - 1:
+                x = self.act(x)
+        x = remove_padding(x,num_pad)
+
+        x = x.permute(0, 2, 1)
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return x
+
+
+class FNO2d(nn.Module):
+    def __init__(self, modes1, modes2,
+                 width=64, fc_dim=128,
+                 layers=None,
+                 in_dim=3, out_dim=1,
+                 act='gelu',
+                 pad_ratio=[0., 0.1]):
+        super(FNO2d, self).__init__()
+        """
+        Args:
+            - modes1: list of int, number of modes in first dimension in each layer
+            - modes2: list of int, number of modes in second dimension in each layer
+            - width: int, optional, if layers is None, it will be initialized as [width] * [len(modes1) + 1]
+            - in_dim: number of input channels
+            - out_dim: number of output channels
+            - act: activation function, {tanh, gelu, relu, leaky_relu}, default: gelu
+            - pad_ratio: list of float, or float; portion of domain to be extended. If float, paddings are added to the right.
+            If list, paddings are added to both sides. pad_ratio[0] pads left, pad_ratio[1] pads right.
+        """
+        if isinstance(pad_ratio, float):
+            pad_ratio = [pad_ratio, pad_ratio]
+        else:
+            assert len(pad_ratio) == 2, 'Cannot add padding in more than 2 directions'
+        self.modes1 = modes1
+        self.modes2 = modes2
+
+        self.pad_ratio = pad_ratio
+        # input channel is 3: (a(x, y), x, y)
+        if layers is None:
+            self.layers = [width] * (len(modes1) + 1)
+        else:
+            self.layers = layers
+        
+        self.fc0 = nn.Linear(in_dim, self.layers[0])
+
+        self.sp_convs = nn.ModuleList([SpectralConv2d(
+            in_size, out_size, mode1_num, mode2_num)
+            for in_size, out_size, mode1_num, mode2_num
+            in zip(self.layers, self.layers[1:], self.modes1, self.modes2)])
+
+        self.ws = nn.ModuleList([nn.Conv1d(in_size, out_size, 1)
+                                 for in_size, out_size in zip(self.layers, self.layers[1:])])
+
+        self.fc1 = nn.Linear(self.layers[-1], fc_dim)
+        self.fc2 = nn.Linear(fc_dim, self.layers[-1])
+        self.fc3 = nn.Linear(self.layers[-1], out_dim)
+        self.act = _get_act(act)
+
+    def forward(self, x):
+        '''
+        Args:
+            - x : (batch size, x_grid, y_grid, 2)
+        Returns:
+            - x: (batch size, x_grid, y_grid, 1)
+        '''
+        size_1, size_2 = x.shape[1], x.shape[2]
+        if max(self.pad_ratio) > 0:
+            num_pad1 = [round(i * size_1) for i in self.pad_ratio]
+            num_pad2 = [round(i * size_2) for i in self.pad_ratio]
+        else:
+            num_pad1 = num_pad2 = [0.,0.]
+
+        length = len(self.ws)
+        batchsize = x.shape[0]
+        x = self.fc0(x)
+        x = x.permute(0, 3, 1, 2)   # B, C, X, Y
+        x = add_padding2(x, num_pad1, num_pad2)
+        size_x, size_y = x.shape[-2], x.shape[-1]
+
+        for i, (speconv, w) in enumerate(zip(self.sp_convs, self.ws)):
+            x1 = speconv(x)
+            x2 = w(x.view(batchsize, self.layers[i], -1)).view(batchsize, self.layers[i+1], size_x, size_y)
+            x = x1 + x2
+            if i != length - 1:
+                x = self.act(x)
+        x = remove_padding2(x, num_pad1, num_pad2)
+        x = x.permute(0, 2, 3, 1)
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        x = self.act(x)
+        x = self.fc3(x)
+        return x
+
+class FNO3d(nn.Module):
+    def __init__(self,
+                 modes1, modes2, modes3,
+                 width=16,
+                 fc_dim=128,
+                 layers=None,
+                 in_dim=4, out_dim=1,
+                 act='gelu',
+                 pad_ratio=[0., 0.05]):
+        '''
+        Args:
+            modes1: list of int, first dimension maximal modes for each layer
+            modes2: list of int, second dimension maximal modes for each layer
+            modes3: list of int, third dimension maximal modes for each layer
+            layers: list of int, channels for each layer
+            fc_dim: dimension of fully connected layers
+            in_dim: int, input dimension
+            out_dim: int, output dimension
+            act: {tanh, gelu, relu, leaky_relu}, activation function
+            pad_ratio: the ratio of the extended domain
+        '''
+        super(FNO3d, self).__init__()
+
+        if isinstance(pad_ratio, float):
+            pad_ratio = [pad_ratio, pad_ratio]
+        else:
+            assert len(pad_ratio) == 2, 'Cannot add padding in more than 2 directions.'
+
+        self.pad_ratio = pad_ratio
+        self.modes1 = modes1
+        self.modes2 = modes2
+        self.modes3 = modes3
+        self.pad_ratio = pad_ratio
+
+        if layers is None:
+            self.layers = [width] * 4
+        else:
+            self.layers = layers
+        self.fc0 = nn.Linear(in_dim, self.layers[0])
+
+        self.sp_convs = nn.ModuleList([SpectralConv3d(
+            in_size, out_size, mode1_num, mode2_num, mode3_num)
+            for in_size, out_size, mode1_num, mode2_num, mode3_num
+            in zip(self.layers, self.layers[1:], self.modes1, self.modes2, self.modes3)])
+
+        self.ws = nn.ModuleList([nn.Conv1d(in_size, out_size, 1)
+                                 for in_size, out_size in zip(self.layers, self.layers[1:])])
+
+        self.fc1 = nn.Linear(self.layers[-1], fc_dim)
+        self.fc2 = nn.Linear(fc_dim, out_dim)
+        self.act = _get_act(act)
+
+    def forward(self, x):
+        '''
+        Args:
+            x: (batchsize, x_grid, y_grid, t_grid, 3)
+
+        Returns:
+            u: (batchsize, x_grid, y_grid, t_grid, 1)
+
+        '''
+        size_x,size_y,size_z = x.shape[1],x.shape[2],x.shape[3]
+        if max(self.pad_ratio) > 0:
+            num_pad1 = [round(size_x * i) for i in self.pad_ratio]
+            num_pad2 = [round(size_y * i) for i in self.pad_ratio]
+            num_pad3 = [round(size_z * i) for i in self.pad_ratio]
+        else:
+            num_pad1 = num_pad2 = num_pad3 = [0., 0.]
+        length = len(self.ws)
+        batchsize = x.shape[0]
+
+        x = self.fc0(x)
+        x = x.permute(0, 4, 1, 2, 3)
+        x = add_padding3(x, num_pad1,num_pad2,num_pad3)
+        size_x, size_y, size_z = x.shape[-3], x.shape[-2], x.shape[-1]
+
+        for i, (speconv, w) in enumerate(zip(self.sp_convs, self.ws)):
+            x1 = speconv(x)
+            x2 = w(x.view(batchsize, self.layers[i], -1)).view(batchsize, self.layers[i+1], size_x, size_y, size_z)
+            x = x1 + x2
+            if i != length - 1:
+                x = self.act(x)
+        x = remove_padding3(x, num_pad1,num_pad2,num_pad3)
+        x = x.permute(0, 2, 3, 4, 1)
+        x = self.fc1(x)
+        x = self.act(x)
+        x = self.fc2(x)
+        return x
+
+

model/unet.py

@@ -0,0 +1,255 @@
+import torch
+import torch.nn as nn
+from collections import OrderedDict
+
+class UNet1d(nn.Module):
+
+    def __init__(self, in_channels=3, out_channels=1, init_features=32):
+        super(UNet1d, self).__init__()
+
+        features = init_features
+        self.encoder1 = UNet1d._block(in_channels, features, name="enc1")
+        self.pool1 = nn.MaxPool1d(kernel_size=2, stride=2)
+        self.encoder2 = UNet1d._block(features, features * 2, name="enc2")
+        self.pool2 = nn.MaxPool1d(kernel_size=2, stride=2)
+        self.encoder3 = UNet1d._block(features * 2, features * 4, name="enc3")
+        self.pool3 = nn.MaxPool1d(kernel_size=2, stride=2)
+        self.encoder4 = UNet1d._block(features * 4, features * 8, name="enc4")
+        self.pool4 = nn.MaxPool1d(kernel_size=2, stride=2)
+
+        self.bottleneck = UNet1d._block(features * 8, features * 16, name="bottleneck")
+
+        self.upconv4 = nn.ConvTranspose1d(features * 16, features * 8, kernel_size=2, stride=2)
+        self.decoder4 = UNet1d._block((features * 8) * 2, features * 8, name="dec4")
+        self.upconv3 = nn.ConvTranspose1d(features * 8, features * 4, kernel_size=2, stride=2)
+        self.decoder3 = UNet1d._block((features * 4) * 2, features * 4, name="dec3")
+        self.upconv2 = nn.ConvTranspose1d(features * 4, features * 2, kernel_size=2, stride=2)
+        self.decoder2 = UNet1d._block((features * 2) * 2, features * 2, name="dec2")
+        self.upconv1 = nn.ConvTranspose1d(features * 2, features, kernel_size=2, stride=2)
+        self.decoder1 = UNet1d._block(features * 2, features, name="dec1")
+
+        self.conv = nn.Conv1d(in_channels=features, out_channels=out_channels, kernel_size=1)
+
+    def forward(self, x):
+        enc1 = self.encoder1(x)
+        enc2 = self.encoder2(self.pool1(enc1))
+        enc3 = self.encoder3(self.pool2(enc2))
+        enc4 = self.encoder4(self.pool3(enc3))
+
+        bottleneck = self.bottleneck(self.pool4(enc4))
+
+        dec4 = self.upconv4(bottleneck)
+        dec4 = torch.cat((dec4, enc4), dim=1)
+        dec4 = self.decoder4(dec4)
+        dec3 = self.upconv3(dec4)
+        dec3 = torch.cat((dec3, enc3), dim=1)
+        dec3 = self.decoder3(dec3)
+        dec2 = self.upconv2(dec3)
+        dec2 = torch.cat((dec2, enc2), dim=1)
+        dec2 = self.decoder2(dec2)
+        dec1 = self.upconv1(dec2)
+        dec1 = torch.cat((dec1, enc1), dim=1)
+        dec1 = self.decoder1(dec1)
+        out = self.conv(dec1)
+        return out
+
+    @staticmethod
+    def _block(in_channels, features, name):
+        return nn.Sequential(
+            OrderedDict(
+                [
+                    (
+                        name + "conv1",
+                        nn.Conv1d(
+                            in_channels=in_channels,
+                            out_channels=features,
+                            kernel_size=3,
+                            padding=1,
+                            bias=False,
+                        ),
+                    ),
+                    (name + "norm1", nn.BatchNorm1d(num_features=features)),
+                    (name + "tanh1", nn.Tanh()),
+                    (
+                        name + "conv2",
+                        nn.Conv1d(
+                            in_channels=features,
+                            out_channels=features,
+                            kernel_size=3,
+                            padding=1,
+                            bias=False,
+                        ),
+                    ),
+                    (name + "norm2", nn.BatchNorm1d(num_features=features)),
+                    (name + "tanh2", nn.Tanh()),
+                ]
+            )
+        )
+
+
+class UNet2d(nn.Module):
+
+    def __init__(self, in_channels=3, out_channels=1, init_features=32):
+        super(UNet2d, self).__init__()
+
+        features = init_features
+        self.encoder1 = UNet2d._block(in_channels, features, name="enc1")
+        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.encoder2 = UNet2d._block(features, features * 2, name="enc2")
+        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.encoder3 = UNet2d._block(features * 2, features * 4, name="enc3")
+        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
+        self.encoder4 = UNet2d._block(features * 4, features * 8, name="enc4")
+        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)
+
+        self.bottleneck = UNet2d._block(features * 8, features * 16, name="bottleneck")
+
+        self.upconv4 = nn.ConvTranspose2d(features * 16, features * 8, kernel_size=2, stride=2)
+        self.decoder4 = UNet2d._block((features * 8) * 2, features * 8, name="dec4")
+        self.upconv3 = nn.ConvTranspose2d(features * 8, features * 4, kernel_size=2, stride=2)
+        self.decoder3 = UNet2d._block((features * 4) * 2, features * 4, name="dec3")
+        self.upconv2 = nn.ConvTranspose2d(features * 4, features * 2, kernel_size=2, stride=2)
+        self.decoder2 = UNet2d._block((features * 2) * 2, features * 2, name="dec2")
+        self.upconv1 = nn.ConvTranspose2d(features * 2, features, kernel_size=2, stride=2)
+        self.decoder1 = UNet2d._block(features * 2, features, name="dec1")
+
+        self.conv = nn.Conv2d(in_channels=features, out_channels=out_channels, kernel_size=1)
+
+    def forward(self, x):
+        enc1 = self.encoder1(x)
+        enc2 = self.encoder2(self.pool1(enc1))
+        enc3 = self.encoder3(self.pool2(enc2))
+        enc4 = self.encoder4(self.pool3(enc3))
+
+        bottleneck = self.bottleneck(self.pool4(enc4))
+
+        dec4 = self.upconv4(bottleneck)
+        dec4 = torch.cat((dec4, enc4), dim=1)
+        dec4 = self.decoder4(dec4)
+        dec3 = self.upconv3(dec4)
+        dec3 = torch.cat((dec3, enc3), dim=1)
+        dec3 = self.decoder3(dec3)
+        dec2 = self.upconv2(dec3)
+        dec2 = torch.cat((dec2, enc2), dim=1)
+        dec2 = self.decoder2(dec2)
+        dec1 = self.upconv1(dec2)
+        dec1 = torch.cat((dec1, enc1), dim=1)
+        dec1 = self.decoder1(dec1)
+        return self.conv(dec1)
+
+    @staticmethod
+    def _block(in_channels, features, name):
+        return nn.Sequential(
+            OrderedDict(
+                [
+                    (
+                        name + "conv1",
+                        nn.Conv2d(
+                            in_channels=in_channels,
+                            out_channels=features,
+                            kernel_size=3,
+                            padding=1,
+                            bias=False,
+                        ),
+                    ),
+                    (name + "norm1", nn.BatchNorm2d(num_features=features)),
+                    (name + "tanh1", nn.Tanh()),
+                    (
+                        name + "conv2",
+                        nn.Conv2d(
+                            in_channels=features,
+                            out_channels=features,
+                            kernel_size=3,
+                            padding=1,
+                            bias=False,
+                        ),
+                    ),
+                    (name + "norm2", nn.BatchNorm2d(num_features=features)),
+                    (name + "tanh2", nn.Tanh()),
+                ]
+            )
+        )
+    
+
+class UNet3d(nn.Module):
+
+    def __init__(self, in_channels=3, out_channels=1, init_features=32):
+        super(UNet3d, self).__init__()
+
+        features = init_features
+        self.encoder1 = UNet3d._block(in_channels, features, name="enc1")
+        self.pool1 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.encoder2 = UNet3d._block(features, features * 2, name="enc2")
+        self.pool2 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.encoder3 = UNet3d._block(features * 2, features * 4, name="enc3")
+        self.pool3 = nn.MaxPool3d(kernel_size=2, stride=2)
+        self.encoder4 = UNet3d._block(features * 4, features * 8, name="enc4")
+        self.pool4 = nn.MaxPool3d(kernel_size=2, stride=2)
+
+        self.bottleneck = UNet3d._block(features * 8, features * 16, name="bottleneck")
+
+        self.upconv4 = nn.ConvTranspose3d(features * 16, features * 8, kernel_size=2, stride=2)
+        self.decoder4 = UNet3d._block((features * 8) * 2, features * 8, name="dec4")
+        self.upconv3 = nn.ConvTranspose3d(features * 8, features * 4, kernel_size=2, stride=2)
+        self.decoder3 = UNet3d._block((features * 4) * 2, features * 4, name="dec3")
+        self.upconv2 = nn.ConvTranspose3d(features * 4, features * 2, kernel_size=2, stride=2)
+        self.decoder2 = UNet3d._block((features * 2) * 2, features * 2, name="dec2")
+        self.upconv1 = nn.ConvTranspose3d(features * 2, features, kernel_size=2, stride=2)
+        self.decoder1 = UNet3d._block(features * 2, features, name="dec1")
+
+        self.conv = nn.Conv3d(in_channels=features, out_channels=out_channels, kernel_size=1)
+
+    def forward(self, x):
+        enc1 = self.encoder1(x)
+        enc2 = self.encoder2(self.pool1(enc1))
+        enc3 = self.encoder3(self.pool2(enc2))
+        enc4 = self.encoder4(self.pool3(enc3))
+
+        bottleneck = self.bottleneck(self.pool4(enc4))
+
+        dec4 = self.upconv4(bottleneck)
+        dec4 = torch.cat((dec4, enc4), dim=1)
+        dec4 = self.decoder4(dec4)
+        dec3 = self.upconv3(dec4)
+        dec3 = torch.cat((dec3, enc3), dim=1)
+        dec3 = self.decoder3(dec3)
+        dec2 = self.upconv2(dec3)
+        dec2 = torch.cat((dec2, enc2), dim=1)
+        dec2 = self.decoder2(dec2)
+        dec1 = self.upconv1(dec2)
+        dec1 = torch.cat((dec1, enc1), dim=1)
+        dec1 = self.decoder1(dec1)
+        return self.conv(dec1)
+
+    @staticmethod
+    def _block(in_channels, features, name):
+        return nn.Sequential(
+            OrderedDict(
+                [
+                    (
+                        name + "conv1",
+                        nn.Conv3d(
+                            in_channels=in_channels,
+                            out_channels=features,
+                            kernel_size=3,
+                            padding=1,
+                            bias=False,
+                        ),
+                    ),
+                    (name + "norm1", nn.BatchNorm3d(num_features=features)),
+                    (name + "tanh1", nn.Tanh()),
+                    (
+                        name + "conv2",
+                        nn.Conv3d(
+                            in_channels=features,
+                            out_channels=features,
+                            kernel_size=3,
+                            padding=1,
+                            bias=False,
+                        ),
+                    ),
+                    (name + "norm2", nn.BatchNorm3d(num_features=features)),
+                    (name + "tanh2", nn.Tanh()),
+                ]
+            )
+        )

model/uno.py

@@ -0,0 +1,278 @@
+# coding=utf-8
+# Codes for section: Results on Darcy Flow Equation
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from onescience.utils.pdenneval.integral_operators import *
+
+###############
+#  UNO^dagger achitechtures
+###############
+class UNO1d(nn.Module):
+    def __init__(self, num_channels, width, pad = 9, factor = 1, initial_step = 10):
+        super(UNO1d, self).__init__()
+
+        self.in_width = num_channels * initial_step + 1 # input channel
+        self.width = width 
+        
+        self.padding = pad  # pad the domain if input is non-periodic
+
+        self.fc_n1 = nn.Linear(self.in_width, self.width//2)
+
+        self.fc0 = nn.Linear(self.width//2, self.width) 
+
+        self.conv0 = OperatorBlock_1D(self.width, 2*factor*self.width,24, 11)
+
+        self.conv1 = OperatorBlock_1D(2*factor*self.width, 4*factor*self.width, 12, 5, Normalize = True)
+
+        self.conv2 = OperatorBlock_1D(4*factor*self.width, 4*factor*self.width, 12, 5)
+
+        self.conv4 = OperatorBlock_1D(4*factor*self.width, 2*factor*self.width, 24, 5, Normalize = True)
+
+        self.conv5 = OperatorBlock_1D(4*factor*self.width, self.width, 57, 11) # will be reshaped
+
+        self.fc1 = nn.Linear(2*self.width, 1*self.width)
+        self.fc2 = nn.Linear(1*self.width, num_channels)
+
+    def forward(self, x, grid):
+        x = torch.cat((x, grid), dim=-1)
+
+        x_fc_1 = self.fc_n1(x)
+        x_fc_1 = F.gelu(x_fc_1)
+
+        x_fc0 = self.fc0(x_fc_1)
+        x_fc0 = F.gelu(x_fc0)
+        
+        x_fc0 = x_fc0.permute(0, 2, 1)
+        # scale = math.ceil(x_fc0.shape[-1]/43)
+        x_fc0 = F.pad(x_fc0, [0,self.padding])
+        
+        D1 = x_fc0.shape[-1]
+        x_c0 = self.conv0(x_fc0,D1//2)
+
+        x_c1 = self.conv1(x_c0,D1//4)
+
+        x_c2 = self.conv2(x_c1,D1//4)
+        
+        x_c4 = self.conv4(x_c2 ,D1//2)
+        x_c4 = torch.cat([x_c4, x_c0], dim=1)
+
+        x_c5 = self.conv5(x_c4,D1)
+        x_c5 = torch.cat([x_c5, x_fc0], dim=1)
+
+
+        if self.padding!=0:
+            x_c5 = x_c5[..., :-self.padding]
+
+
+        x_c5 = x_c5.permute(0, 2, 1)
+
+        x_fc1 = self.fc1(x_c5)
+        x_fc1 = F.gelu(x_fc1)
+
+        x_out = self.fc2(x_fc1)
+        return x_out.unsqueeze(-2)
+
+class UNO2d(nn.Module):
+    def __init__(self, num_channels, width, pad = 6, factor = 1, initial_step = 10):
+        super(UNO2d, self).__init__()
+
+        self.in_width = num_channels * initial_step + 2 # input channel
+        self.width = width 
+        
+        self.padding = pad  # pad the domain if input is non-periodic
+
+        self.fc_n1 = nn.Linear(self.in_width, self.width//2)
+
+        self.fc0 = nn.Linear(self.width//2, self.width) 
+
+        self.conv0 = OperatorBlock_2D(self.width, 2*factor*self.width,36, 36, 17, 17)
+
+        self.conv1 = OperatorBlock_2D(2*factor*self.width, 4*factor*self.width, 18, 18, 7, 7, Normalize = True)
+
+        self.conv2 = OperatorBlock_2D(4*factor*self.width, 4*factor*self.width, 18, 18, 7, 7)
+
+        self.conv4 = OperatorBlock_2D(4*factor*self.width, 2*factor*self.width, 36, 36, 7, 7, Normalize = True)
+
+        self.conv5 = OperatorBlock_2D(4*factor*self.width, self.width, 81, 81, 17, 17) # will be reshaped
+
+        self.fc1 = nn.Linear(2*self.width, 1*self.width)
+        self.fc2 = nn.Linear(1*self.width, num_channels)
+
+    def forward(self, x, grid):
+        x = torch.cat((x, grid), dim=-1)
+
+        x_fc_1 = self.fc_n1(x)
+        x_fc_1 = F.gelu(x_fc_1)
+
+        x_fc0 = self.fc0(x_fc_1)
+        x_fc0 = F.gelu(x_fc0)
+        
+        x_fc0 = x_fc0.permute(0, 3, 1, 2)
+        
+        # scale = math.ceil(x_fc0.shape[-1]/85)
+        x_fc0 = F.pad(x_fc0, [0,self.padding, 0,self.padding])
+        
+        D1,D2 = x_fc0.shape[-2],x_fc0.shape[-1]
+
+        x_c0 = self.conv0(x_fc0,D1//2,D2//2)
+
+        x_c1 = self.conv1(x_c0,D1//4,D2//4)
+
+        x_c2 = self.conv2(x_c1,D1//4,D2//4)
+        
+        x_c4 = self.conv4(x_c2 ,D1//2,D2//2)
+        x_c4 = torch.cat([x_c4, x_c0], dim=1)
+
+        x_c5 = self.conv5(x_c4,D1,D2)
+        x_c5 = torch.cat([x_c5, x_fc0], dim=1)
+
+
+        if self.padding!=0:
+            x_c5 = x_c5[..., :-self.padding, :-self.padding]
+
+
+        x_c5 = x_c5.permute(0, 2, 3, 1)
+
+        x_fc1 = self.fc1(x_c5)
+        x_fc1 = F.gelu(x_fc1)
+
+        x_out = self.fc2(x_fc1)
+        
+        return x_out.unsqueeze(-2)
+
+class UNO3d(nn.Module):
+    def __init__(self, num_channels, width, pad = 5, factor = 1, initial_step = 10):
+        super(UNO3d, self).__init__()
+
+        self.in_width = num_channels * initial_step + 3 # input channel
+        self.width = width 
+        
+        self.padding = pad  # pad the domain if input is non-periodic
+
+        self.fc_n1 = nn.Linear(self.in_width, self.width//2)
+
+        self.fc0 = nn.Linear(self.width//2, self.width) 
+
+        self.conv0 = OperatorBlock_3D(self.width, 2*factor*self.width,36, 36, 36, 17, 17, 17)
+
+        self.conv1 = OperatorBlock_3D(2*factor*self.width, 4*factor*self.width, 18, 18, 18, 7, 7, 7, Normalize = True)
+
+        self.conv2 = OperatorBlock_3D(4*factor*self.width, 4*factor*self.width, 18, 18, 18, 7, 7, 7)
+
+        self.conv4 = OperatorBlock_3D(4*factor*self.width, 2*factor*self.width, 36, 36, 36, 7, 7, 7, Normalize = True)
+
+        self.conv5 = OperatorBlock_3D(4*factor*self.width, self.width, 81, 81, 81, 17, 17, 17) # will be reshaped
+
+        self.fc1 = nn.Linear(2*self.width, 1*self.width)
+        self.fc2 = nn.Linear(1*self.width, num_channels)
+
+    def forward(self, x, grid):
+        x = torch.cat((x, grid), dim=-1)
+
+        x_fc_1 = self.fc_n1(x)
+        x_fc_1 = F.gelu(x_fc_1)
+
+        x_fc0 = self.fc0(x_fc_1)
+        x_fc0 = F.gelu(x_fc0)
+        
+        x_fc0 = x_fc0.permute(0, 4, 1, 2, 3)
+        
+        # scale = math.ceil(x_fc0.shape[-1]/85)
+        x_fc0 = F.pad(x_fc0, [0,self.padding, 0,self.padding, 0,self.padding])
+        
+        D1,D2,D3 = x_fc0.shape[-3],x_fc0.shape[-2],x_fc0.shape[-1]
+
+        x_c0 = self.conv0(x_fc0,D1//2,D2//2,D3//2)
+
+        x_c1 = self.conv1(x_c0,D1//4,D2//4,D3//4)
+
+        x_c2 = self.conv2(x_c1,D1//4,D2//4,D3//4)
+        
+        x_c4 = self.conv4(x_c2 ,D1//2,D2//2,D3//2)
+        x_c4 = torch.cat([x_c4, x_c0], dim=1)
+
+        x_c5 = self.conv5(x_c4,D1,D2,D3)
+        x_c5 = torch.cat([x_c5, x_fc0], dim=1)
+
+
+        if self.padding!=0:
+            x_c5 = x_c5[..., :-self.padding, :-self.padding, :-self.padding]
+
+
+        x_c5 = x_c5.permute(0, 2, 3, 4, 1)
+
+        x_fc1 = self.fc1(x_c5)
+        x_fc1 = F.gelu(x_fc1)
+
+        x_out = self.fc2(x_fc1)
+        
+        return x_out.unsqueeze(-2)
+
+class UNO_maxwell(nn.Module):
+    def __init__(self, num_channels, width, pad = 5, factor = 1, initial_step = 10):
+        super(UNO_maxwell, self).__init__()
+
+        self.in_width = num_channels * initial_step # input channel
+        self.width = width 
+        
+        self.padding = pad  # pad the domain if input is non-periodic
+
+        self.fc_n1 = nn.Linear(self.in_width, self.width//2)
+
+        self.fc0 = nn.Linear(self.width//2, self.width) 
+
+        self.conv0 = OperatorBlock_3D(self.width, 2*factor*self.width,24, 24, 24, 10, 10, 10)
+
+        self.conv1 = OperatorBlock_3D(2*factor*self.width, 4*factor*self.width, 12, 12, 12, 5, 5, 5, Normalize = True)
+
+        self.conv2 = OperatorBlock_3D(4*factor*self.width, 4*factor*self.width, 12, 12, 12, 5, 5, 5)
+
+        self.conv4 = OperatorBlock_3D(4*factor*self.width, 2*factor*self.width, 12, 12, 12, 5, 5, 5, Normalize = True)
+
+        self.conv5 = OperatorBlock_3D(4*factor*self.width, self.width, 53, 53, 53, 10, 10, 10) # will be reshaped
+
+        self.fc1 = nn.Linear(2*self.width, 1*self.width)
+        self.fc2 = nn.Linear(1*self.width, num_channels)
+
+    def forward(self, x, grid):
+        x_fc_1 = self.fc_n1(x)
+        x_fc_1 = F.gelu(x_fc_1)
+
+        x_fc0 = self.fc0(x_fc_1)
+        x_fc0 = F.gelu(x_fc0)
+        
+        x_fc0 = x_fc0.permute(0, 4, 1, 2, 3)
+        
+        # scale = math.ceil(x_fc0.shape[-1]/85)
+        x_fc0 = F.pad(x_fc0, [0,self.padding, 0,self.padding, 0,self.padding])
+        
+        D1,D2,D3 = x_fc0.shape[-3],x_fc0.shape[-2],x_fc0.shape[-1]
+
+        x_c0 = self.conv0(x_fc0,D1//2,D2//2,D3//2)
+
+        x_c1 = self.conv1(x_c0,D1//4,D2//4,D3//4)
+
+        x_c2 = self.conv2(x_c1,D1//4,D2//4,D3//4)
+        
+        x_c4 = self.conv4(x_c2 ,D1//2,D2//2,D3//2)
+        x_c4 = torch.cat([x_c4, x_c0], dim=1)
+
+        x_c5 = self.conv5(x_c4,D1,D2,D3)
+        x_c5 = torch.cat([x_c5, x_fc0], dim=1)
+
+
+        if self.padding!=0:
+            x_c5 = x_c5[..., :-self.padding, :-self.padding, :-self.padding]
+
+
+        x_c5 = x_c5.permute(0, 2, 3, 4, 1)
+
+        x_fc1 = self.fc1(x_c5)
+        x_fc1 = F.gelu(x_fc1)
+
+        x_out = self.fc2(x_fc1)
+        
+        return x_out.unsqueeze(-2)
+    

scripts/common.py

@@ -0,0 +1,337 @@
+from __future__ import annotations
+
+import os
+import sys
+import warnings
+from pathlib import Path
+from types import SimpleNamespace
+from typing import Any, Optional
+
+import h5py
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, Dataset
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+DEFAULT_CONFIG = PROJECT_ROOT / "conf" / "config.yaml"
+
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+from model.fno import FNO1d, FNO2d, FNO3d, FNO_maxwell
+
+try:
+    from onescience.utils.YParams import YParams
+except ModuleNotFoundError as exc:
+    YParams = None
+    YPARAMS_IMPORT_ERROR = exc
+
+try:
+    from onescience.distributed.manager import DistributedManager
+except ModuleNotFoundError as exc:
+    DistributedManager = None
+    DISTRIBUTED_IMPORT_ERROR = exc
+
+try:
+    from onescience.datapipes.cfd.PDENNEval import PDEBenchFNODatapipe
+except (ImportError, ModuleNotFoundError, OSError) as exc:
+    PDEBenchFNODatapipe = None
+    DATAPIPE_IMPORT_ERROR = exc
+
+
+def get_attr(obj: Any, name: str, default: Any = None) -> Any:
+    if isinstance(obj, dict):
+        return obj.get(name, default)
+    return getattr(obj, name, default)
+
+
+def set_attr(obj: Any, name: str, value: Any) -> None:
+    if isinstance(obj, dict):
+        obj[name] = value
+    else:
+        setattr(obj, name, value)
+
+
+def to_namespace(value: Any) -> Any:
+    if isinstance(value, dict):
+        return SimpleNamespace(**{key: to_namespace(val) for key, val in value.items()})
+    if isinstance(value, list):
+        return [to_namespace(item) for item in value]
+    return value
+
+
+def load_config(config_path: Path | str = DEFAULT_CONFIG) -> Any:
+    config_path = Path(config_path).expanduser().resolve()
+    if YParams is not None:
+        return YParams(str(config_path), "fno_config")
+
+    import yaml
+
+    with config_path.open("r", encoding="utf-8") as handle:
+        raw = yaml.safe_load(handle)
+    if not isinstance(raw, dict) or "fno_config" not in raw:
+        raise ValueError(f"{config_path} must contain root key 'fno_config'")
+    return to_namespace(raw["fno_config"])
+
+
+def resolve_path(value: Any, root: Path = PROJECT_ROOT) -> Path:
+    if value is None:
+        raise ValueError("path value can not be null")
+    expanded = os.path.expandvars(os.path.expanduser(str(value)))
+    path = Path(expanded)
+    if not path.is_absolute():
+        path = root / path
+    return path.resolve()
+
+
+def prepare_config(
+    cfg: Any,
+    data_dir: Optional[str] = None,
+    output_dir: Optional[str] = None,
+    checkpoint: Optional[str] = None,
+) -> Any:
+    source = cfg.datapipe.source
+    training = cfg.training
+    inference = get_attr(cfg, "inference", None)
+
+    resolved_data = resolve_path(data_dir or source.data_dir)
+    set_attr(source, "data_dir", str(resolved_data))
+
+    resolved_output = resolve_path(output_dir or training.output_dir)
+    set_attr(training, "output_dir", str(resolved_output))
+
+    model_path = get_attr(training, "model_path", None)
+    if model_path:
+        set_attr(training, "model_path", str(resolve_path(model_path)))
+
+    if inference is not None:
+        infer_output = get_attr(inference, "output_dir", "./result/output")
+        infer_checkpoint = checkpoint or get_attr(inference, "checkpoint", None)
+        set_attr(inference, "output_dir", str(resolve_path(infer_output)))
+        if infer_checkpoint:
+            set_attr(inference, "checkpoint", str(resolve_path(infer_checkpoint)))
+
+    return cfg
+
+
+class SingleProcessManager:
+    rank = 0
+    local_rank = 0
+    world_size = 1
+    distributed = False
+    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+
+def initialize_distributed() -> Any:
+    if DistributedManager is None:
+        warnings.warn(
+            f"Falling back to a single-process manager: {DISTRIBUTED_IMPORT_ERROR}",
+            RuntimeWarning,
+        )
+        return SingleProcessManager()
+    DistributedManager.initialize()
+    return DistributedManager()
+
+
+def cleanup_distributed() -> None:
+    if DistributedManager is not None and DistributedManager.is_initialized():
+        DistributedManager.cleanup()
+
+
+class LocalPDEBenchFNODataset(Dataset):
+    """Small FNO-compatible reader used when optional OneScience graph deps are absent."""
+
+    def __init__(self, datapipe_cfg: Any, mode: str):
+        self.mode = mode
+        self.data_cfg = datapipe_cfg.data
+        self.source_cfg = datapipe_cfg.source
+        self.initial_step = int(self.data_cfg.initial_step)
+        self.reduced_resolution = int(self.data_cfg.reduced_resolution)
+        self.reduced_resolution_t = int(self.data_cfg.reduced_resolution_t)
+        self.reduced_batch = int(self.data_cfg.reduced_batch)
+        self.test_ratio = float(get_attr(self.data_cfg, "test_ratio", 0.1))
+        self.file_path = Path(self.source_cfg.data_dir) / self.source_cfg.file_name
+        self._load_single_file()
+
+    def _load_single_file(self) -> None:
+        if not self.file_path.is_file():
+            raise FileNotFoundError(f"HDF5 file not found: {self.file_path}")
+
+        with h5py.File(self.file_path, "r") as handle:
+            if "tensor" not in handle:
+                raise ValueError("local fallback datapipe requires HDF5 dataset 'tensor'")
+            tensor = np.asarray(handle["tensor"], dtype=np.float32)
+            if tensor.ndim == 3:
+                data = tensor[
+                    :: self.reduced_batch,
+                    :: self.reduced_resolution_t,
+                    :: self.reduced_resolution,
+                ]
+                data = np.transpose(data, (0, 2, 1))
+                self.data = data[:, :, :, None]
+                x = np.asarray(handle["x-coordinate"], dtype=np.float32)
+                self.grid = torch.tensor(
+                    x[:: self.reduced_resolution], dtype=torch.float32
+                ).unsqueeze(-1)
+            elif tensor.ndim == 4:
+                data = tensor[
+                    :: self.reduced_batch,
+                    :,
+                    :: self.reduced_resolution,
+                    :: self.reduced_resolution,
+                ]
+                data = np.transpose(data, (0, 2, 3, 1))
+                if "nu" in handle:
+                    nu = np.asarray(handle["nu"], dtype=np.float32)[
+                        :: self.reduced_batch,
+                        :: self.reduced_resolution,
+                        :: self.reduced_resolution,
+                    ]
+                    data = np.concatenate([nu[:, :, :, None], data], axis=-1)
+                self.data = data[:, :, :, :, None]
+                x = torch.tensor(np.asarray(handle["x-coordinate"], dtype=np.float32))
+                y = torch.tensor(np.asarray(handle["y-coordinate"], dtype=np.float32))
+                xx, yy = torch.meshgrid(x, y, indexing="ij")
+                self.grid = torch.stack((xx, yy), dim=-1)[
+                    :: self.reduced_resolution, :: self.reduced_resolution
+                ].float()
+            else:
+                raise ValueError(f"unsupported tensor ndim for FNO fallback: {tensor.ndim}")
+
+        sample_count = self.data.shape[0]
+        val_count = max(1, int(sample_count * self.test_ratio)) if sample_count > 1 else 0
+        if self.mode == "train":
+            self.data = self.data[val_count:]
+        else:
+            self.data = self.data[:val_count]
+        self.data = torch.tensor(self.data, dtype=torch.float32)
+        self.spatial_dim = len(self.data.shape) - 3
+
+    def __len__(self) -> int:
+        return len(self.data)
+
+    def __getitem__(self, idx: int):
+        sample = self.data[idx]
+        return sample[..., : self.initial_step, :], sample, self.grid
+
+
+class LocalPDEBenchFNODatapipe:
+    def __init__(self, cfg: Any, distributed: bool = False):
+        self.config = cfg
+        self.distributed = distributed
+        self.train_dataset = LocalPDEBenchFNODataset(cfg.datapipe, "train")
+        self.val_dataset = LocalPDEBenchFNODataset(cfg.datapipe, "val")
+        self.spatial_dim = self.train_dataset.spatial_dim
+
+    def train_dataloader(self):
+        loader_args = self.config.datapipe.dataloader
+        return DataLoader(
+            self.train_dataset,
+            batch_size=int(loader_args.batch_size),
+            num_workers=int(loader_args.num_workers),
+            pin_memory=bool(loader_args.pin_memory),
+            shuffle=True,
+            drop_last=True,
+        ), None
+
+    def val_dataloader(self):
+        loader_args = self.config.datapipe.dataloader
+        return DataLoader(
+            self.val_dataset,
+            batch_size=int(loader_args.batch_size),
+            num_workers=int(loader_args.num_workers),
+            pin_memory=bool(loader_args.pin_memory),
+            shuffle=False,
+            drop_last=False,
+        ), None
+
+
+def build_datapipe(cfg: Any, distributed: bool = False, force_local: bool = False) -> Any:
+    if PDEBenchFNODatapipe is not None and not force_local:
+        return PDEBenchFNODatapipe(cfg, distributed=distributed)
+    if PDEBenchFNODatapipe is None:
+        warnings.warn(
+            f"Using local FNO HDF5 datapipe because OneScience datapipe import failed: "
+            f"{DATAPIPE_IMPORT_ERROR}",
+            RuntimeWarning,
+        )
+    return LocalPDEBenchFNODatapipe(cfg, distributed=distributed)
+
+
+def build_model(spatial_dim: int, cfg: Any) -> torch.nn.Module:
+    model_args = cfg.model
+    data_cfg = cfg.datapipe.data
+    initial_step = int(data_cfg.initial_step)
+    pde_name = get_attr(data_cfg, "pde_name", "")
+    modes = int(model_args.modes)
+
+    if pde_name == "3D_Maxwell":
+        return FNO_maxwell(
+            num_channels=int(model_args.num_channels),
+            width=int(model_args.width),
+            modes1=modes,
+            modes2=modes,
+            modes3=modes,
+            initial_step=initial_step,
+        )
+    if spatial_dim == 1:
+        return FNO1d(
+            num_channels=int(model_args.num_channels),
+            width=int(model_args.width),
+            modes=modes,
+            initial_step=initial_step,
+        )
+    if spatial_dim == 2:
+        return FNO2d(
+            num_channels=int(model_args.num_channels),
+            width=int(model_args.width),
+            modes1=modes,
+            modes2=modes,
+            initial_step=initial_step,
+        )
+    if spatial_dim == 3:
+        return FNO3d(
+            num_channels=int(model_args.num_channels),
+            width=int(model_args.width),
+            modes1=modes,
+            modes2=modes,
+            modes3=modes,
+            initial_step=initial_step,
+        )
+    raise ValueError(f"unsupported spatial dimension: {spatial_dim}")
+
+
+def predict_batch(model: torch.nn.Module, x: torch.Tensor, y: torch.Tensor, grid: torch.Tensor, cfg: Any):
+    data_cfg = cfg.datapipe.data
+    train_cfg = cfg.training
+    initial_step = int(data_cfg.initial_step)
+    t_train = min(int(train_cfg.t_train), y.shape[-2])
+    input_shape = list(x.shape)[:-2] + [-1]
+
+    if get_attr(train_cfg, "training_type", "single") == "autoregressive":
+        pred = y[..., :initial_step, :]
+        for _ in range(initial_step, t_train):
+            model_input = x.reshape(input_shape)
+            model_output = model(model_input, grid)
+            if model_output.dim() == pred.dim() - 1:
+                model_output = model_output.unsqueeze(-2)
+            pred = torch.cat((pred, model_output), dim=-2)
+            x = torch.cat((x[..., 1:, :], model_output), dim=-2)
+        return pred, y[..., :t_train, :]
+
+    model_input = x.reshape(input_shape)
+    target = y[..., t_train - 1 : t_train, :]
+    pred = model(model_input, grid)
+    if pred.dim() == target.dim() - 1:
+        pred = pred.unsqueeze(-2)
+    return pred, target
+
+
+def load_model_state(path: Path, device: torch.device) -> dict[str, torch.Tensor]:
+    checkpoint = torch.load(path, map_location=device, weights_only=False)
+    if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint:
+        return checkpoint["model_state_dict"]
+    if isinstance(checkpoint, dict):
+        return checkpoint
+    raise ValueError(f"unsupported checkpoint format: {path}")

scripts/fake_data.py

@@ -0,0 +1,59 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import h5py
+import numpy as np
+
+from common import DEFAULT_CONFIG, load_config, prepare_config
+
+
+def generate_darcy_file(path: Path, samples: int, grid_size: int, seed: int, overwrite: bool) -> None:
+    if path.exists() and not overwrite:
+        print(f"Fake data already exists: {path}")
+        return
+
+    path.parent.mkdir(parents=True, exist_ok=True)
+    rng = np.random.default_rng(seed)
+    x = np.linspace(0.0, 1.0, grid_size, dtype=np.float32)
+    y = np.linspace(0.0, 1.0, grid_size, dtype=np.float32)
+    xx, yy = np.meshgrid(x, y, indexing="ij")
+
+    nu = rng.uniform(0.05, 0.25, size=(samples, grid_size, grid_size)).astype(np.float32)
+    tensor = np.empty((samples, 2, grid_size, grid_size), dtype=np.float32)
+    base = np.sin(np.pi * xx) * np.sin(np.pi * yy)
+    for idx in range(samples):
+        tensor[idx, 0] = base + 0.05 * rng.standard_normal(base.shape)
+        tensor[idx, 1] = tensor[idx, 0] * (1.0 + nu[idx]) + 0.01 * rng.standard_normal(base.shape)
+
+    with h5py.File(path, "w") as handle:
+        handle.create_dataset("tensor", data=tensor)
+        handle.create_dataset("nu", data=nu)
+        handle.create_dataset("x-coordinate", data=x)
+        handle.create_dataset("y-coordinate", data=y)
+        handle.create_dataset("t-coordinate", data=np.array([0.0, 1.0], dtype=np.float32))
+
+    print(f"Fake Darcy HDF5 written: {path}")
+    print(f"tensor={tensor.shape}, nu={nu.shape}")
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Generate a small PDENNEval FNO fake dataset.")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
+    parser.add_argument("--data-dir", default=None, help="Override datapipe.source.data_dir")
+    parser.add_argument("--samples", type=int, default=20)
+    parser.add_argument("--grid-size", type=int, default=16)
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--overwrite", action="store_true")
+    args = parser.parse_args()
+
+    cfg = prepare_config(load_config(args.config), data_dir=args.data_dir)
+    data_dir = Path(cfg.datapipe.source.data_dir)
+    file_path = data_dir / cfg.datapipe.source.file_name
+    generate_darcy_file(file_path, args.samples, args.grid_size, args.seed, args.overwrite)
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/inference.py

@@ -0,0 +1,83 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from common import (
+    DEFAULT_CONFIG,
+    build_datapipe,
+    build_model,
+    cleanup_distributed,
+    get_attr,
+    initialize_distributed,
+    load_config,
+    load_model_state,
+    predict_batch,
+    prepare_config,
+)
+
+
+@torch.no_grad()
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Run PDENNEval FNO inference.")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
+    parser.add_argument("--data-dir", default=None, help="Override datapipe.source.data_dir")
+    parser.add_argument("--checkpoint", default=None, help="Checkpoint path")
+    parser.add_argument("--output-dir", default=None, help="Directory for npz predictions")
+    parser.add_argument("--max-batches", type=int, default=None)
+    parser.add_argument("--force-local-datapipe", action="store_true")
+    args = parser.parse_args()
+
+    dist = initialize_distributed()
+    device = dist.device
+    cfg = prepare_config(load_config(args.config), data_dir=args.data_dir, checkpoint=args.checkpoint)
+    if args.output_dir:
+        cfg.inference.output_dir = str(Path(args.output_dir).expanduser().resolve())
+
+    checkpoint = Path(args.checkpoint or cfg.inference.checkpoint)
+    if not checkpoint.is_file():
+        raise FileNotFoundError(f"checkpoint not found: {checkpoint}")
+
+    datapipe = build_datapipe(
+        cfg,
+        distributed=False,
+        force_local=args.force_local_datapipe,
+    )
+    val_loader, _ = datapipe.val_dataloader()
+    model = build_model(datapipe.spatial_dim, cfg).to(device)
+    model.load_state_dict(load_model_state(checkpoint, device))
+    model.eval()
+
+    output_dir = Path(args.output_dir or cfg.inference.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    max_batches = args.max_batches or int(get_attr(cfg.inference, "max_batches", 1))
+
+    written = 0
+    for batch_idx, (x, y, grid) in enumerate(val_loader):
+        if batch_idx >= max_batches:
+            break
+        x = x.to(device)
+        y = y.to(device)
+        grid = grid.to(device)
+        pred, target = predict_batch(model, x, y, grid, cfg)
+        path = output_dir / f"batch_{batch_idx:04d}.npz"
+        np.savez_compressed(
+            path,
+            prediction=pred.detach().cpu().numpy(),
+            target=target.detach().cpu().numpy(),
+        )
+        print(f"wrote {path}")
+        written += 1
+
+    if written == 0:
+        raise RuntimeError("no validation batches were available for inference")
+
+    cleanup_distributed()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/preflight_check.py

@@ -0,0 +1,42 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import h5py
+
+from common import DEFAULT_CONFIG, load_config, prepare_config
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Preflight the PDENNEval standard package.")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
+    parser.add_argument("--data-dir", default=None, help="Override datapipe.source.data_dir")
+    args = parser.parse_args()
+
+    cfg = prepare_config(load_config(args.config), data_dir=args.data_dir)
+    data_path = Path(cfg.datapipe.source.data_dir) / cfg.datapipe.source.file_name
+    if not data_path.is_file():
+        raise FileNotFoundError(f"missing HDF5 data file: {data_path}")
+
+    with h5py.File(data_path, "r") as handle:
+        for key in ("tensor", "nu", "x-coordinate", "y-coordinate"):
+            if key not in handle:
+                raise ValueError(f"{data_path.name} missing required key: {key}")
+        if handle["tensor"].ndim != 4:
+            raise ValueError("2D Darcy fake/default config expects tensor ndim=4")
+        if handle["nu"].ndim != 3:
+            raise ValueError("2D Darcy fake/default config expects nu ndim=3")
+
+    output_dir = Path(cfg.training.output_dir)
+    output_dir.mkdir(parents=True, exist_ok=True)
+    probe = output_dir / ".preflight_write_test"
+    probe.write_text("ok", encoding="utf-8")
+    probe.unlink()
+    print(f"[OK] data file: {data_path}")
+    print(f"[OK] output directory writable: {output_dir}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/result.py

@@ -0,0 +1,50 @@
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+from common import DEFAULT_CONFIG, load_config, prepare_config
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Summarize PDENNEval FNO inference outputs.")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
+    parser.add_argument("--prediction-dir", default=None)
+    parser.add_argument("--output-dir", default=None)
+    args = parser.parse_args()
+
+    cfg = prepare_config(load_config(args.config))
+    prediction_dir = Path(args.prediction_dir or cfg.result.prediction_dir).expanduser().resolve()
+    output_dir = Path(args.output_dir or cfg.result.output_dir).expanduser().resolve()
+    files = sorted(prediction_dir.glob("*.npz"))
+    if not files:
+        raise FileNotFoundError(f"no prediction files found in {prediction_dir}")
+
+    mse_values = []
+    mae_values = []
+    for path in files:
+        data = np.load(path)
+        pred = data["prediction"]
+        target = data["target"]
+        diff = pred - target
+        mse_values.append(float(np.mean(diff ** 2)))
+        mae_values.append(float(np.mean(np.abs(diff))))
+
+    metrics = {
+        "num_files": len(files),
+        "mse": float(np.mean(mse_values)),
+        "mae": float(np.mean(mae_values)),
+    }
+    output_dir.mkdir(parents=True, exist_ok=True)
+    metrics_path = output_dir / "metrics.json"
+    metrics_path.write_text(json.dumps(metrics, indent=2), encoding="utf-8")
+    print(json.dumps(metrics, indent=2))
+    print(f"wrote {metrics_path}")
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/smoke_models.py

@@ -0,0 +1,138 @@
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+from types import SimpleNamespace
+
+import torch
+
+
+PROJECT_ROOT = Path(__file__).resolve().parents[1]
+if str(PROJECT_ROOT) not in sys.path:
+    sys.path.insert(0, str(PROJECT_ROOT))
+
+
+def check(name: str, tensor: torch.Tensor, expected_last_dim: int | None = None) -> None:
+    if not torch.isfinite(tensor).all():
+        raise RuntimeError(f"{name} produced non-finite values")
+    if expected_last_dim is not None and tensor.shape[-1] != expected_last_dim:
+        raise RuntimeError(
+            f"{name} last dim expected {expected_last_dim}, got {tensor.shape[-1]}"
+        )
+    print(f"[OK] {name}: shape={tuple(tensor.shape)}")
+
+
+@torch.no_grad()
+def smoke_deeponet() -> None:
+    from model.deeponet import DeepONetCartesianProd1D
+
+    model = DeepONetCartesianProd1D(
+        size=8,
+        in_channel_branch=1,
+        query_dim=1,
+        out_channel=1,
+        activation="relu",
+        base_model="MLP",
+    ).eval()
+    x_func = torch.randn(2, 8, 1)
+    x_loc = torch.linspace(0.0, 1.0, 8).unsqueeze(-1)
+    out = model((x_func, x_loc))
+    check("deeponet.DeepONetCartesianProd1D", out, expected_last_dim=1)
+
+
+@torch.no_grad()
+def smoke_fno() -> None:
+    from model.fno import FNO2d
+
+    model = FNO2d(num_channels=1, modes1=4, modes2=4, width=8, initial_step=1).eval()
+    x = torch.randn(2, 16, 16, 1)
+    coords = torch.linspace(0.0, 1.0, 16)
+    xx, yy = torch.meshgrid(coords, coords, indexing="ij")
+    grid = torch.stack((xx, yy), dim=-1).unsqueeze(0).repeat(2, 1, 1, 1)
+    out = model(x, grid)
+    check("fno.FNO2d", out, expected_last_dim=1)
+
+
+@torch.no_grad()
+def smoke_mpnn() -> None:
+    from torch_geometric.data import Data
+
+    from model.mpnn import MPNN
+
+    pde = SimpleNamespace(
+        name="smoke",
+        tmin=0.0,
+        tmax=1.0,
+        resolution_t=10,
+        spatial_domain=[(0.0, 1.0)],
+        resolution=[4],
+        spatial_dim=1,
+        variables={},
+    )
+    model = MPNN(pde=pde, time_window=10, hidden_features=128, hidden_layers=1).eval()
+    edge_index = torch.tensor(
+        [[0, 1, 2, 3, 1, 2, 3, 0], [1, 2, 3, 0, 0, 1, 2, 3]],
+        dtype=torch.long,
+    )
+    data = Data(
+        x=torch.randn(4, 10, 1),
+        x_pos=torch.linspace(0.0, 1.0, 4).unsqueeze(-1),
+        t_pos=torch.zeros(4),
+        edge_index=edge_index,
+        batch=torch.zeros(4, dtype=torch.long),
+        variables=torch.empty(4, 0),
+    )
+    out = model(data, v=0)
+    check("mpnn.MPNN", out)
+
+
+@torch.no_grad()
+def smoke_pino_fno() -> None:
+    from model.pino_fno import FNO1d
+
+    model = FNO1d(
+        modes=[4, 4],
+        width=4,
+        layers=[4, 4, 4],
+        fc_dim=8,
+        in_dim=2,
+        out_dim=1,
+        pad_ratio=[0.0, 0.0],
+    ).eval()
+    out = model(torch.randn(2, 16, 2))
+    check("pino_fno.FNO1d", out, expected_last_dim=1)
+
+
+@torch.no_grad()
+def smoke_unet() -> None:
+    from model.unet import UNet2d
+
+    model = UNet2d(in_channels=1, out_channels=1, init_features=2).eval()
+    out = model(torch.randn(1, 1, 16, 16))
+    check("unet.UNet2d", out)
+
+
+@torch.no_grad()
+def smoke_uno() -> None:
+    from model.uno import UNO1d
+
+    model = UNO1d(num_channels=1, width=4, initial_step=1).eval()
+    x = torch.randn(1, 64, 1)
+    grid = torch.linspace(0.0, 1.0, 64).view(1, 64, 1)
+    out = model(x, grid)
+    check("uno.UNO1d", out, expected_last_dim=1)
+
+
+def main() -> int:
+    torch.manual_seed(0)
+    smoke_deeponet()
+    smoke_fno()
+    smoke_mpnn()
+    smoke_pino_fno()
+    smoke_unet()
+    smoke_uno()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

scripts/train.py

@@ -0,0 +1,143 @@
+from __future__ import annotations
+
+import argparse
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from common import (
+    DEFAULT_CONFIG,
+    build_datapipe,
+    build_model,
+    cleanup_distributed,
+    get_attr,
+    initialize_distributed,
+    load_config,
+    load_model_state,
+    predict_batch,
+    prepare_config,
+)
+
+
+def train_one_epoch(model, loader, optimizer, loss_fn, device, cfg) -> float:
+    model.train()
+    losses = []
+    for x, y, grid in loader:
+        x = x.to(device)
+        y = y.to(device)
+        grid = grid.to(device)
+        pred, target = predict_batch(model, x, y, grid, cfg)
+        if pred.numel() == 0:
+            continue
+        loss = loss_fn(pred.reshape(pred.size(0), -1), target.reshape(target.size(0), -1))
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+        losses.append(float(loss.item()))
+    return float(np.mean(losses)) if losses else 0.0
+
+
+@torch.no_grad()
+def validate(model, loader, loss_fn, device, cfg) -> float:
+    model.eval()
+    losses = []
+    for x, y, grid in loader:
+        x = x.to(device)
+        y = y.to(device)
+        grid = grid.to(device)
+        pred, target = predict_batch(model, x, y, grid, cfg)
+        if pred.numel() == 0:
+            continue
+        loss = loss_fn(pred.reshape(pred.size(0), -1), target.reshape(target.size(0), -1))
+        losses.append(float(loss.item()))
+    return float(np.mean(losses)) if losses else 0.0
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(description="Train the standard PDENNEval FNO package.")
+    parser.add_argument("--config", default=str(DEFAULT_CONFIG), help="Path to conf/config.yaml")
+    parser.add_argument("--data-dir", default=None, help="Override datapipe.source.data_dir")
+    parser.add_argument("--output-dir", default=None, help="Override training.output_dir")
+    parser.add_argument("--force-local-datapipe", action="store_true")
+    args = parser.parse_args()
+
+    dist = initialize_distributed()
+    device = dist.device
+    cfg = prepare_config(load_config(args.config), data_dir=args.data_dir, output_dir=args.output_dir)
+
+    seed = int(get_attr(cfg.training, "seed", 0))
+    torch.manual_seed(seed + dist.rank)
+    np.random.seed(seed + dist.rank)
+
+    output_dir = Path(cfg.training.output_dir)
+    if dist.rank == 0:
+        output_dir.mkdir(parents=True, exist_ok=True)
+        print(f"Output directory: {output_dir}")
+        print(f"Device: {device}, world_size={dist.world_size}")
+
+    datapipe = build_datapipe(
+        cfg,
+        distributed=(dist.world_size > 1),
+        force_local=args.force_local_datapipe,
+    )
+    train_loader, train_sampler = datapipe.train_dataloader()
+    val_loader, _ = datapipe.val_dataloader()
+
+    model = build_model(datapipe.spatial_dim, cfg).to(device)
+    if dist.world_size > 1:
+        model = DDP(model, device_ids=[dist.local_rank], output_device=dist.local_rank)
+
+    optimizer_cfg = cfg.training.optimizer
+    scheduler_cfg = cfg.training.scheduler
+    optimizer = getattr(torch.optim, optimizer_cfg.name)(
+        model.parameters(),
+        lr=float(optimizer_cfg.lr),
+        weight_decay=float(optimizer_cfg.weight_decay),
+    )
+    scheduler = getattr(torch.optim.lr_scheduler, scheduler_cfg.name)(
+        optimizer,
+        step_size=int(scheduler_cfg.step_size),
+        gamma=float(scheduler_cfg.gamma),
+    )
+    loss_fn = nn.MSELoss()
+
+    if bool(get_attr(cfg.training, "continue_training", False)) and cfg.training.model_path:
+        model_to_load = model.module if hasattr(model, "module") else model
+        model_to_load.load_state_dict(load_model_state(Path(cfg.training.model_path), device))
+
+    best_val = float("inf")
+    epochs = int(cfg.training.epochs)
+    save_period = max(1, int(cfg.training.save_period))
+
+    for epoch in range(epochs):
+        if train_sampler is not None:
+            train_sampler.set_epoch(epoch)
+        train_loss = train_one_epoch(model, train_loader, optimizer, loss_fn, device, cfg)
+        scheduler.step()
+        val_loss = validate(model, val_loader, loss_fn, device, cfg)
+
+        if dist.rank == 0:
+            print(f"epoch={epoch} train_loss={train_loss:.6e} val_loss={val_loss:.6e}")
+            model_to_save = model.module if hasattr(model, "module") else model
+            state = {
+                "epoch": epoch,
+                "model_state_dict": model_to_save.state_dict(),
+                "optimizer_state_dict": optimizer.state_dict(),
+                "loss": val_loss,
+            }
+            torch.save(state, output_dir / "latest_model.pt")
+            if (epoch + 1) % save_period == 0:
+                torch.save(state, output_dir / f"model_epoch_{epoch}.pt")
+            if val_loss <= best_val:
+                best_val = val_loss
+                torch.save(state, output_dir / "best_model.pt")
+
+    cleanup_distributed()
+    return 0
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

weight/.gitkeep

@@ -0,0 +1 @@
+