KernelNeuralOperator#

class KernelNeuralOperator(lifting_operator, integral_kernels, projection_operator)[source]#

Bases: Module

Base class for Neural Operators with integral kernels.

This class serves as a foundation for building Neural Operators that incorporate multiple integral kernels. All Neural Operator models in PINA inherit from this class. The design follows the framework proposed by Kovachki et al., as illustrated in Figure 2 of their work.

Neural Operators derived from this class can be expressed as:

\[G_\theta := P \circ K_m \circ \cdot \circ K_1 \circ L\]

where:

  • \(G_\theta: \mathcal{A}\subset \mathbb{R}^{\rm{in}} \rightarrow \mathcal{D}\subset \mathbb{R}^{\rm{out}}\) is the neural operator approximation of the unknown real operator \(G\), that is \(G \approx G_\theta\)

  • \(L: \mathcal{A}\subset \mathbb{R}^{\rm{in}} \rightarrow \mathbb{R}^{\rm{emb}}\) is a lifting operator mapping the input from its domain \(\mathcal{A}\subset \mathbb{R}^{\rm{in}}\) to its embedding dimension \(\mathbb{R}^{\rm{emb}}\)

  • \(\{K_i : \mathbb{R}^{\rm{emb}} \rightarrow \mathbb{R}^{\rm{emb}} \}_{i=1}^m\) are \(m\) integral kernels mapping each hidden representation to the next one.

  • \(P : \mathbb{R}^{\rm{emb}} \rightarrow \mathcal{D}\subset \mathbb{R}^{\rm{out}}\) is a projection operator mapping the hidden representation to the output function.

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 KernelNeuralOperator class.

Parameters:

  • lifting_operator (torch.nn.Module) – The lifting operator mapping the input to its hidden dimension.

  • integral_kernels (torch.nn.Module) – List of integral kernels mapping each hidden representation to the next one.

  • projection_operator (torch.nn.Module) – The projection operator mapping the hidden representation to the output function.

property lifting_operator#

The lifting operator module.

Returns:

The lifting operator module.

Return type:

torch.nn.Module

property projection_operator#

The projection operator module.

Returns:

The projection operator module.

Return type:

torch.nn.Module

property integral_kernels#

The integral kernels operator module.

Returns:

The integral kernels operator module.

Return type:

torch.nn.Module

forward(x)[source]#

Forward pass for the KernelNeuralOperator model.

The lifting_operator maps the input to the hidden dimension. The integral_kernels apply the integral kernels to the hidden representation. The projection_operator maps the hidden representation to the output function.

Parameters:

x (torch.Tensor | 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, and \(D\) the dimension of the problem. In particular, \(D\) is the number of spatial, parametric, and/or temporal variables plus the field variables. For instance, for 2D problems with 2 output variables, \(D=4\).

Returns:

The output tensor.

Return type:

torch.Tensor