🎨 Format Python code with psf/black

pina/model/__init__.py

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

pina/model/layers/lowrank_layer.py

@@ -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,
+            inner_size=inner_size,
-                                    func=func, bias=bias)
+            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)
 

pina/model/lno.py

@@ -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)