torchphysics.models.deeponet namespace

Submodules

torchphysics.models.deeponet.branchnets module

class torchphysics.models.deeponet.branchnets.BranchNet(function_space, discretization_sampler)[source]

Bases: Model

A neural network that can be used inside a DeepONet-model.

Parameters:

function_space (Space) – The space of functions that can be put in this network.
discretization_sampler (torchphysics.sampler) – A sampler that will create the points at which the input functions should evaluated, to create a discrete input for the network. The number of input neurons will be equal to the number of sampled points. Therefore, the sampler should always return the same number of points!

finalize(output_space, output_neurons)[source]

Method to set the output space and output neurons of the network. Will be called once the BranchNet is connected to the TrunkNet, so that both will have a fitting output shape.

output_spaceSpace: The space in which the final output of the DeepONet will belong to.
output_neuronsint: The number of output neurons. Will be multiplied my the dimension of the output space, so each dimension will have the same number of intermediate neurons.

fix_input(function, device='cpu')[source]

Fixes the branch net for a given function. The branch net will be evaluated for the given function and the output saved in current_out.

Parameters:

function (callable, torchphysics.domains.FunctionSet, torch.Tensor,) – torchphysics.spaces.Points The function(s) for which the network should be evaluaded.
device (str, optional) – The device where the data lays. Default is ‘cpu’.

Notes

To overwrite the data current_out (the fixed function) just call .fix_input again with a new function.

abstract forward(discrete_function_batch, device='cpu')[source]

Evaluated the network at a given function batch. Should not be called directly, rather use the method .fix_input.

Parameters:

discrete_function_batch (tp.space.Points) – The points object of discrete function values to evaluate the model.
device (str, optional) – The device where the data lays. Default is ‘cpu’.

Notes

Will, in general, not return anything. The output of the network will be saved internally to be used multiple times.

class torchphysics.models.deeponet.branchnets.ConvBranchNet1D(function_space, discretization_sampler, convolutional_network, hidden=(20, 20, 20), activations=Tanh(), xavier_gains=1.6666666666666667)[source]

Bases: BranchNet

A branch network that first applies a convolution to the input functions and afterwards linear FC-layers.

Parameters:

function_space (Space) – The space of functions that can be put in this network.
discretization_sampler (torchphysics.sampler) – A sampler that will create the points at which the input functions should evaluated, to create a discrete input for the network. The number of input neurons will be equal to the number of sampled points. Therefore, the sampler should always return the same number of points!
convolutional_network (torch.nn.module) –
The user defined convolutional network, that should be applied to the branch input. Inside this network, the input can be transformed arbitrary, e.g. you can also apply pooling or other layers. We only expect that the network gets the input in the shape:

[batch_dim, function_space.output_space.dim (channels_in),
len(discretization_sampler)]

You have to make sure, that the number of output dimension is compatible with the following linear layers.
hidden (list or tuple) – The number and size of the hidden layers of the neural network. The lenght of the list/tuple will be equal to the number of hidden layers, while the i-th entry will determine the number of neurons of each layer.
activations (torch.nn or list, optional) – The activation functions of this network. Deafult is nn.Tanh().
xavier_gains (float or list, optional) – For the weight initialization a Xavier/Glorot algorithm will be used. Default is 5/3.

finalize(output_space, output_neurons)[source]

Method to set the output space and output neurons of the network. Will be called once the BranchNet is connected to the TrunkNet, so that both will have a fitting output shape.

output_spaceSpace: The space in which the final output of the DeepONet will belong to.
output_neuronsint: The number of output neurons. Will be multiplied my the dimension of the output space, so each dimension will have the same number of intermediate neurons.

forward(discrete_function_batch)[source]

Evaluated the network at a given function batch. Should not be called directly, rather use the method .fix_input.

Parameters:

discrete_function_batch (tp.space.Points) – The points object of discrete function values to evaluate the model.
device (str, optional) – The device where the data lays. Default is ‘cpu’.

Notes

Will, in general, not return anything. The output of the network will be saved internally to be used multiple times.

class torchphysics.models.deeponet.branchnets.FCBranchNet(function_space, discretization_sampler, hidden=(20, 20, 20), activations=Tanh(), xavier_gains=1.6666666666666667)[source]

Bases: BranchNet

A neural network that can be used inside a DeepONet-model.

Parameters:

