DeepONet

Deep Operator Networks (DeepONet)

neojax implements a general modular DeepONet (arbitrary branch and trunk networks) and a preconfigured MLPDeepONet (MLP branch and MLP trunk network).

The DeepONets are implemented to evaluate the underlying operator with input function \(u \in \mathbb{R}^m\) at the output \(y \in \mathbb{R}^d\). In case of higher dimensional input spaces, e.g. \(\mathbb{R}^2, \mathbb{R}^3, \dots\) the sensor locations m may not be 1D arrays, as long as the branch network handles this correctly.

In practice, it is more useful to evaluate the output across a grid of points y. See below for details.

How to evaluate the input function across a grid of points

To evaluate a single input function \(u\) across a grid (or batch) of points \(y\), use jax.vmap. Here is how you would evaluate on a \(256 \times 256\) grid:

import jax
import jax.numpy as jnp
from neojax.models import DeepONet

# Initialize model
model = DeepONet(...)

# Define inputs
u = jnp.ones((100,))  # One function input
y_grid = jnp.meshgrid(jnp.linspace(0, 1, 256), jnp.linspace(0, 1, 256))
y_points = jnp.stack(y_grid, axis=-1)  # Shape: (256, 256, 2)

# Vectorize the model over the coordinate axes
# in_axes: (None, 0) means 'u' is fixed, 'y' is mapped over its 0-th dimension
vmapped_inner = jax.vmap(model, in_axes=(None, 0))          # Maps (256, 2) -> (256, 1)
vmapped_outer = jax.vmap(vmapped_inner, in_axes=(None, 0))  # Maps (256, 256, 2) -> (256, 256, 1)

# Generate predictions
predictions = vmapped_outer(u, y_points)  # Shape: (256, 256, 1)

Pro Tip

If you want to evaluate a batch of functions across a batch of points, just add a third jax.vmap call.

# Batch of 32 functions, each evaluated at 256 points
u_batch = jnp.ones((32, 100))
y_batch = jnp.zeros((32, 256, 2))

# vmap over functions (axis 0) and points (axis 0)
batch_model = jax.vmap(vmapped_inner, in_axes=(0, 0))
batch_preds = batch_model(u_batch, y_batch) # Shape: (32, 256, 1)

---

General DeepONet

neojax.models.deeponet.DeepONet

General DeepONet (Deep Operator Network).

Learns a PDE operator through two separate branch and trunk networks. The branch net encodes the discrete function space and the trunk net encodes the domain of output functions. Implemented as in the original publication [1].

Parameters:

Name	Type	Description	Default
branch_net	Module	The network representing the branch (encodes u). Should map shape (m_sensors,) to (p,). May also take different input shapes, as long as it outputs (p,), e.g. CNN branch net.	required
trunk_net	Module	The network representing the trunk (encodes y). Should map shape (d_dim,) to (p,).	required
out_activation	Callable | None	Optional output activation after dot product.	None
bias	Float[Array, 1] | None	Optional learnable bias after dot product.	None

Internal Attributes

These fields store the internal layers state (and weights).

• branch_net (eqx.Module): Arbitrary NN that maps (m_sensors,) -> (p,).
• trunk_net (eqx.Module): Arbitrary NN that maps (d_dim,) -> (p,).
• out_activation (Callable | None): Optional output activation after dot product.
• bias (Float[Array, "1"] | None): Optional learnable bias after dot product.

Info

This is the Unstacked (Standard) implementation as described in Lu et al. (2021). It uses a single branch network and a single trunk network for maximum computational efficiency.

For vector-valued operators (the "Stacked" variant), it is recommended to either:

1. Increase the output dimension of the branch/trunk networks.
2. Use jax.vmap to wrap this class for multiple independent operators.

Cite

Learning nonlinear operators via DeepONet based on the universal approximation theorem of operators

@article{lu2021learning,
    title={Learning nonlinear operators via DeepONet
    based on the universal approximation theorem of operators},
    author={Lu, Lu and Jin, Pengzhan and Pang,
    Guofei and Zhang, Zhongqiang and Karniadakis, George Em},
    journal={Nature machine intelligence},
    volume={3},
    number={3},
    pages={218--229},
    year={2021},
    publisher={Nature Publishing Group}
}

__call__(u: Float[Array, m_sensors], y: Float[Array, d_dim]) -> Float[Array, '']

Forward pass of DeepONet.

Parameters:

Name	Type	Description	Default
u	Float[Array, m_sensors]	Input function discretized at sensor points.	required
y	Float[Array, d_dim]	Coordinates in the output domain.	required

Returns:

Type	Description
Float[Array, '']	Scalar value of operator at y.

---

MLPDeepONet

neojax.models.deeponet.MLPDeepONet

DeepONet with MLP branch and MLP trunk network.

The branch MLP maps (m_sensors,) function values to the (p,) latent vector and the trunk MLP maps the (d,) evaluation vector to the (p,) latent vector.

Parameters:

Name	Type	Description	Default
key	PRNGKeyArray	Jax random key.	required
m_sensors	int	Number of function values u.	required
d_dim	int	Output space dimension.	required
p_latent	int	Latent vector dimension.	required
branch_hidden_dims	Sequence[int]	Hidden dimensions of the branch MLP. First layer is (m_sensors, branch_hidden_dims[0]) and last layer is (branch_hidden_dims[-1], p_latent).	required
trunk_hidden_dims	Sequence[int]	Hidden dimensions of the trunk MLP. First layers is (d_dim, trunk_hidden_dims[0]) and last layer is (trunk_hidden_dims[-1], p_latent).	required
branch_activations	Callable | Sequence[Callable]	Activation functions of branch MLP.	gelu
trunk_activations	Callable | Sequence[Callable]	Activation functions of trunk MLP. Single callable means, all activations use this function.	gelu
out_activation	Callable | None	Optional output activation after dot product. Default is None.	None
bias	Float[Array, ''] | None	Optional learnable bias after dot product. Default is no bias.	None

Internal Attributes

These fields store the internal layers state (and weights).

• branch_net (PointwiseMLP): MLP that maps (m_sensors,) -> (p,).
• trunk_net (PointwiseMLP): MLP that maps (d_dim,) -> (p,).
• out_activation (Callable | None): Optional output activation after dot product.
• bias (Float[Array, ""] | None): Optional learnable bias after dot product.