"""
Fourier Neural Operator Module.
"""
import torch
import torch.nn as nn
from pina import LabelTensor
import warnings
from ..utils import check_consistency
from .layers.fourier import FourierBlock1D, FourierBlock2D, FourierBlock3D
from .base_no import KernelNeuralOperator
[docs]
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.
<https://arxiv.org/abs/2010.08895>`_
"""
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,
):
"""
: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__()
# check type consistency
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)
else:
raise ValueError("layers must be tuple or list of int.")
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."
)
# assign padding
self._padding = padding
# initialize fourier layer for each dimension
if dimensions == 1:
fourier_layer = FourierBlock1D
elif dimensions == 2:
fourier_layer = FourierBlock2D
elif dimensions == 3:
fourier_layer = FourierBlock3D
else:
raise NotImplementedError("FNO implemented only for 1D/2D/3D data.")
# Here we build the FNO kernels by stacking Fourier Blocks
# 1. Assign output dimensions for each FNO layer
if layers is None:
layers = [inner_size] * n_layers
# 2. Assign activation functions for each FNO layer
if isinstance(func, list):
if len(layers) != len(func):
raise RuntimeError(
"Uncosistent number of layers and functions."
)
_functions = func
else:
_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
):
raise RuntimeError(
"Uncosistent number of layers and functions."
)
elif all(isinstance(i, int) for i in n_modes):
n_modes = [n_modes] * len(layers)
else:
n_modes = [n_modes] * len(layers)
# 4. Build the FNO network
_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):
padding = [padding] * dimensions
self._ipad = [-pad if pad > 0 else None for pad in padding[:dimensions]]
self._padding_type = padding_type
self._pad = [
val for pair in zip([0] * dimensions, padding) for val in pair
]
[docs]
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 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)
# 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)
# padding the input
x = torch.nn.functional.pad(x, pad=self._pad, mode=self._padding_type)
# apply fourier layers
x = self._layers(x)
# remove padding
idxs = [slice(None), slice(None)] + [slice(pad) for pad in self._ipad]
x = x[idxs]
# permuting back [batch, x, y, ..., channels]
permutation_idx = [0, *[i for i in range(2, x.ndim)], 1]
x = x.permute(permutation_idx)
return x
[docs]
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.
<https://arxiv.org/abs/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,
):
"""
: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]
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,
),
)
[docs]
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 FNO.
:rtype: torch.Tensor
"""
return super().forward(x)
