Deep Operator Network (DeepONet)

Implementation of the Deep Operator Network (DeepONet) as defined by Lu et al [42]. The deep operator network is defined by its branch and trunk networks, as shown below.

A diagram showing the architecture of a deep operator network.

The architecture of a generic deep operator network.

The DeepOperatorNetwork class is imported using the following command:

>>> from UQpy.scientific_machine_learning import DeepOperatorNetwork

Methods

class DeepOperatorNetwork(branch_network, trunk_network, out_channels=1)[source]

Construct a Deep Operator Network via its branch and trunk networks

Parameters:

  • branch_network (Module) – Encodes mapping of the function \(f(x)\)

  • trunk_network (Module) – Encodes mapping of the domain \(y\) for \(\mathcal{L}f(y)\)

  • out_channels (int) – Number of channels produced by the Deep Operator Network

Note

The last layer of the branch and trunk network must have the same number of neurons so the last dimension of their outputs match, i.e. both outputs have shape \((*, \text{width})\). Additionally, \(\text{width}\) must be divisible by \(C_\text{out}\).

Shape:

  • Input: Two tensors representing \(x\) and \(f(x)\)

    • Branch Network (\(f(x)\)): Any shape (can be different from trunk)

    • Trunk Network (\(x\)): Any shape (can be different from branch)

  • Intermediary: The output from the branch and trunk network must be of shapes \((*, \text{width})\) and \((*, \text{width})\). Where \(*\) refers to any broadcastable dimensions. Both the tensors are viewed reshaped as \((*, C_\text{out}. \frac{\text{width}}{C_\text{out}}\) before the dot product is computed.

  • Output: Tensor of shape \((*, C_\text{out})\)

forward(x, f_x)[source]

Compute the dot product of branch and trunk outputs

Parameters:

  • x (Tensor) – Input to the trunk_network

  • f_x (Tensor) – Input to the branch_network

Return type:

Tensor

Returns:

Dot product of the branch and trunk outputs

Raises:

RuntimeError – If incompatible trunk and branch outputs are encountered. See Shape for details.

summary(**kwargs)

Call torchinfo.summary() on self. See torchinfo documentation for details.

Parameters:

kwargs – Keyword arguments passed to torchinfo.summary.

Returns:

Model statistics

count_parameters()

Get the total number of parameters that require a gradient computation in the model

drop(mode=True)

Set dropping mode.

Note

This method and self.dropping only affects UQpy’s dropout layers

Parameters:

mode (bool) – If True perform dropout, otherwise act as the identity function.

sample(mode=True)

Set sampling mode.

Note

This method and self.sampling only affects UQpy’s Bayesian layers

Parameters:

mode (bool) – If True sample from distributions, otherwise use distribution means.

Returns:

self

is_deterministic()

Check if neural network is behaving deterministically or probabilistically.

Note

This flag may be incorrect if the model has sources of randomness that do not depend on the attributes training, dropping, or sampling.

Return type:

bool

Returns:

True if sampling, dropping, and training are all False. Otherwise, returns False.

set_deterministic(mode=True)

Set training, dropping, and sampling to the opposite of mode.

This is equivalent to

>>> model.train(not mode)
>>> model.drop(not mode)
>>> model.sample(not mode)

If the model has sources of randomness that do not depend on the training, dropping, or sampling attributes, they will not be affected.

Attributes

DeepOperatorNetwork.branch_network: nn.Module

Architecture of the branch neural network defined by a torch.nn.Module

DeepOperatorNetwork.trunk_network: nn.Module

Architecture of the trunk neural network defined by a torch.nn.Module

Examples