From 3423c449231acfcbee296ac823f493f252e01e75 Mon Sep 17 00:00:00 2001 From: Dario Coscia Date: Fri, 9 Feb 2024 19:28:24 +0100 Subject: [PATCH 1/4] * Unifying integral kernel NO architectures with BaseNO (maybe rename it better) * Implement FNO based on BaseNO --- pina/model/__init__.py | 17 +-- pina/model/base_no.py | 90 ++++++++++++++++ pina/model/fno.py | 229 ++++++++++++++++++++++++++++------------- 3 files changed, 256 insertions(+), 80 deletions(-) create mode 100644 pina/model/base_no.py diff --git a/pina/model/__init__.py b/pina/model/__init__.py index aab81bf8f..c1a174411 100644 --- a/pina/model/__init__.py +++ b/pina/model/__init__.py @@ -1,13 +1,16 @@ __all__ = [ - "FeedForward", - "ResidualFeedForward", - "MultiFeedForward", - "DeepONet", - "MIONet", - "FNO", + 'FeedForward', + 'ResidualFeedForward', + 'MultiFeedForward', + 'DeepONet', + 'MIONet', + 'FNO', + 'FourierIntegralKernel', + 'BaseNO' ] from .feed_forward import FeedForward, ResidualFeedForward from .multi_feed_forward import MultiFeedForward from .deeponet import DeepONet, MIONet -from .fno import FNO +from .fno import FNO, FourierIntegralKernel +from .base_no import BaseNO diff --git a/pina/model/base_no.py b/pina/model/base_no.py new file mode 100644 index 000000000..460330dab --- /dev/null +++ b/pina/model/base_no.py @@ -0,0 +1,90 @@ +import torch +from pina.utils import check_consistency + +class BaseNO(torch.nn.Module): + def __init__(self, lifting_operator, integral_kernels, projection_operator): + 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: + $$ G_\theta := P \circ K_L \circ\cdot\circ K_1 \circ L,$$ + where $L$ is a lifting operator mapping the input to its hidden + dimension. The $\{K_i\}_{i=1}^L$ are integral kernels mapping + each hidden representation to the next one. Finally $P$ 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. + + :param lifting_operator: The lifting operator mapping the input to its hidden dimension. + :type lifting_operator: torch.nn.Module + :param integral_kernels: List of integral kernels mapping each hidden representation to the next one. + :type integral_kernels: torch.nn.Module + :param projection_operator: The projection operator mapping the hidden representation to the output function. + :type projection_operator: torch.nn.Module + """ + + super().__init__() + + self._lifting_operator = lifting_operator + self._integral_kernels = integral_kernels + self._projection_operator = projection_operator + + @property + def lifting_operator(self): + return self._lifting_operator + + @lifting_operator.setter + def lifting_operator(self, value): + check_consistency(value, torch.nn.Module) + self._lifting_operator = value + + @property + def projection_operator(self): + return self._projection_operator + + @projection_operator.setter + def projection_operator(self, value): + check_consistency(value, torch.nn.Module) + self._projection_operator = value + + @property + def integral_kernels(self): + return self._integral_kernels + + @integral_kernels.setter + def integral_kernels(self, value): + check_consistency(value, torch.nn.Module) + self._integral_kernels = value + + + def forward(self, x): + """ + 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: $$[B \times \times N \times D]$$ + where $B$ is the batch_size, $N$ the number of points in the mesh, + $D$ the dimension of the problem. In particular $D$ is the number + of spatial/paramtric/temporal variables + field variables. + For example for 2D problems with 2 output variables $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 diff --git a/pina/model/fno.py b/pina/model/fno.py index 93168e81a..9d5a32527 100644 --- a/pina/model/fno.py +++ b/pina/model/fno.py @@ -1,54 +1,70 @@ import torch import torch.nn as nn -from ..utils import check_consistency -from .layers.fourier import FourierBlock1D, FourierBlock2D, FourierBlock3D from pina import LabelTensor import warnings +from ..utils import check_consistency +from .layers.fourier import FourierBlock1D, FourierBlock2D, FourierBlock3D +from .base_no import BaseNO + +class FourierIntegralKernel(torch.nn.Module): + def __init__(self, + input_numb_fields, + output_numb_fields, + n_modes, + dimensions=3, + padding=8, + padding_type="constant", + inner_size=20, + n_layers=2, + func=nn.Tanh, + layers=None): + """ + Implementation of Fourier Integral Kernel network. -class FNO(torch.nn.Module): - """ - The PINA implementation of Fourier Neural Operator network. - - Fourier Neural Operator (FNO) is a general architecture for learning Operators. - Unlike traditional machine learning methods FNO is designed to map - entire functions to other functions. It can be trained both with - Supervised learning strategies. FNO does global convolution by performing the - operation on the Fourier space. - - .. seealso:: - - **Original reference**: Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., - Bhattacharya, K., Stuart, A., & Anandkumar, A. (2020). *Fourier neural operator for - parametric partial differential equations*. - DOI: `arXiv preprint arXiv:2010.08895. - `_ - """ - - def __init__( - self, - lifting_net, - projecting_net, - n_modes, - dimensions=3, - padding=8, - padding_type="constant", - inner_size=20, - n_layers=2, - func=nn.Tanh, - layers=None, - ): + This class implements the Fourier Integral Kernel network, which is a + PINA implementation of Fourier Neural Operator kernel network. + It performs global convolution by operating in the Fourier space. + + .. seealso:: + + **Original reference**: Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., + Bhattacharya, K., Stuart, A., & Anandkumar, A. (2020). *Fourier neural operator for + parametric partial differential equations*. + DOI: `arXiv preprint arXiv:2010.08895. + `_ + + :param input_numb_fields: Number of input fields. + :type input_numb_fields: int + :param output_numb_fields: Number of output fields. + :type output_numb_fields: int + :param n_modes: Number of modes. + :type n_modes: int or list[int] + :param dimensions: Number of dimensions (1, 2, or 3). + :type dimensions: int, optional + :param padding: Padding size, defaults to 8. + :type padding: int, optional + :param padding_type: Type of padding, defaults to "constant". + :type padding_type: str, optional + :param inner_size: Inner size, defaults to 20. + :type inner_size: int, optional + :param n_layers: Number of layers, defaults to 2. + :type n_layers: int, optional + :param func: Activation function, defaults to nn.Tanh. + :type func: torch.nn.Module, optional + :param layers: List of layer sizes, defaults to None. + :type layers: list[int], optional + """ super().__init__() # check type consistency - check_consistency(lifting_net, nn.Module) - check_consistency(projecting_net, nn.Module) check_consistency(dimensions, int) check_consistency(padding, int) check_consistency(padding_type, str) check_consistency(inner_size, int) check_consistency(n_layers, int) check_consistency(func, nn.Module, subclass=True) + if layers is not None: if isinstance(layers, (tuple, list)): check_consistency(layers, int) @@ -60,10 +76,7 @@ def __init__( " More information on the official documentation." ) - # assign variables - # TODO check input lifting net and input projecting net - self._lifting_net = lifting_net - self._projecting_net = projecting_net + # assign padding self._padding = padding # initialize fourier layer for each dimension @@ -76,7 +89,7 @@ def __init__( else: raise NotImplementedError("FNO implemented only for 1D/2D/3D data.") - # Here we build the FNO by stacking Fourier Blocks + # Here we build the FNO kernels by stacking Fourier Blocks # 1. Assign output dimensions for each FNO layer if layers is None: @@ -86,11 +99,11 @@ def __init__( if isinstance(func, list): if len(layers) != len(func): raise RuntimeError( - "Uncosistent number of layers and functions." - ) - self._functions = func + 'Uncosistent number of layers and functions.') + _functions = func else: - self._functions = [func for _ in range(len(layers))] + _functions = [func for _ in range(len(layers)-1)] + _functions.append(torch.nn.Identity) # 3. Assign modes functions for each FNO layer if isinstance(n_modes, list): @@ -106,23 +119,15 @@ def __init__( n_modes = [n_modes] * len(layers) # 4. Build the FNO network - tmp_layers = layers.copy() - first_parameter = next(lifting_net.parameters()) - input_shape = first_parameter.size() - out_feats = lifting_net(torch.rand(size=input_shape)).shape[-1] - tmp_layers.insert(0, out_feats) - - self._layers = [] - for i in range(len(tmp_layers) - 1): - self._layers.append( - fourier_layer( - input_numb_fields=tmp_layers[i], - output_numb_fields=tmp_layers[i + 1], - n_modes=n_modes[i], - activation=self._functions[i], - ) - ) - self._layers = nn.Sequential(*self._layers) + _layers = [] + tmp_layers = [input_numb_fields] + layers + [output_numb_fields] + for i in range(len(layers)): + _layers.append( + fourier_layer(input_numb_fields=tmp_layers[i], + output_numb_fields=tmp_layers[i+1], + n_modes=n_modes[i], + activation=_functions[i])) + self._layers = nn.Sequential(*_layers) # 5. Padding values for spectral conv if isinstance(padding, int): @@ -148,15 +153,9 @@ def forward(self, x): :return: The output tensor obtained from the FNO. :rtype: torch.Tensor """ - if isinstance(x, LabelTensor): # TODO remove when Network is fixed - warnings.warn( - "LabelTensor passed as input is not allowed, casting LabelTensor to Torch.Tensor" - ) - x = x.as_subclass(torch.Tensor) - - # lifting the input in higher dimensional space - x = self._lifting_net(x) - + if isinstance(x, LabelTensor): #TODO remove when Network is fixed + warnings.warn('LabelTensor passed as input is not allowed, casting LabelTensor to Torch.Tensor') + x = x.as_subclass(torch.Tensor) # permuting the input [batch, channels, x, y, ...] permutation_idx = [0, x.ndim - 1, *[i for i in range(1, x.ndim - 1)]] x = x.permute(permutation_idx) @@ -175,5 +174,89 @@ def forward(self, x): permutation_idx = [0, *[i for i in range(2, x.ndim)], 1] x = x.permute(permutation_idx) - # apply projecting operator and return - return self._projecting_net(x) + return x + +class FNO(BaseNO): + def __init__(self, + lifting_net, + projecting_net, + n_modes, + dimensions=3, + padding=8, + padding_type="constant", + inner_size=20, + n_layers=2, + func=nn.Tanh, + layers=None): + """ + The PINA implementation of Fourier Neural Operator network. + + Fourier Neural Operator (FNO) is a general architecture for learning Operators. + Unlike traditional machine learning methods FNO is designed to map + entire functions to other functions. It can be trained both with + Supervised learning strategies. FNO does global convolution by performing the + operation on the Fourier space. + + .. seealso:: + + **Original reference**: Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., + Bhattacharya, K., Stuart, A., & Anandkumar, A. (2020). *Fourier neural operator for + parametric partial differential equations*. + DOI: `arXiv preprint arXiv:2010.08895. + `_ + + :param lifting_net: The neural network for lifting the input. + :type lifting_net: torch.nn.Module + :param projecting_net: The neural network for projecting the output. + :type projecting_net: torch.nn.Module + :param n_modes: Number of modes. + :type n_modes: int + :param dimensions: Number of dimensions (1, 2, or 3), defaults to 3. + :type dimensions: int, optional + :param padding: Padding size, defaults to 8. + :type padding: int, optional + :param padding_type: Type of padding, defaults to "constant". + :type padding_type: str, optional + :param inner_size: Inner size, defaults to 20. + :type inner_size: int, optional + :param n_layers: Number of layers, defaults to 2. + :type n_layers: int, optional + :param func: Activation function, defaults to nn.Tanh. + :type func: torch.nn.Module, optional + :param layers: List of layer sizes, defaults to None. + :type layers: list[int], optional + """ + lifting_operator_out = lifting_net( + torch.rand(size=next(lifting_net.parameters()).size())).shape[-1] + super().__init__(lifting_operator=lifting_net, + projection_operator=projecting_net, + integral_kernels=FourierIntegralKernel( + input_numb_fields=lifting_operator_out, + output_numb_fields=next(projecting_net.parameters()).size(), + n_modes=n_modes, + dimensions=dimensions, + padding=padding, + padding_type=padding_type, + inner_size=inner_size, + n_layers=n_layers, + func=func, + layers=layers) + ) + + def forward(self, x): + """ + Forward computation for Fourier Neural Operator. It performs a + lifting of the input by the ``lifting_net``. Then different layers + of Fourier Blocks are applied. Finally the output is projected + to the final dimensionality by the ``projecting_net``. + + :param torch.Tensor x: The input tensor for fourier block, depending on + ``dimension`` in the initialization. In particular it is expected + * 1D tensors: ``[batch, X, channels]`` + * 2D tensors: ``[batch, X, Y, channels]`` + * 3D tensors: ``[batch, X, Y, Z, channels]`` + :return: The output tensor obtained from the FNO. + :rtype: torch.Tensor + """ + return super().forward(x) + From b5f7740b8d7d01ee2c37d3ea2e85bc8e2333ab72 Mon Sep 17 00:00:00 2001 From: Dario Coscia Date: Fri, 16 Feb 2024 13:22:55 +0100 Subject: [PATCH 2/4] * better naming * modify doc for FNO and adding for FourierIntegralKernel, NeuralKernelOperator * adding tests --- docs/source/_rst/_code.rst | 2 ++ docs/source/_rst/models/base_no.rst | 7 ++++ docs/source/_rst/models/fourier_kernel.rst | 7 ++++ pina/model/__init__.py | 4 +-- pina/model/base_no.py | 40 ++++++++++++++------- pina/model/fno.py | 42 +++++++++++----------- tests/test_model/test_base_no.py | 40 +++++++++++++++++++++ 7 files changed, 107 insertions(+), 35 deletions(-) create mode 100644 docs/source/_rst/models/base_no.rst create mode 100644 docs/source/_rst/models/fourier_kernel.rst create mode 100644 tests/test_model/test_base_no.py diff --git a/docs/source/_rst/_code.rst b/docs/source/_rst/_code.rst index bdbe70c85..4156279a8 100644 --- a/docs/source/_rst/_code.rst +++ b/docs/source/_rst/_code.rst @@ -48,11 +48,13 @@ Models :maxdepth: 5 Network + KernelNeuralOperator FeedForward MultiFeedForward ResidualFeedForward DeepONet MIONet + FourierIntegralKernel FNO Layers diff --git a/docs/source/_rst/models/base_no.rst b/docs/source/_rst/models/base_no.rst new file mode 100644 index 000000000..772261c5c --- /dev/null +++ b/docs/source/_rst/models/base_no.rst @@ -0,0 +1,7 @@ +KernelNeuralOperator +======================= +.. currentmodule:: pina.model.base_no + +.. autoclass:: KernelNeuralOperator + :members: + :show-inheritance: \ No newline at end of file diff --git a/docs/source/_rst/models/fourier_kernel.rst b/docs/source/_rst/models/fourier_kernel.rst new file mode 100644 index 000000000..e45ba174d --- /dev/null +++ b/docs/source/_rst/models/fourier_kernel.rst @@ -0,0 +1,7 @@ +FourierIntegralKernel +========================= +.. currentmodule:: pina.model.fno + +.. autoclass:: FourierIntegralKernel + :members: + :show-inheritance: \ No newline at end of file diff --git a/pina/model/__init__.py b/pina/model/__init__.py index c1a174411..d2a6f08b8 100644 --- a/pina/model/__init__.py +++ b/pina/model/__init__.py @@ -6,11 +6,11 @@ 'MIONet', 'FNO', 'FourierIntegralKernel', - 'BaseNO' + 'KernelNeuralOperator' ] from .feed_forward import FeedForward, ResidualFeedForward from .multi_feed_forward import MultiFeedForward from .deeponet import DeepONet, MIONet from .fno import FNO, FourierIntegralKernel -from .base_no import BaseNO +from .base_no import KernelNeuralOperator diff --git a/pina/model/base_no.py b/pina/model/base_no.py index 460330dab..cac15a4de 100644 --- a/pina/model/base_no.py +++ b/pina/model/base_no.py @@ -1,7 +1,7 @@ import torch from pina.utils import check_consistency -class BaseNO(torch.nn.Module): +class KernelNeuralOperator(torch.nn.Module): def __init__(self, lifting_operator, integral_kernels, projection_operator): r""" Base class for composing Neural Operators with integral kernels. @@ -11,11 +11,25 @@ def __init__(self, lifting_operator, integral_kernels, projection_operator): 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: - $$ G_\theta := P \circ K_L \circ\cdot\circ K_1 \circ L,$$ - where $L$ is a lifting operator mapping the input to its hidden - dimension. The $\{K_i\}_{i=1}^L$ are integral kernels mapping - each hidden representation to the next one. Finally $P$ is a projection - operator mapping the hidden representation to the output function. + + .. 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:: @@ -67,7 +81,7 @@ def integral_kernels(self, value): 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 @@ -75,12 +89,12 @@ def forward(self, x): to the final dimensionality by the ``projection_operator``. :param torch.Tensor x: The input tensor for performing the - computation. It expects a tensor: $$[B \times \times N \times D]$$ - where $B$ is the batch_size, $N$ the number of points in the mesh, - $D$ the dimension of the problem. In particular $D$ is the number - of spatial/paramtric/temporal variables + field variables. - For example for 2D problems with 2 output variables $D=4$. - + 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 """ diff --git a/pina/model/fno.py b/pina/model/fno.py index 9d5a32527..7797bd425 100644 --- a/pina/model/fno.py +++ b/pina/model/fno.py @@ -4,7 +4,7 @@ import warnings from ..utils import check_consistency from .layers.fourier import FourierBlock1D, FourierBlock2D, FourierBlock3D -from .base_no import BaseNO +from .base_no import KernelNeuralOperator class FourierIntegralKernel(torch.nn.Module): @@ -41,19 +41,19 @@ def __init__(self, :param n_modes: Number of modes. :type n_modes: int or list[int] :param dimensions: Number of dimensions (1, 2, or 3). - :type dimensions: int, optional + :type dimensions: int :param padding: Padding size, defaults to 8. - :type padding: int, optional + :type padding: int :param padding_type: Type of padding, defaults to "constant". - :type padding_type: str, optional + :type padding_type: str :param inner_size: Inner size, defaults to 20. - :type inner_size: int, optional + :type inner_size: int :param n_layers: Number of layers, defaults to 2. - :type n_layers: int, optional + :type n_layers: int :param func: Activation function, defaults to nn.Tanh. - :type func: torch.nn.Module, optional + :type func: torch.nn.Module :param layers: List of layer sizes, defaults to None. - :type layers: list[int], optional + :type layers: list[int] """ super().__init__() @@ -146,11 +146,12 @@ def forward(self, x): to the final dimensionality by the ``projecting_net``. :param torch.Tensor x: The input tensor for fourier block, depending on - ``dimension`` in the initialization. In particular it is expected + ``dimension`` in the initialization. In particular it is expected: + * 1D tensors: ``[batch, X, channels]`` * 2D tensors: ``[batch, X, Y, channels]`` * 3D tensors: ``[batch, X, Y, Z, channels]`` - :return: The output tensor obtained from the FNO. + :return: The output tensor obtained from the kernels convolution. :rtype: torch.Tensor """ if isinstance(x, LabelTensor): #TODO remove when Network is fixed @@ -176,7 +177,7 @@ def forward(self, x): return x -class FNO(BaseNO): +class FNO(KernelNeuralOperator): def __init__(self, lifting_net, projecting_net, @@ -212,19 +213,19 @@ def __init__(self, :param n_modes: Number of modes. :type n_modes: int :param dimensions: Number of dimensions (1, 2, or 3), defaults to 3. - :type dimensions: int, optional + :type dimensions: int :param padding: Padding size, defaults to 8. - :type padding: int, optional + :type padding: int :param padding_type: Type of padding, defaults to "constant". - :type padding_type: str, optional + :type padding_type: str :param inner_size: Inner size, defaults to 20. - :type inner_size: int, optional + :type inner_size: int :param n_layers: Number of layers, defaults to 2. - :type n_layers: int, optional + :type n_layers: int :param func: Activation function, defaults to nn.Tanh. - :type func: torch.nn.Module, optional + :type func: torch.nn.Module :param layers: List of layer sizes, defaults to None. - :type layers: list[int], optional + :type layers: list[int] """ lifting_operator_out = lifting_net( torch.rand(size=next(lifting_net.parameters()).size())).shape[-1] @@ -251,11 +252,12 @@ def forward(self, x): to the final dimensionality by the ``projecting_net``. :param torch.Tensor x: The input tensor for fourier block, depending on - ``dimension`` in the initialization. In particular it is expected + ``dimension`` in the initialization. In particular it is expected: + * 1D tensors: ``[batch, X, channels]`` * 2D tensors: ``[batch, X, Y, channels]`` * 3D tensors: ``[batch, X, Y, Z, channels]`` - :return: The output tensor obtained from the FNO. + :return: The output tensor obtained from FNO. :rtype: torch.Tensor """ return super().forward(x) diff --git a/tests/test_model/test_base_no.py b/tests/test_model/test_base_no.py new file mode 100644 index 000000000..4a14fd1e4 --- /dev/null +++ b/tests/test_model/test_base_no.py @@ -0,0 +1,40 @@ +import torch +from pina.model import KernelNeuralOperator, FeedForward + +input_dim = 2 +output_dim = 4 +embedding_dim = 24 +batch_size = 10 +numb = 256 +data = torch.rand(size=(batch_size, numb, input_dim), requires_grad=True) +output_shape = torch.Size([batch_size, numb, output_dim]) + + +lifting_operator = FeedForward(input_dimensions=input_dim, output_dimensions=embedding_dim) +projection_operator = FeedForward(input_dimensions=embedding_dim, output_dimensions=output_dim) +integral_kernels = torch.nn.Sequential(FeedForward(input_dimensions=embedding_dim, + output_dimensions=embedding_dim), + FeedForward(input_dimensions=embedding_dim, + output_dimensions=embedding_dim),) + +def test_constructor(): + KernelNeuralOperator(lifting_operator=lifting_operator, + integral_kernels=integral_kernels, + projection_operator=projection_operator) + +def test_forward(): + operator = KernelNeuralOperator(lifting_operator=lifting_operator, + integral_kernels=integral_kernels, + projection_operator=projection_operator) + out = operator(data) + assert out.shape == output_shape + +def test_backward(): + operator = KernelNeuralOperator(lifting_operator=lifting_operator, + integral_kernels=integral_kernels, + projection_operator=projection_operator) + out = operator(data) + loss = torch.nn.functional.mse_loss(out, torch.zeros_like(out)) + loss.backward() + grad = data.grad + assert grad.shape == data.shape From ecb3c94d0d019ea02e09f4a4794e24380959f24d Mon Sep 17 00:00:00 2001 From: Dario Coscia Date: Fri, 16 Feb 2024 13:51:06 +0100 Subject: [PATCH 3/4] codacy --- pina/model/base_no.py | 7 +++--- pina/model/fno.py | 51 ++++++++++++++++++++++--------------------- 2 files changed, 30 insertions(+), 28 deletions(-) diff --git a/pina/model/base_no.py b/pina/model/base_no.py index cac15a4de..3859ae7b4 100644 --- a/pina/model/base_no.py +++ b/pina/model/base_no.py @@ -1,7 +1,9 @@ import torch from pina.utils import check_consistency + class KernelNeuralOperator(torch.nn.Module): + def __init__(self, lifting_operator, integral_kernels, projection_operator): r""" Base class for composing Neural Operators with integral kernels. @@ -54,7 +56,7 @@ def __init__(self, lifting_operator, integral_kernels, projection_operator): @property def lifting_operator(self): - return self._lifting_operator + return self._lifting_operator @lifting_operator.setter def lifting_operator(self, value): @@ -63,7 +65,7 @@ def lifting_operator(self, value): @property def projection_operator(self): - return self._projection_operator + return self._projection_operator @projection_operator.setter def projection_operator(self, value): @@ -79,7 +81,6 @@ def integral_kernels(self, value): check_consistency(value, torch.nn.Module) self._integral_kernels = value - def forward(self, x): r""" Forward computation for Base Neural Operator. It performs a diff --git a/pina/model/fno.py b/pina/model/fno.py index 7797bd425..ae20a2498 100644 --- a/pina/model/fno.py +++ b/pina/model/fno.py @@ -8,6 +8,7 @@ class FourierIntegralKernel(torch.nn.Module): + def __init__(self, input_numb_fields, output_numb_fields, @@ -73,8 +74,7 @@ def __init__(self, if not isinstance(n_modes, (list, tuple, int)): raise ValueError( "n_modes must be a int or list or tuple of valid modes." - " More information on the official documentation." - ) + " More information on the official documentation.") # assign padding self._padding = padding @@ -102,17 +102,15 @@ def __init__(self, 'Uncosistent number of layers and functions.') _functions = func else: - _functions = [func for _ in range(len(layers)-1)] + _functions = [func for _ in range(len(layers) - 1)] _functions.append(torch.nn.Identity) # 3. Assign modes functions for each FNO layer if isinstance(n_modes, list): - if all(isinstance(i, list) for i in n_modes) and len(layers) != len( - n_modes - ): + if all(isinstance(i, list) + for i in n_modes) and len(layers) != len(n_modes): raise RuntimeError( - "Uncosistent number of layers and functions." - ) + "Uncosistent number of layers and functions.") elif all(isinstance(i, int) for i in n_modes): n_modes = [n_modes] * len(layers) else: @@ -124,7 +122,7 @@ def __init__(self, for i in range(len(layers)): _layers.append( fourier_layer(input_numb_fields=tmp_layers[i], - output_numb_fields=tmp_layers[i+1], + output_numb_fields=tmp_layers[i + 1], n_modes=n_modes[i], activation=_functions[i])) self._layers = nn.Sequential(*_layers) @@ -154,9 +152,11 @@ def forward(self, x): :return: The output tensor obtained from the kernels convolution. :rtype: torch.Tensor """ - if isinstance(x, LabelTensor): #TODO remove when Network is fixed - warnings.warn('LabelTensor passed as input is not allowed, casting LabelTensor to Torch.Tensor') - x = x.as_subclass(torch.Tensor) + if isinstance(x, LabelTensor): #TODO remove when Network is fixed + warnings.warn( + 'LabelTensor passed as input is not allowed, casting LabelTensor to Torch.Tensor' + ) + x = x.as_subclass(torch.Tensor) # permuting the input [batch, channels, x, y, ...] permutation_idx = [0, x.ndim - 1, *[i for i in range(1, x.ndim - 1)]] x = x.permute(permutation_idx) @@ -177,7 +177,9 @@ def forward(self, x): return x + class FNO(KernelNeuralOperator): + def __init__(self, lifting_net, projecting_net, @@ -232,18 +234,18 @@ def __init__(self, super().__init__(lifting_operator=lifting_net, projection_operator=projecting_net, integral_kernels=FourierIntegralKernel( - input_numb_fields=lifting_operator_out, - output_numb_fields=next(projecting_net.parameters()).size(), - n_modes=n_modes, - dimensions=dimensions, - padding=padding, - padding_type=padding_type, - inner_size=inner_size, - n_layers=n_layers, - func=func, - layers=layers) - ) - + input_numb_fields=lifting_operator_out, + output_numb_fields=next( + projecting_net.parameters()).size(), + n_modes=n_modes, + dimensions=dimensions, + padding=padding, + padding_type=padding_type, + inner_size=inner_size, + n_layers=n_layers, + func=func, + layers=layers)) + def forward(self, x): """ Forward computation for Fourier Neural Operator. It performs a @@ -261,4 +263,3 @@ def forward(self, x): :rtype: torch.Tensor """ return super().forward(x) - From 0e62e668f7f9aedf43039aff505d69c1e36cfe75 Mon Sep 17 00:00:00 2001 From: Dario Coscia Date: Fri, 16 Feb 2024 16:14:53 +0100 Subject: [PATCH 4/4] codacy 2 --- pina/model/base_no.py | 121 +++++++++++++++++++++------------- pina/model/fno.py | 149 ++++++++++++++++++++---------------------- 2 files changed, 147 insertions(+), 123 deletions(-) diff --git a/pina/model/base_no.py b/pina/model/base_no.py index 3859ae7b4..57434377b 100644 --- a/pina/model/base_no.py +++ b/pina/model/base_no.py @@ -1,51 +1,57 @@ +""" +Kernel Neural Operator Module. +""" + import torch from pina.utils import check_consistency 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): - 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. - - :param lifting_operator: The lifting operator mapping the input to its hidden dimension. - :type lifting_operator: torch.nn.Module - :param integral_kernels: List of integral kernels mapping each hidden representation to the next one. - :type integral_kernels: torch.nn.Module - :param projection_operator: The projection operator mapping the hidden representation to the output function. - :type projection_operator: torch.nn.Module + """ + :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__() @@ -56,28 +62,53 @@ def __init__(self, lifting_operator, integral_kernels, 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 @@ -93,9 +124,9 @@ def forward(self, x): 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`. + :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 """ diff --git a/pina/model/fno.py b/pina/model/fno.py index ae20a2498..e320383f1 100644 --- a/pina/model/fno.py +++ b/pina/model/fno.py @@ -1,3 +1,7 @@ +""" +Fourier Neural Operator Module. +""" + import torch import torch.nn as nn from pina import LabelTensor @@ -8,7 +12,22 @@ class FourierIntegralKernel(torch.nn.Module): - + """ + Implementation of Fourier Integral Kernel network. + + This class implements the Fourier Integral Kernel network, which is a + PINA implementation of Fourier Neural Operator kernel network. + It performs global convolution by operating in the Fourier space. + + .. seealso:: + + **Original reference**: Li, Z., Kovachki, N., Azizzadenesheli, + K., Liu, B., Bhattacharya, K., Stuart, A., & Anandkumar, A. + (2020). *Fourier neural operator for parametric partial + differential equations*. + DOI: `arXiv preprint arXiv:2010.08895. + `_ + """ def __init__(self, input_numb_fields, output_numb_fields, @@ -21,40 +40,16 @@ def __init__(self, func=nn.Tanh, layers=None): """ - Implementation of Fourier Integral Kernel network. - - This class implements the Fourier Integral Kernel network, which is a - PINA implementation of Fourier Neural Operator kernel network. - It performs global convolution by operating in the Fourier space. - - .. seealso:: - - **Original reference**: Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., - Bhattacharya, K., Stuart, A., & Anandkumar, A. (2020). *Fourier neural operator for - parametric partial differential equations*. - DOI: `arXiv preprint arXiv:2010.08895. - `_ - - :param input_numb_fields: Number of input fields. - :type input_numb_fields: int - :param output_numb_fields: Number of output fields. - :type output_numb_fields: int - :param n_modes: Number of modes. - :type n_modes: int or list[int] - :param dimensions: Number of dimensions (1, 2, or 3). - :type dimensions: int - :param padding: Padding size, defaults to 8. - :type padding: int - :param padding_type: Type of padding, defaults to "constant". - :type padding_type: str - :param inner_size: Inner size, defaults to 20. - :type inner_size: int - :param n_layers: Number of layers, defaults to 2. - :type n_layers: int - :param func: Activation function, defaults to nn.Tanh. - :type func: torch.nn.Module - :param layers: List of layer sizes, defaults to None. - :type layers: list[int] + :param int input_numb_fields: Number of input fields. + :param int output_numb_fields: Number of output fields. + :param int | list[int] n_modes: Number of modes. + :param int dimensions: Number of dimensions (1, 2, or 3). + :param int padding: Padding size, defaults to 8. + :param str padding_type: Type of padding, defaults to "constant". + :param int inner_size: Inner size, defaults to 20. + :param int n_layers: Number of layers, defaults to 2. + :param torch.nn.Module func: Activation function, defaults to nn.Tanh. + :param list[int] layers: List of layer sizes, defaults to None. """ super().__init__() @@ -87,7 +82,9 @@ def __init__(self, elif dimensions == 3: fourier_layer = FourierBlock3D else: - raise NotImplementedError("FNO implemented only for 1D/2D/3D data.") + raise NotImplementedError( + "FNO implemented only for 1D/2D/3D data." + ) # Here we build the FNO kernels by stacking Fourier Blocks @@ -143,8 +140,9 @@ def forward(self, x): of Fourier Blocks are applied. Finally the output is projected to the final dimensionality by the ``projecting_net``. - :param torch.Tensor x: The input tensor for fourier block, depending on - ``dimension`` in the initialization. In particular it is expected: + :param torch.Tensor x: The input tensor for fourier block, + depending on ``dimension`` in the initialization. + In particular it is expected: * 1D tensors: ``[batch, X, channels]`` * 2D tensors: ``[batch, X, Y, channels]`` @@ -154,7 +152,8 @@ def forward(self, x): """ if isinstance(x, LabelTensor): #TODO remove when Network is fixed warnings.warn( - 'LabelTensor passed as input is not allowed, casting LabelTensor to Torch.Tensor' + 'LabelTensor passed as input is not allowed,' + ' casting LabelTensor to Torch.Tensor' ) x = x.as_subclass(torch.Tensor) # permuting the input [batch, channels, x, y, ...] @@ -179,7 +178,24 @@ def forward(self, x): class FNO(KernelNeuralOperator): - + """ + The PINA implementation of Fourier Neural Operator network. + + Fourier Neural Operator (FNO) is a general architecture for + learning Operators. Unlike traditional machine learning methods + FNO is designed to map entire functions to other functions. It + can be trained with Supervised learning strategies. FNO does global + convolution by performing the operation on the Fourier space. + + .. seealso:: + + **Original reference**: Li, Z., Kovachki, N., Azizzadenesheli, + K., Liu, B., Bhattacharya, K., Stuart, A., & Anandkumar, A. + (2020). *Fourier neural operator for parametric partial + differential equations*. + DOI: `arXiv preprint arXiv:2010.08895. + `_ + """ def __init__(self, lifting_net, projecting_net, @@ -192,42 +208,18 @@ def __init__(self, func=nn.Tanh, layers=None): """ - The PINA implementation of Fourier Neural Operator network. - - Fourier Neural Operator (FNO) is a general architecture for learning Operators. - Unlike traditional machine learning methods FNO is designed to map - entire functions to other functions. It can be trained both with - Supervised learning strategies. FNO does global convolution by performing the - operation on the Fourier space. - - .. seealso:: - - **Original reference**: Li, Z., Kovachki, N., Azizzadenesheli, K., Liu, B., - Bhattacharya, K., Stuart, A., & Anandkumar, A. (2020). *Fourier neural operator for - parametric partial differential equations*. - DOI: `arXiv preprint arXiv:2010.08895. - `_ - - :param lifting_net: The neural network for lifting the input. - :type lifting_net: torch.nn.Module - :param projecting_net: The neural network for projecting the output. - :type projecting_net: torch.nn.Module - :param n_modes: Number of modes. - :type n_modes: int - :param dimensions: Number of dimensions (1, 2, or 3), defaults to 3. - :type dimensions: int - :param padding: Padding size, defaults to 8. - :type padding: int - :param padding_type: Type of padding, defaults to "constant". - :type padding_type: str - :param inner_size: Inner size, defaults to 20. - :type inner_size: int - :param n_layers: Number of layers, defaults to 2. - :type n_layers: int - :param func: Activation function, defaults to nn.Tanh. - :type func: torch.nn.Module - :param layers: List of layer sizes, defaults to None. - :type layers: list[int] + :param torch.nn.Module lifting_net: The neural network for lifting + the input. + :param torch.nn.Module projecting_net: The neural network for + projecting the output. + :param int | list[int] n_modes: Number of modes. + :param int dimensions: Number of dimensions (1, 2, or 3). + :param int padding: Padding size, defaults to 8. + :param str padding_type: Type of padding, defaults to `constant`. + :param int inner_size: Inner size, defaults to 20. + :param int n_layers: Number of layers, defaults to 2. + :param torch.nn.Module func: Activation function, defaults to nn.Tanh. + :param list[int] layers: List of layer sizes, defaults to None. """ lifting_operator_out = lifting_net( torch.rand(size=next(lifting_net.parameters()).size())).shape[-1] @@ -253,8 +245,9 @@ def forward(self, x): of Fourier Blocks are applied. Finally the output is projected to the final dimensionality by the ``projecting_net``. - :param torch.Tensor x: The input tensor for fourier block, depending on - ``dimension`` in the initialization. In particular it is expected: + :param torch.Tensor x: The input tensor for fourier block, + depending on ``dimension`` in the initialization. In + particular it is expected: * 1D tensors: ``[batch, X, channels]`` * 2D tensors: ``[batch, X, Y, channels]``