function_space (Space) – The space of functions that can be put in this network.
discretization_sampler (torchphysics.sampler) – A sampler that will create the points at which the input functions should evaluated, to create a discrete input for the network. The number of input neurons will be equal to the number of sampled points. Therefore, the sampler should always return the same number of points!
hidden (list or tuple) – The number and size of the hidden layers of the neural network. The lenght of the list/tuple will be equal to the number of hidden layers, while the i-th entry will determine the number of neurons of each layer.
activations (torch.nn or list, optional) – The activation functions of this network. Deafult is nn.Tanh().
xavier_gains (float or list, optional) – For the weight initialization a Xavier/Glorot algorithm will be used. Default is 5/3.

finalize(output_space, output_neurons)[source]

Method to set the output space and output neurons of the network. Will be called once the BranchNet is connected to the TrunkNet, so that both will have a fitting output shape.

output_spaceSpace: The space in which the final output of the DeepONet will belong to.
output_neuronsint: The number of output neurons. Will be multiplied my the dimension of the output space, so each dimension will have the same number of intermediate neurons.

forward(discrete_function_batch)[source]

Evaluated the network at a given function batch. Should not be called directly, rather use the method .fix_input.

Parameters:

discrete_function_batch (tp.space.Points) – The points object of discrete function values to evaluate the model.
device (str, optional) – The device where the data lays. Default is ‘cpu’.

Notes

Will, in general, not return anything. The output of the network will be saved internally to be used multiple times.

torchphysics.models.deeponet.deeponet module

class torchphysics.models.deeponet.deeponet.DeepONet(trunk_net, branch_net, output_space, output_neurons)[source]

Bases: Model

Implementation of the architecture used in the DeepONet paper [1]. Consists of two single neural networks. One for the inputs of the function space (branch net) and one for the inputs of the variables (trunk net).

Parameters:

trunk_net (torchphysics.models.TrunkNet) – The neural network that will get the space/time/… variables as an input.
branch_net (torchphysics.models.BranchNet) – The neural network that will get the function variables as an input.
output_space (Space) – The space in which the final output of the DeepONet will belong to.
output_neurons (int) – The number of output neurons, that will be the output of the TrunkNet and BranchNet. The corresponding outputs of both networks are then connected with the inner product. For higher dimensional outputs, will be multiplied my the dimension of the output space, so each dimension will have the same number of intermediate neurons.

Notes

The number of output neurons in the branch and trunk net have to be the same!

fix_branch_input(function, device='cpu')[source]

Fixes the branch net for a given function. this function will then be used in every following forward call. To set a new function just call this method again.

Parameters:

function (callable, torchphysics.domains.FunctionSet) – The function(s) for which the branch should be evaluaded.
device (str, optional) – The device where the data lays. Default is ‘cpu’.

forward(trunk_inputs, branch_inputs=None, device='cpu')[source]

Apply the network to the given inputs.

Parameters:

trunk_inputs (torchphysics.spaces.Points) – The inputs for the trunk net.
branch_inputs (callable, torchphysics.domains.FunctionSet, optional) – The function(s) for which the branch should be evaluaded. If no input is given, the branch net has to be fixed before hand!
device (str, optional) – The device where the data lays. Default is ‘cpu’.

Returns:

A point object containing the output.

Return type:

torchphysics.spaces.Points

torchphysics.models.deeponet.layers module

class torchphysics.models.deeponet.layers.TrunkLinear(in_features: int, out_features: int, bias: bool = True, device=None, dtype=None)[source]

Bases: Module

Applies a linear transformation to the incoming data: \(y = xA^T + b\), similar

to torch.nn.Linear, but assumes the input x to be identical along the first batch axis, since this is the case in our implementation of trunk nets.

Args:

in_features: size of each input sample out_features: size of each output sample bias: If set to False, the layer will not learn an additive bias.

Default: True

Shape:

Input: \((*, H_{in})\) where \(*\) means any number of dimensions including none and \(H_{in} = ext{in\_features}\).
Output: \((*, H_{out})\) where all but the last dimension are the same shape as the input and \(H_{out} = ext{out\_features}\).

Attributes:

weight: the learnable weights of the module of shape: \(( ext{out\_features}, ext{in\_features})\). The values are initialized from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\), where :math:`k =

rac{1}{ ext{in_features}}`

bias: the learnable bias of the module of shape \(( ext{out\_features})\).: If bias is True, the values are initialized from \(\mathcal{U}(-\sqrt{k}, \sqrt{k})\) where :math:`k =

rac{1}{ ext{in_features}}`

