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.
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:
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
1modelscope download --model OneScience/MP_PDE --local_dir ./MP_PDE
2cd MP_PDE
Set Up the Runtime Environment
DCU environment
1# Activate DTK and Conda first
2conda create -n onescience311 python=3.11 -y
3conda activate onescience311
4# Installation with uv is also supported
5pip install onescience[cfd-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
GPU environment
1# Activate Conda first
2conda create -n onescience311 python=3.11 -y libstdcxx-ng=12 libgcc-ng=12 gcc_linux-64=12 gxx_linux-64=12
3conda activate onescience311
4# Installation with uv is also supported
5pip 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\left(
-\frac{\lVert x_1-x_2\rVert^2}{2l^2}
\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:
$$
\left(u,\ y,\ G(u)(y)\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:
1python scripts/train.py \
2 --config config/config.yaml \
3 --generate-data \
4 --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.
1python scripts/inference.py \
2 --config config/config.yaml \
3 --device auto
Evaluation and Visualization
After training and inference, generate prediction visualizations with:
1python scripts/result.py \
2 --config config/config.yaml
Official OneScience Resources
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.
