pde-steady-solver

PDE SolversNeural-Solver-Libraryrigorous codebase

Description

Steady-State PDE Solving: Custom Neural Operator Design

Objective

Design and implement a custom neural operator for steady-state PDE solving on structured 2D meshes. Your code goes in the Model class in models/Custom.py. Three reference implementations (FNO, Transformer, Transolver) are provided as read-only context.

Model Interface

Your model receives args at initialization and must implement:

forward(self, x, fx, T=None, geo=None) -> output

x: spatial coordinates, shape (B, N, space_dim) where space_dim=2
fx: input function values, shape (B, N, fun_dim) — can be None if fun_dim=0
T: time embedding (unused for steady tasks, always None)
geo: geometry info (unused, always None)
output: predicted solution, shape (B, N, out_dim)

Key args attributes: n_hidden, n_layers, n_heads, space_dim, fun_dim, out_dim, act, mlp_ratio, dropout, geotype (always structured_2D), shapelist (grid dimensions [H, W]), unified_pos, ref, slice_num, modes.

Evaluation

Trained and evaluated on three structured 2D PDE benchmarks (relative L2 error, lower is better):

Airfoil (NACA airflow, fun_dim=2, out_dim=1, 221x51 grid)
Darcy (Darcy flow, fun_dim=1, out_dim=1, 85x85 grid after downsampling)
Pipe (pipe flow, fun_dim=2, out_dim=1, 129x129 grid)

Training epochs vary per dataset: Airfoil (100), Darcy (300), Pipe (50), all with OneCycleLR scheduler.

Code

Custom.py

EditableRead-only

1import torch
2import torch.nn as nn
3import numpy as np
4from timm.models.layers import trunc_normal_
5from layers.Basic import MLP
6from layers.Embedding import timestep_embedding, unified_pos_embedding
7
8
9class Model(nn.Module):
10    def __init__(self, args):
11        super(Model, self).__init__()
12        self.__name__ = 'Custom'
13        self.args = args
14        ## embedding
15        if args.unified_pos and args.geotype != 'unstructured':

Additional context files (read-only):

Neural-Solver-Library/models/Transolver.py
Neural-Solver-Library/models/FNO.py
Neural-Solver-Library/models/Transformer.py
Neural-Solver-Library/layers/Basic.py
Neural-Solver-Library/layers/Embedding.py
Neural-Solver-Library/layers/Physics_Attention.py
Neural-Solver-Library/layers/FNO_Layers.py

Results

Show per-seed results

Model	Type	rel err Airfoil ↓	rel err Darcy ↓	rel err Pipe ↓
fno	baseline	0.006	0.010	0.007
transformer	baseline	0.015	0.004	0.008
transolver	baseline	0.005	0.006	0.004