Examples::

>>> m = nn.Linear(20, 30)
>>> input = torch.randn(128, 20)
>>> output = m(input)
>>> print(output.size())
torch.Size([128, 30])

extra_repr() → str[source]

Set the extra representation of the module.

To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.

forward(input: Tensor) → Tensor[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

reset_parameters() → None[source]

class torchphysics.models.deeponet.layers.linear(*args, **kwargs)[source]

Bases: Function

static backward(ctx, grad_output)[source]

Define a formula for differentiating the operation with backward mode automatic differentiation.

This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the vjp function.)

It must accept a context ctx as the first argument, followed by as many outputs as the forward() returned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to forward(). Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.

The context can be used to retrieve tensors saved during the forward pass. It also has an attribute ctx.needs_input_grad as a tuple of booleans representing whether each input needs gradient. E.g., backward() will have ctx.needs_input_grad[0] = True if the first input to forward() needs gradient computed w.r.t. the output.

static forward(ctx, input, weight, bias=None)[source]

Define the forward of the custom autograd Function.

This function is to be overridden by all subclasses. There are two ways to define forward:

Usage 1 (Combined forward and ctx):

@staticmethod
def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any:
    pass

It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
See combining-forward-context for more details

Usage 2 (Separate forward and ctx):

@staticmethod
def forward(*args: Any, **kwargs: Any) -> Any:
    pass

@staticmethod
def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None:
    pass

The forward no longer accepts a ctx argument.
Instead, you must also override the torch.autograd.Function.setup_context() staticmethod to handle setting up the ctx object. output is the output of the forward, inputs are a Tuple of inputs to the forward.
See extending-autograd for more details

The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on ctx (though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with ctx.save_for_backward() if they are intended to be used in backward (equivalently, vjp) or ctx.save_for_forward() if they are intended to be used for in jvp.

torchphysics.models.deeponet.trunknets module

class torchphysics.models.deeponet.trunknets.FCTrunkNet(input_space, hidden=(20, 20, 20), activations=Tanh(), xavier_gains=1.6666666666666667, trunk_input_copied=True)[source]

Bases: TrunkNet

A fully connected neural network that can be used inside a DeepONet. :param input_space: The space of the points the can be put into this model. :type input_space: Space :param hidden: The number and size of the hidden layers of the neural network.

The lenght of the list/tuple will be equal to the number of hidden layers, while the i-th entry will determine the number of neurons of each layer.

Parameters:

activations (torch.nn or list, optional) – The activation functions of this network. Deafult is nn.Tanh().
xavier_gains (float or list, optional) – For the weight initialization a Xavier/Glorot algorithm will be used. Default is 5/3.

finalize(output_space, output_neurons)[source]

Method to set the output space and output neurons of the network. Will be called once the BranchNet is connected to the TrunkNet, so that both will have a fitting output shape.

Parameters:

output_space (Space) – The space in which the final output of the DeepONet will belong to.
output_neurons (int) – The number of output neurons. Will be multiplied my the dimension of the output space, so each dimension will have the same number of intermediate neurons.

forward(points)[source]

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

class torchphysics.models.deeponet.trunknets.TrunkNet(input_space, trunk_input_copied=True)[source]

Bases: Model

A neural network that can be used inside a DeepONet-model.

Parameters:

input_space (Space) – The space of the points that can be put into this model.
trunk_input_copied (bool, optional) – If every sample function of the branch input gets evaluated at the same trunk input, the evaluation process can be speed up, since the trunk only has to evaluated once for the whole data batch of branch inputs. If this is the case, set trunk_input_copied = True. If for example a dataset with different trunk inputs for each branch function is used, set trunk_input_copied = False. Else this may lead to unexpected behavior.

finalize(output_space, output_neurons)[source]

Method to set the output space and output neurons of the network. Will be called once the BranchNet is connected to the TrunkNet, so that both will have a fitting output shape.

Parameters:

output_space (Space) – The space in which the final output of the DeepONet will belong to.
output_neurons (int) – The number of output neurons. Will be multiplied my the dimension of the output space, so each dimension will have the same number of intermediate neurons.

torchphysics.models.deeponet.trunknets.construct_FC_trunk_layers(hidden, input_dim, output_dim, activations, xavier_gains)[source]: Constructs the layer structure for a fully connected neural network.