🎨 Format Python code with psf/black

pina/model/layers/__init__.py

@@ -9,7 +9,7 @@ __all__ = [
     "FourierBlock2D",
     "FourierBlock3D",
     "PODBlock",
-    "PeriodicBoundaryEmbedding"
+    "PeriodicBoundaryEmbedding",
 ]
 
 from .convolution_2d import ContinuousConvBlock

pina/model/layers/embedding.py

@@ -52,6 +52,7 @@ class PeriodicBoundaryEmbedding(torch.nn.Module):
         :math:`>3`. The PINA code is tested only for function PBC and not for
         its derivatives.
     """
+
     def __init__(self, input_dimension, periods, output_dimension=None):
         """
         :param int input_dimension: The dimension of the input tensor, it can
@@ -82,16 +83,18 @@ class PeriodicBoundaryEmbedding(torch.nn.Module):
 
         # checks on the periods
         if isinstance(periods, dict):
-            if not all(isinstance(dim, (str, int)) and
+            if not all(
-                       isinstance(period, (float, int))
+                isinstance(dim, (str, int)) and isinstance(period, (float, int))
-                       for dim, period in periods.items()):
+                for dim, period in periods.items()
-                raise TypeError('In dictionary periods, keys must be integers'
+            ):
-                                ' or strings, and values must be float or int.')
+                raise TypeError(
+                    "In dictionary periods, keys must be integers"
+                    " or strings, and values must be float or int."
+                )
             self._period = periods
         else:
             self._period = {k: periods for k in range(input_dimension)}
 
-
     def forward(self, x):
         """
         Forward pass to compute the periodic boundary conditions embedding.
@@ -100,14 +103,24 @@ class PeriodicBoundaryEmbedding(torch.nn.Module):
         :return: Fourier embeddings of the input.
         :rtype: torch.Tensor
         """
-        omega = torch.stack([torch.pi * 2. / torch.tensor([val],
+        omega = torch.stack(
-                                                          device=x.device)
+            [
-                                   for val in self._period.values()],
+                torch.pi * 2.0 / torch.tensor([val], device=x.device)
-                                   dim=-1)
+                for val in self._period.values()
+            ],
+            dim=-1,
+        )
         x = self._get_vars(x, list(self._period.keys()))
-        return self._layer(torch.cat([torch.ones_like(x),
+        return self._layer(
+            torch.cat(
+                [
+                    torch.ones_like(x),
                     torch.cos(omega * x),
-                          torch.sin(omega * x)], dim=-1))
+                    torch.sin(omega * x),
+                ],
+                dim=-1,
+            )
+        )
 
     def _get_vars(self, x, indeces):
         """
@@ -123,16 +136,18 @@ class PeriodicBoundaryEmbedding(torch.nn.Module):
                 return x.extract(indeces)
             except AttributeError:
                 raise RuntimeError(
-                    'Not possible to extract input variables from tensor.'
+                    "Not possible to extract input variables from tensor."
-                    ' Ensure that the passed tensor is a LabelTensor or'
+                    " Ensure that the passed tensor is a LabelTensor or"
-                    ' pass list of integers to extract variables. For'
+                    " pass list of integers to extract variables. For"
-                    ' more information refer to warning in the documentation.')
+                    " more information refer to warning in the documentation."
+                )
         elif isinstance(indeces[0], int):
             return x[..., indeces]
         else:
             raise RuntimeError(
-                'Not able to extract right indeces for tensor.'
+                "Not able to extract right indeces for tensor."
-                ' For more information refer to warning in the documentation.')
+                " For more information refer to warning in the documentation."
+            )
 
     @property
     def period(self):