Low Rank Neural Operator#

class LowRankNeuralOperator(lifting_net, projecting_net, field_indices, coordinates_indices, n_kernel_layers, rank, inner_size=20, n_layers=2, func=<class 'torch.nn.modules.activation.Tanh'>, bias=True)[source]#

Bases: KernelNeuralOperator

Low Rank Neural Operator model class.

The Low Rank Neural Operator is a general architecture for learning operators, which map functions to functions. It can be trained both with Supervised and Physics-Informed learning strategies. The Low Rank Neural Operator performs convolution by means of a low rank approximation.

See also

Original reference: Kovachki, N., Li, Z., Liu, B., Azizzadenesheli, K., Bhattacharya, K., Stuart, A., & Anandkumar, A. (2023). Neural operator: Learning maps between function spaces with applications to PDEs. Journal of Machine Learning Research, 24(89), 1-97.

Initialization of the LowRankNeuralOperator class.

Parameters:

  • lifting_net (torch.nn.Module) – The lifting neural network mapping the input to its hidden dimension. It must take as input the input field and the coordinates at which the input field is evaluated.

  • projecting_net (torch.nn.Module) – The projection neural network mapping the hidden representation to the output function. It must take as input the embedding dimension plus the dimension of the coordinates.

  • field_indices (list[str]) – The labels of the fields in the input tensor.

  • coordinates_indices (list[str]) – The labels of the coordinates in the input tensor.

  • n_kernel_layers (int) – The number of hidden kernel layers.

  • rank (int) – The rank of the low rank approximation.

  • inner_size (int) – The number of neurons for each hidden layer in the basis function neural network. Default is 20.

  • n_layers (int) – The number of hidden layers in the basis function neural network. Default is 2.

  • func (torch.nn.Module | list[torch.nn.Module]) – The activation function. If a list is passed, it must have the same length as n_layers. If a single function is passed, it is used for all layers, except for the last one. Default is torch.nn.Tanh.

  • bias (bool) – If True bias is considered for the basis function neural network. Default is True.

Raises:

  • ValueError – If the input dimension does not match with the labels of the fields and coordinates.

  • ValueError – If the input dimension of the projecting network does not match with the hidden dimension of the lifting network.

forward(x)[source]#

Forward pass for the LowRankNeuralOperator model.

The lifting_net maps the input to the hidden dimension. Then, several layers of LowRankBlock are applied. Finally, the projecting_net maps the hidden representation to the output function.

Parameters:

x (LabelTensor) – The input tensor for performing the computation. It expects a tensor \(B \times N \times D\), where \(B\) is the batch_size, \(N\) the number of points in the mesh, \(D\) the dimension of the problem, i.e. the sum of len(coordinates_indices) and len(field_indices).

Returns:

The output tensor.

Return type:

torch.Tensor