OneScience-Group/DeepONet

DeepONet

Model Introduction

DeepONet is a deep neural network for operator learning proposed by a research team affiliated with Brown University and published in Nature Machine Intelligence in 2021. Unlike conventional neural networks, which primarily learn mappings between finite-dimensional vectors, DeepONet learns nonlinear operator mappings directly between function spaces, establishing a relationship between an input function and its corresponding output function. It can approximate solution operators defined by ordinary differential equations, partial differential equations, and other physical systems, providing an efficient data-driven modeling approach for complex dynamical systems and scientific computing.

Using the OneScience skill workflow, this project independently reproduces experiments related to the DeepONet paper.

Paper: DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators

Model Description

DeepONet uses a dual-network architecture consisting of a Branch Net and a Trunk Net. The Branch Net encodes discrete samples of an input function at fixed sensor locations, while the Trunk Net encodes the spatial or spatiotemporal coordinates at which predictions are requested. The output of the target operator at a specified location is obtained from the inner product of the two feature vectors plus a bias term.

Use Cases

Use case	Description
Operator learning	Learns mappings between function spaces, directly predicting an output function (G(u)) from an input function (u). This is useful for function-to-function mappings that conventional neural networks cannot easily handle.
Spatiotemporal field prediction	Provides fast surrogate predictions for fluid problems such as Navier–Stokes and compressible Euler equations.
Multiscale physical-field modeling	The Trunk Net can directly accept multidimensional coordinates such as ((x,t)), making it suitable for predicting temperature, concentration, diffusion, and other fields that vary in space and time.
Multiple query-point prediction	For a fixed input function, solutions at different spatial or temporal locations can be predicted by changing only the query coordinates supplied to the Trunk Net.

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/DeepONet --local_dir ./DeepONet
cd DeepONet

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. The E3 data is generated on the fly by models/dataset.py according to config/config.yaml and corresponds to the following parameterized one-dimensional PDE:

$$\frac{\mathrm{\partial }u}{\mathrm{\partial }t}+\alpha \frac{\mathrm{\partial }(u^2)}{\mathrm{\partial }x}-\beta \frac{{\mathrm{\partial }}^{2}u}{\mathrm{\partial }{x}^2}+\gamma \frac{{\mathrm{\partial }}^{3}u}{\mathrm{\partial }{x}^3}=\delta (t,x),\backslash frac\{\backslash partial\; u\}\{\backslash partial\; t\}\; +\backslash alpha\backslash frac\{\backslash partial(u^2)\}\{\backslash partial\; x\}\; -\backslash beta\backslash frac\{\backslash partial^2u\}\{\backslash partial\; x^2\}\; +\backslash gamma\backslash frac\{\backslash partial^3u\}\{\backslash partial\; x^3\}\; =\backslash delta(t,x),$$

where $x\in[0,16)$ and $t\in[0,4]$, with periodic boundary conditions in the spatial dimension. The equation parameters are sampled independently:

$$\alpha \sim \mathcal{U}(0,3),\beta \sim \mathcal{U}(0,0.4),\gamma \sim \mathcal{U}(0,1)\mathrm{.}\backslash alpha\backslash sim\backslash mathcal\; U(0,3),\backslash quad\; \backslash beta\backslash sim\backslash mathcal\; U(0,0.4),\backslash quad\; \backslash gamma\backslash sim\backslash mathcal\; U(0,1).$$

The forcing term and initial condition are defined as:

$$\delta (t,x)=\sum _{j=1}^5{A}_j\sin ({\omega }_{j}t+\frac{2\pi k_{j}x}{16}+{\varphi }_j),u(0,x)=\delta (0,x),\backslash delta(t,x)=\backslash sum\_\{j=1\}^\{5\}\; A\_j\backslash sin\backslash left(\; \backslash omega\_jt+\backslash frac\{2\backslash pi\; k\_jx\}\{16\}+\backslash phi\_j\; \backslash right),\; \backslash qquad\; u(0,x)=\backslash delta(0,x),$$

where:

• $A_j\sim\mathcal U(-0.5,0.5)$;
• $\omega_j=-0.4$;
• $k_j\in{1,2,3}$;
• $\phi_j\sim\mathcal U(0,2\pi)$.

Reference solutions are first generated on 200 spatial grid points and then downsampled to 100 points. The nonlinear flux is discretized using a fifth-order WENO scheme, and time integration uses a fourth-order Runge–Kutta method. Each trajectory contains 250 time points and has shape:

$$u\in {\mathbb{R}}^{250\times 100}\mathrm{.}u\backslash in\backslash mathbb\; R^\{250\backslash times100\}.$$

The model uses a history window of length $K=25$ to predict the next 25 time steps. One supervised sample can be written as:

$$(u_{i:i+K-1},x,t,(\alpha ,\beta ,\gamma );u_{i+K:i+2K-1})\mathrm{.}\backslash left(\; u\_\{i:i+K-1\},\backslash ,x,\backslash ,t,\backslash ,(\backslash alpha,\backslash beta,\backslash gamma);\; \backslash \; u\_\{i+K:i+2K-1\}\; \backslash right).$$

Training

The default configuration reproduces four operator-learning experiments from the DeepONet paper: antiderivative, nonlinear ODE, forced pendulum, and diffusion–reaction equation. By default, the training script runs the antiderivative experiment, with input functions and reference solutions generated on the fly by models/dataset.py.

Run the default antiderivative experiment:

python scripts/train.py \
    --config config/config.yaml \
    --experiment antiderivative \
    --device auto

Set --experiment to all to run all four main experiments from the paper sequentially.

During training, the training loss, test MSE, and relative L2 error are printed at the intervals configured in config/config.yaml. When the test MSE improves, the model weights, optimizer state, current iteration, evaluation metrics, and effective runtime configuration are saved to weight/best_model.pth.

Trained Weights

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

Inference

Before running inference, make sure the configured data path is valid and weight/best_model.pth exists.

python scripts/inference.py \
    --config config/config.yaml \
    --experiment <experiment-name> \
    --variant unstacked_bias \
    --mode <inference-mode> \
    --device auto

Available inference modes are:

• random_test: random test set;
• ood: out-of-distribution ODE inputs;
• pde_grid: two-dimensional spatiotemporal field for the diffusion–reaction equation.

The default batch size is 8,192 and can be changed with --batch-size. Predictions and evaluation metrics are saved to the results directory.

Evaluation and Visualization

After training and inference, summarize existing experiment results and generate training curves, prediction comparisons, and an evaluation report 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: DeepONet: Learning nonlinear operators for identifying differential equations based on the universal approximation theorem of operators
• This project is an independent reproduction of the DeepONet paper. The official implementation is licensed under the MIT License. The project code, model weights, training data, and third-party dependencies remain subject to their respective license terms.