deeponet for neural operator learning

DeepONet is a neural network architecture for operator learning, mapping functions to functions. it’s especially useful for infinite-dimensional problems like PDEs and scientific simulations where a standard network doesn’t quite fit.

this Julia implementation adds a few extra layers after combining the branch and trunk outputs, just to give the model a bit more flexibility.

neural operator learning

operator learning uses machine learning to approximate mathematical operators. unlike finite-dimensional ML tasks, it works with transformations in infinite-dimensional spaces, which is exactly what you need for PDEs and other function-to-function problems.

universal approximation theorem for operators, in a nutshell

https://arxiv.org/pdf/1910.03193: … \(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 main parts.

branch network: processes the input function.
trunk network: processes spatial or temporal coordinates.

the two outputs are combined to approximate the target operator.

installation

install the package like this:

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

usage

here is a minimal setup:

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)

there are more concrete examples below.

examples

integration on \([0, 1]\)

train DeepONet to approximate the integration operator over \([0, 1]\):

\[ F(x) = \int_0^1 f(x) \, dx \]

the network learns to integrate a given function \( f(x) \) over this interval, and we compare the predictions against the exact values.

code for this example is here.

Burger’s equation

approximate the solution to the 1D Burger’s equation:

\[ \frac{\partial u}{\partial t} + u \frac{\partial u}{\partial x} = \nu \frac{\partial^2 u}{\partial x^2} \]

DeepONet is trained on initial conditions to predict the solution over time.

code for this example is here.

Darcy flow

learn Darcy’s law for fluid flow in porous media:

\[ \mathbf{v} = -\frac{\kappa}{\mu} \nabla p \]

given inputs like permeability (\(\kappa\)) and pressure gradient (\(\nabla p\)), the network predicts fluid velocity (\(\mathbf{v}\)).

code for this example is here.