🎨 Format Python code with psf/black

This commit is contained in:
ndem0
2024-04-02 08:24:40 +00:00
committed by Nicola Demo
parent 766856494a
commit 14dfa62402
5 changed files with 40 additions and 29 deletions

View File

@@ -8,7 +8,7 @@ __all__ = [
"FourierIntegralKernel",
"KernelNeuralOperator",
"AveragingNeuralOperator",
"LowRankNeuralOperator"
"LowRankNeuralOperator",
]
from .feed_forward import FeedForward, ResidualFeedForward

View File

@@ -42,14 +42,16 @@ class LowRankBlock(torch.nn.Module):
"""
def __init__(self,
def __init__(
self,
input_dimensions,
embedding_dimenion,
rank,
inner_size=20,
n_layers=2,
func=torch.nn.Tanh,
bias=True):
bias=True,
):
"""
:param int input_dimensions: The number of input components of the
model.
@@ -78,10 +80,14 @@ class LowRankBlock(torch.nn.Module):
super().__init__()
# Assignment (check consistency inside FeedForward)
self._basis = pm.FeedForward(input_dimensions=input_dimensions,
self._basis = pm.FeedForward(
input_dimensions=input_dimensions,
output_dimensions=2 * rank * embedding_dimenion,
inner_size=inner_size, n_layers=n_layers,
func=func, bias=bias)
inner_size=inner_size,
n_layers=n_layers,
func=func,
bias=bias,
)
self._nn = torch.nn.Linear(embedding_dimenion, embedding_dimenion)
check_consistency(rank, int)
@@ -121,9 +127,9 @@ class LowRankBlock(torch.nn.Module):
psi = basis[..., : self.rank]
phi = basis[..., self.rank :]
# compute dot product
coeff = torch.einsum('...dr,...d->...r', psi,x)
coeff = torch.einsum("...dr,...d->...r", psi, x)
# expand the basis
expansion = torch.einsum('...r,...dr->...d', coeff,phi)
expansion = torch.einsum("...r,...dr->...d", coeff, phi)
# apply linear layer and return
return self._func(self._nn(x) + expansion)

View File

@@ -41,7 +41,7 @@ class LowRankNeuralOperator(KernelNeuralOperator):
inner_size=20,
n_layers=2,
func=torch.nn.Tanh,
bias=True
bias=True,
):
"""
:param torch.nn.Module lifting_net: The neural network for lifting
@@ -105,13 +105,18 @@ class LowRankNeuralOperator(KernelNeuralOperator):
self.coordinates_indices = coordinates_indices
self.field_indices = field_indices
integral_net = nn.Sequential(
*[LowRankBlock(input_dimensions=len(coordinates_indices),
*[
LowRankBlock(
input_dimensions=len(coordinates_indices),
embedding_dimenion=output_lifting_net,
rank=rank,
inner_size=inner_size,
n_layers=n_layers,
func=func,
bias=bias) for _ in range(n_kernel_layers)]
bias=bias,
)
for _ in range(n_kernel_layers)
]
)
super().__init__(lifting_net, integral_net, projecting_net)