OneScience-Group/MP_PDE

MP_PDE

Model Introduction

MP-PDE (Message Passing Neural PDE Solver) is a neural PDE solver based on message passing, proposed by researchers from the University of Amsterdam and other institutions and published at ICLR 2022. It targets the spatiotemporal prediction of parameterized partial differential equations (PDEs). By learning local interactions between spatial nodes, it predicts numerical PDE solutions and adapts well to different equation parameters, spatial resolutions, and mesh structures.

Using the OneScience skill workflow, this project independently reproduces the E3 parameterized PDE prediction experiment from the MP-PDE paper to evaluate spatiotemporal evolution predictions under different PDE parameter settings.

Paper: Message Passing Neural PDE Solvers

Model Description

MP-PDE uses an encoder-processor-decoder architecture consisting of an MLP encoder, a multilayer message-passing graph neural network (GNN processor), and a one-dimensional CNN decoder. It encodes PDE states from a historical time window together with spatial coordinates and equation parameters as node features. Multiple message-passing layers learn local spatial relationships, while temporal bundling predicts several future time steps at once for autoregressive spatiotemporal PDE forecasting.

Use Cases

Use case	Description
Cross-resolution PDE prediction	Graph structures describe relationships between spatial nodes with less dependence on fixed convolutional grids, supporting training and generalization across spatial discretization resolutions.
Irregular-mesh computation	Message passing uses relative spatial positions between nodes and can handle nonuniform or irregular spatial discretizations, as demonstrated in the paper's wave-equation experiment.
Long-horizon autoregressive prediction	Techniques such as temporal bundling and the pushforward trick reduce error accumulation during autoregressive prediction and improve the stability of long PDE rollouts.

Usage

1. Using OneCode

Try intelligent, one-click AI4S programming in the OneCode online environment:

Try intelligent, one-click AI4S programming

2. Manual Installation and Usage

Hardware requirements

• A GPU or DCU is recommended.
• A CPU can be used for import checks and small-scale connectivity tests, 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 current cluster, is recommended.

Download the Model Package

modelscope download --model OneScience/MP_PDE --local_dir ./MP_PDE
cd MP_PDE

Set Up the Runtime Environment

DCU environment

# 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

# 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

This project does not depend on an external dataset. Training data is synthesized on the fly by models/dataset.py according to config.yaml. The default input function $u$ is sampled from a zero-mean Gaussian random field (GRF) with covariance kernel:

$$k_l(x_1,x_2)=\exp (-\frac{\parallel x_1-x_2{\parallel }^2}{2l^2}),k\_l(x\_1,x\_2)\; =\; \backslash exp\backslash left(\; -\backslash frac\{\backslash lVert\; x\_1-x\_2\backslash rVert^2\}\{2l^2\}\; \backslash right),$$

where the correlation length is $l=0.2$. Input functions are first generated on a fine grid of 1,000 points, then sampled at 100 equally spaced sensor locations using cubic interpolation. These discrete samples are used as the Branch Net input in MP_PDE.

Each supervised-learning sample is represented as:

$$(u,y,G(u)(y)),\backslash left(u,\backslash \; y,\backslash \; G(u)(y)\backslash right),$$

where:

• $u$: discrete values of the input function at 100 sensor locations;
• $y$: the query location to predict—one-dimensional coordinates for ODE experiments and two-dimensional spatiotemporal coordinates $(x,t)$ for PDE experiments;
• $G(u)(y)$: the reference solution computed by the corresponding differential-equation solver, representing the output of operator $G$ applied to input function $u$ at location $y$.

Training

The default configuration reproduces the E3 parameterized PDE experiment from the MP-PDE paper. Training data is generated on the fly by models/dataset.py according to config/config.yaml, and the model uses a historical window of 25 time steps to predict the next 25 steps.

Run E3 training with:

python scripts/train.py \
    --config config/config.yaml \
    --generate-data \
    --device auto

Trained Weights

weight/best_model.pth contains the best weights from the full training run and can be used directly for inference or fine-tuning.

Inference

Before running inference, make sure the E3 dataset data/e3.h5 and trained weights weight/best_model.pth exist.

python scripts/inference.py \
    --config config/config.yaml \
    --device auto

Evaluation and Visualization

After training and inference, generate prediction visualizations with:

python scripts/result.py \
    --config config/config.yaml

Official OneScience Resources

Platform	OneScience Main 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

Citation and License

• Original paper: Message Passing Neural PDE Solvers
• This project is an independent reproduction of the MP_PDE paper. The project code, model weights, training data, and third-party dependencies remain subject to their respective license terms.