fix embedding by removing linear layer

pina/model/layers/embedding.py

@@ -161,18 +161,17 @@ class PeriodicBoundaryEmbedding(torch.nn.Module):
 
 class FourierFeatureEmbedding(torch.nn.Module):
     def __init__(self,
-                 input_dimension : int,
+                 input_dimension,
-                 output_dimension : int,
+                 output_dimension,
-                 sigmas : Union[float, int, Sequence[float], Sequence[int]],
+                 sigma):
-                 embedding_output_dimension : int = None):
         r"""
         Fourier Feature Embedding class for encoding input features
         using random Fourier features.This class applies a Fourier
         transformation to the input features,
         which can help in learning high-frequency variations in data.
-        If multiple sigmas are provided, the class 
+        If multiple sigma are provided, the class 
         supports multiscale feature embedding, creating embeddings for
-        each scale specified by the sigmas.
+        each scale specified by the sigma.
 
         The :obj:`FourierFeatureEmbedding` augments the input
         by the following formula (3.10 of original paper):
@@ -184,7 +183,7 @@ class FourierFeatureEmbedding(torch.nn.Module):
 
         where :math:`\mathbf{B}_{ij} \sim \mathcal{N}(0, \sigma^2)`.
 
-        In case multiple ``sigmas`` are passed, the resulting embeddings
+        In case multiple ``sigma`` are passed, the resulting embeddings
         are concateneted:
 
         .. math::
@@ -211,45 +210,32 @@ class FourierFeatureEmbedding(torch.nn.Module):
             <https://doi.org/10.1016/j.cma.2021.113938>`_
 
         :param int input_dimension: The input vector dimension of the layer.
-        :param int output_dimension: The output dimension of the layer.
+        :param int output_dimension: The output dimension of the layer. The
-        :param sigmas: The standard deviation(s) used for the Fourier embedding.
+            output is obtained as a concatenation of the cosine and sine
-            This can be a single float or integer, or a sequence of floats
+            embedding, hence it must be a multiple of two (even number).
-            or integers. If a sequence is provided, the embedding will be
+        :param int | float sigma: The standard deviation used for the
-            computed for each sigma separately, enabling multiscale embeddings.
+            Fourier Embedding. This value must reflect the granularity of the
-        :type sigmas: Union[float, int, Sequence[float], Sequence[int]]
+            scale in the differential equation solution.
-        :param int output_dimension: The emebedding output dimension of the
-            random matrix use to compute the fourier feature. If ``None``, it
-            will be the same as ``output_dimension``, default ``None``.
         """
         super().__init__()
 
         # check consistency
-        check_consistency(sigmas, (int, float))
+        check_consistency(sigma, (int, float))
-        if isinstance(sigmas, (int, float)):
-            sigmas = [sigmas]
         check_consistency(output_dimension, int)
         check_consistency(input_dimension, int)
+        if output_dimension % 2:
+            raise RuntimeError('Expected output_dimension to be a even number, '
+                               f'got {output_dimension}.')
 
-        if embedding_output_dimension is None:
+        # assign sigma
-            embedding_output_dimension = output_dimension
+        self._sigma = sigma
-        check_consistency(embedding_output_dimension, int)
-
-        # assign
-        self.sigmas = sigmas
 
         # create non-trainable matrices
-        self._matrices = [
+        self._matrix = torch.rand(
-            torch.rand(
                         size = (input_dimension,
-                        embedding_output_dimension),
+                                output_dimension // 2),
-                requires_grad = False) * sigma for sigma in sigmas
+                        requires_grad = False
-                      ]
+                                )
-        
-        # create linear layer to map to the output dimension
-        self._linear = torch.nn.Linear(
-            in_features=2*len(sigmas)*embedding_output_dimension,
-            out_features=output_dimension)
-
 
     def forward(self, x):
         """
@@ -260,8 +246,14 @@ class FourierFeatureEmbedding(torch.nn.Module):
         :rtype: torch.Tensor
         """
         # compute random matrix multiplication
-        out = torch.cat([torch.mm(x, m) for m in self._matrices], dim=-1)
+        out = torch.mm(x, self._matrix)
-        # compute cos/sin emebedding
+        # return embedding
-        out = torch.cat([torch.cos(out), torch.sin(out)], dim=-1)
+        return torch.cat([torch.cos(out), torch.sin(out)], dim=-1)
-        # return linear layer mapping
+    
-        return self._linear(out)
+
+    @property
+    def sigma(self):
+        """
+        Returning the variance of the sampled matrix for Fourier Embedding.
+        """
+        return self._sigma

tests/test_layers/test_embedding.py

@@ -65,35 +65,27 @@ def test_forward_backward_same_period_PeriodicBoundaryEmbedding(input_dimension,
 
 def test_constructor_FourierFeatureEmbedding():
     FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
-                            sigmas=1)
+                            sigma=1)
-    FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
-                            sigmas=[0.01, 0.1, 1])
-    FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
-                            sigmas=[0.01, 0.1, 1])
-    FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
-                            sigmas=1, embedding_output_dimension=20)
     with pytest.raises(TypeError): 
         FourierFeatureEmbedding()
+    with pytest.raises(RuntimeError): 
+        FourierFeatureEmbedding(input_dimension=1, output_dimension=3, sigma=1)
     with pytest.raises(ValueError):
         FourierFeatureEmbedding(input_dimension='x', output_dimension=20,
-                                sigmas=1)
+                                sigma=1)
         FourierFeatureEmbedding(input_dimension=1, output_dimension='x',
-                                sigmas=1)
+                                sigma=1)
         FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
-                                sigmas='x')
+                                sigma='x')
-        FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
-                                sigmas=1, embedding_output_dimension='x')
 
-@pytest.mark.parametrize("output_dimension", [1, 2, 2])
+@pytest.mark.parametrize("output_dimension", [2, 4, 6])
 @pytest.mark.parametrize("input_dimension", [1, 2, 3])
-@pytest.mark.parametrize("sigmas", [1, [0.01, 0.1, 1]])
+@pytest.mark.parametrize("sigma", [10, 1, 0.1])
-@pytest.mark.parametrize("embedding_output_dimension", [1, 2, 3])
 def test_forward_backward_FourierFeatureEmbedding(input_dimension,
                                                   output_dimension,
-                                         sigmas,
+                                                  sigma):
-                                         embedding_output_dimension):
     func = FourierFeatureEmbedding(input_dimension, output_dimension,
-                                   sigmas, embedding_output_dimension)
+                                   sigma)
     # coordinates
     x = torch.rand((10, input_dimension), requires_grad=True)
     # output