folivo/PINA
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

🎨 Format Python code with psf/black

Browse Source
This commit is contained in:
ndem0
2024-02-09 11:25:00 +00:00
committed by Nicola Demo
parent 591aeeb02b
commit cbb43a5392
64 changed files with 1323 additions and 955 deletions
Download Patch File Download Diff File Expand all files Collapse all files

73
pina/model/layers/fourier.py
View File

@@ -2,14 +2,18 @@ import torch
import torch.nn as nn
from ...utils import check_consistency
from pina.model.layers import SpectralConvBlock1D, SpectralConvBlock2D, SpectralConvBlock3D
from pina.model.layers import (
SpectralConvBlock1D,
SpectralConvBlock2D,
SpectralConvBlock3D,
)
class FourierBlock1D(nn.Module):
"""
Fourier block implementation for three dimensional
input tensor. The combination of Fourier blocks
make up the Fourier Neural Operator
make up the Fourier Neural Operator
.. seealso::
@@ -21,11 +25,13 @@ class FourierBlock1D(nn.Module):
"""
def __init__(self,
input_numb_fields,
output_numb_fields,
n_modes,
activation=torch.nn.Tanh):
def __init__(
self,
input_numb_fields,
output_numb_fields,
n_modes,
activation=torch.nn.Tanh,
):
super().__init__()
"""
PINA implementation of Fourier block one dimension. The module computes
@@ -51,17 +57,18 @@ class FourierBlock1D(nn.Module):
self._spectral_conv = SpectralConvBlock1D(
input_numb_fields=input_numb_fields,
output_numb_fields=output_numb_fields,
n_modes=n_modes)
n_modes=n_modes,
)
self._activation = activation()
self._linear = nn.Conv1d(input_numb_fields, output_numb_fields, 1)
def forward(self, x):
"""
Forward computation for Fourier Block. It performs a spectral
Forward computation for Fourier Block. It performs a spectral
convolution and a linear transformation of the input and sum the
results.
:param x: The input tensor for fourier block, expect of size
:param x: The input tensor for fourier block, expect of size
``[batch, input_numb_fields, x]``.
:type x: torch.Tensor
:return: The output tensor obtained from the
@@ -75,7 +82,7 @@ class FourierBlock2D(nn.Module):
"""
Fourier block implementation for two dimensional
input tensor. The combination of Fourier blocks
make up the Fourier Neural Operator
make up the Fourier Neural Operator
.. seealso::
@@ -87,18 +94,20 @@ class FourierBlock2D(nn.Module):
"""
def __init__(self,
input_numb_fields,
output_numb_fields,
n_modes,
activation=torch.nn.Tanh):
def __init__(
self,
input_numb_fields,
output_numb_fields,
n_modes,
activation=torch.nn.Tanh,
):
"""
PINA implementation of Fourier block two dimensions. The module computes
the spectral convolution of the input with a linear kernel in the
fourier space, and then it maps the input back to the physical
space. The output is then added to a Linear tranformation of the
input in the physical space. Finally an activation function is
applied to the output.
applied to the output.
The block expects an input of size ``[batch, input_numb_fields, Nx, Ny]``
and returns an output of size ``[batch, output_numb_fields, Nx, Ny]``.
@@ -118,17 +127,18 @@ class FourierBlock2D(nn.Module):
self._spectral_conv = SpectralConvBlock2D(
input_numb_fields=input_numb_fields,
output_numb_fields=output_numb_fields,
n_modes=n_modes)
n_modes=n_modes,
)
self._activation = activation()
self._linear = nn.Conv2d(input_numb_fields, output_numb_fields, 1)
def forward(self, x):
"""
Forward computation for Fourier Block. It performs a spectral
Forward computation for Fourier Block. It performs a spectral
convolution and a linear transformation of the input and sum the
results.
:param x: The input tensor for fourier block, expect of size
:param x: The input tensor for fourier block, expect of size
``[batch, input_numb_fields, x, y]``.
:type x: torch.Tensor
:return: The output tensor obtained from the
@@ -142,7 +152,7 @@ class FourierBlock3D(nn.Module):
"""
Fourier block implementation for three dimensional
input tensor. The combination of Fourier blocks
make up the Fourier Neural Operator
make up the Fourier Neural Operator
.. seealso::
@@ -154,18 +164,20 @@ class FourierBlock3D(nn.Module):
"""
def __init__(self,
input_numb_fields,
output_numb_fields,
n_modes,
activation=torch.nn.Tanh):
def __init__(
self,
input_numb_fields,
output_numb_fields,
n_modes,
activation=torch.nn.Tanh,
):
"""
PINA implementation of Fourier block three dimensions. The module computes
the spectral convolution of the input with a linear kernel in the
fourier space, and then it maps the input back to the physical
space. The output is then added to a Linear tranformation of the
input in the physical space. Finally an activation function is
applied to the output.
applied to the output.
The block expects an input of size ``[batch, input_numb_fields, Nx, Ny, Nz]``
and returns an output of size ``[batch, output_numb_fields, Nx, Ny, Nz]``.
@@ -186,17 +198,18 @@ class FourierBlock3D(nn.Module):
self._spectral_conv = SpectralConvBlock3D(
input_numb_fields=input_numb_fields,
output_numb_fields=output_numb_fields,
n_modes=n_modes)
n_modes=n_modes,
)
self._activation = activation()
self._linear = nn.Conv3d(input_numb_fields, output_numb_fields, 1)
def forward(self, x):
"""
Forward computation for Fourier Block. It performs a spectral
Forward computation for Fourier Block. It performs a spectral
convolution and a linear transformation of the input and sum the
results.
:param x: The input tensor for fourier block, expect of size
:param x: The input tensor for fourier block, expect of size
``[batch, input_numb_fields, x, y, z]``.
:type x: torch.Tensor
:return: The output tensor obtained from the