pde-autoregressive-solver

PDE SolversNeural-Solver-Libraryrigorous codebase

Description

Autoregressive PDE Solving: Custom Neural Operator Design

Objective

Design and implement a custom neural operator for autoregressive time-dependent PDE solving. 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)
fx: input function values (previous timesteps), shape (B, N, fun_dim) where fun_dim=10
T: time embedding (unused for autoregressive, always None)
geo: geometry info (unused, always None)
output: predicted next timestep, shape (B, N, out_dim) where out_dim=1

The model is called autoregressively: at each step, the oldest input channel is dropped, and the new prediction is appended as the newest channel in fx.

Key args attributes: n_hidden, n_layers, n_heads, space_dim, fun_dim, out_dim, act, mlp_ratio, dropout, geotype (structured_2D or structured_1D), shapelist, unified_pos, ref, slice_num, modes.

Evaluation

Trained and evaluated on four autoregressive PDE benchmarks (relative L2 error, lower is better):

NS (2D Navier-Stokes vorticity, structured_2D, 64x64, viscosity=1e-5, 10-step rollout)
DiffSorp (1D diffusion-sorption, structured_1D, PDEBench, 10-step rollout)
Advection (1D advection equation, structured_1D, PDEBench, beta=1.0, 10-step rollout)
Burgers (1D Burgers equation, structured_1D, PDEBench, viscosity=1e-3, 10-step rollout)

Your model must handle both structured_1D and structured_2D geometry types via args.geotype.

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 NS ↑	rel err DiffSorp ↑	rel err Advection ↑	rel err Burgers ↑
fno	baseline	0.089	0.023	0.071	0.100
transformer	baseline	0.063	0.029	0.886	0.920
transolver	baseline	0.132	0.028	7.858	0.488