DeepONet for neural operator learning

I implemented DeepONet in Julia and tested it on an integration operator and two PDEs. DeepONet learns an operator, a map from an input function to an output function, rather than a map between finite-dimensional vectors.

One difference from the original architecture: I added a few dense layers after the branch and trunk outputs are combined, to give the model more capacity.

neural operator learning

Operator learning approximates mappings between function spaces, which are infinite-dimensional. This is the natural setting for PDEs, where both the input (an initial condition, a coefficient field) and the output (the solution) are functions.

universal approximation theorem for operators

From the DeepONet paper: … \(G\) is a nonlinear continuous operator. Then for any \(\epsilon>0\), there exist parameters such that the following holds for all \(u\) and \(y\):

\[ \left|G(u)(y) - \sum_{k=1}^p \underbrace{\sum_{i=1}^n c_i^k \sigma\left(\sum_{j=1}^m \xi_{ij}^ku(x_j)+\theta_i^k\right)}_\text{branch} \underbrace{\sigma(w_k \cdot y+\zeta_k)}_\text{trunk} \right|<\epsilon \]

DeepONet architecture

DeepONet has two subnetworks:

Branch network: encodes the input function, sampled at fixed sensor points.
Trunk network: encodes the coordinates at which the output function is evaluated.

Their outputs are combined to approximate the target operator.

installation

The package can be installed from GitHub:

pkg> add https://github.com/chutommy/DeepONet.jl

usage

A minimal training setup. The full examples below include the data loaders and model settings.

model = DeepONetModel( M, 2, 20, activations;
    branch_sizes = branch_sizes,
    trunk_sizes = trunk_sizes,
    output_sizes = output_sizes,
)

opt_state = Flux.setup(Flux.AdamW(0.0003), model)
train!(model, opt_state, train_loader, test_loader)