Source code for pina.model.base_no
"""
Kernel Neural Operator Module.
"""
import torch
from pina.utils import check_consistency
[docs]
class KernelNeuralOperator(torch.nn.Module):
r"""
Base class for composing Neural Operators with integral kernels.
This is a base class for composing neural operators with multiple
integral kernels. All neural operator models defined in PINA inherit
from this class. The structure is inspired by the work of Kovachki, N.
et al. see Figure 2 of the reference for extra details. The Neural
Operators inheriting from this class can be written as:
.. math::
G_\theta := P \circ K_m \circ \cdot \circ K_1 \circ L
where:
* :math:`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 :math:`G`, that is
:math:`G \approx G_\theta`
* :math:`L: \mathcal{A}\subset \mathbb{R}^{\rm{in}} \rightarrow
\mathbb{R}^{\rm{emb}}` is a lifting operator mapping the input
from its domain :math:`\mathcal{A}\subset \mathbb{R}^{\rm{in}}`
to its embedding dimension :math:`\mathbb{R}^{\rm{emb}}`
* :math:`\{K_i : \mathbb{R}^{\rm{emb}} \rightarrow
\mathbb{R}^{\rm{emb}} \}_{i=1}^m` are :math:`m` integral kernels
mapping each hidden representation to the next one.
* :math:`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.
.. seealso::
**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.
"""
def __init__(self, lifting_operator, integral_kernels, projection_operator):
"""
:param torch.nn.Module lifting_operator: The lifting operator
mapping the input to its hidden dimension.
:param torch.nn.Module integral_kernels: List of integral kernels
mapping each hidden representation to the next one.
:param torch.nn.Module projection_operator: The projection operator
mapping the hidden representation to the output function.
"""
super().__init__()
self._lifting_operator = lifting_operator
self._integral_kernels = integral_kernels
self._projection_operator = projection_operator
@property
def lifting_operator(self):
"""
The lifting operator property.
"""
return self._lifting_operator
@lifting_operator.setter
def lifting_operator(self, value):
"""
The lifting operator setter
:param torch.nn.Module value: The lifting operator torch module.
"""
check_consistency(value, torch.nn.Module)
self._lifting_operator = value
@property
def projection_operator(self):
"""
The projection operator property.
"""
return self._projection_operator
@projection_operator.setter
def projection_operator(self, value):
"""
The projection operator setter
:param torch.nn.Module value: The projection operator torch module.
"""
check_consistency(value, torch.nn.Module)
self._projection_operator = value
@property
def integral_kernels(self):
"""
The integral kernels operator property.
"""
return self._integral_kernels
@integral_kernels.setter
def integral_kernels(self, value):
"""
The integral kernels operator setter
:param torch.nn.Module value: The integral kernels operator torch
module.
"""
check_consistency(value, torch.nn.Module)
self._integral_kernels = value
[docs]
def forward(self, x):
r"""
Forward computation for Base Neural Operator. It performs a
lifting of the input by the ``lifting_operator``.
Then different layers integral kernels are applied using
``integral_kernels``. Finally the output is projected
to the final dimensionality by the ``projection_operator``.
:param torch.Tensor x: The input tensor for performing the
computation. It expects a tensor :math:`B \times N \times D`,
where :math:`B` is the batch_size, :math:`N` the number of points
in the mesh, :math:`D` the dimension of the problem. In particular
:math:`D` is the number of spatial/paramtric/temporal variables
plus the field variables. For example for 2D problems with 2
output\ variables :math:`D=4`.
:return: The output tensor obtained from the NO.
:rtype: torch.Tensor
"""
x = self.lifting_operator(x)
x = self.integral_kernels(x)
x = self.projection_operator(x)
return x