🎨 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/__init__.py

@@ -27,4 +27,4 @@ from .pod import PODBlock
 from .embedding import PeriodicBoundaryEmbedding
 from .avno_layer import AVNOBlock
 from .lowrank_layer import LowRankBlock
-from .adaptive_func import AdaptiveActivationFunction
+from .adaptive_func import AdaptiveActivationFunction

pina/model/layers/lowrank_layer.py

@@ -3,7 +3,7 @@
 import torch
 
 from pina.utils import check_consistency
-import pina.model as pm                     # avoid circular import
+import pina.model as pm  # avoid circular import
 
 
 class LowRankBlock(torch.nn.Module):
@@ -42,14 +42,16 @@ class LowRankBlock(torch.nn.Module):
 
     """
 
-    def __init__(self,
-                 input_dimensions,
-                 embedding_dimenion,
-                 rank,
-                 inner_size=20,
-                 n_layers=2,
-                 func=torch.nn.Tanh,
-                 bias=True):
+    def __init__(
+        self,
+        input_dimensions,
+        embedding_dimenion,
+        rank,
+        inner_size=20,
+        n_layers=2,
+        func=torch.nn.Tanh,
+        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,
-                                    output_dimensions=2*rank*embedding_dimenion,
-                                    inner_size=inner_size, n_layers=n_layers,
-                                    func=func, bias=bias)
+        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,
+        )
         self._nn = torch.nn.Linear(embedding_dimenion, embedding_dimenion)
 
         check_consistency(rank, int)
@@ -115,15 +121,15 @@ class LowRankBlock(torch.nn.Module):
         # extract basis
         basis = self._basis(coords)
         # reshape [B, N, D, 2*rank]
-        shape = list(basis.shape[:-1]) + [-1, 2*self.rank]
+        shape = list(basis.shape[:-1]) + [-1, 2 * self.rank]
         basis = basis.reshape(shape)
         # divide
-        psi = basis[..., :self.rank]
-        phi = basis[..., self.rank:]
+        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),
-                           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)]
+            *[
+                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)
+            ]
         )
         super().__init__(lifting_net, integral_net, projecting_net)