* Fourier feature embedding added, and modify typos in doc Periodic Boundary Embedding.

* Fixing doc for Periodic Boundary Embedding.
* Creating doc for Fourier Feature Embedding.

docs/source/_rst/_code.rst

@@ -79,7 +79,8 @@ Layers
     Low Rank layer <layers/lowrank_layer.rst>
     Continuous convolution <layers/convolution.rst>
     Proper Orthogonal Decomposition <layers/pod.rst>
-    Periodic Boundary Condition embeddings <layers/embedding.rst>
+    Periodic Boundary Condition Embedding <layers/pbc_embedding.rst>
+    Fourier Feature Embedding <layers/fourier_embedding.rst>
 
 Adaptive Activation Functions
 -------------------------------

docs/source/_rst/layers/fourier_embedding.rst

@@ -0,0 +1,8 @@
+Fourier Feature Embedding
+=======================================
+.. currentmodule:: pina.model.layers.embedding
+    
+.. autoclass:: FourierFeatureEmbedding
+    :members:
+    :show-inheritance:
+

docs/source/_rst/layers/pbc_embedding.rst

@@ -1,4 +1,4 @@
-Periodic Boundary Condition embeddings
+Periodic Boundary Condition Embedding
 =======================================
 .. currentmodule:: pina.model.layers.embedding
     

pina/model/layers/__init__.py

@@ -10,6 +10,7 @@ __all__ = [
     "FourierBlock3D",
     "PODBlock",
     "PeriodicBoundaryEmbedding",
+    "FourierFeatureEmbedding",
     "AVNOBlock",
     "LowRankBlock",
 ]
@@ -23,6 +24,6 @@ from .spectral import (
 )
 from .fourier import FourierBlock1D, FourierBlock2D, FourierBlock3D
 from .pod import PODBlock
-from .embedding import PeriodicBoundaryEmbedding
+from .embedding import PeriodicBoundaryEmbedding, FourierFeatureEmbedding
 from .avno_layer import AVNOBlock
 from .lowrank_layer import LowRankBlock

pina/model/layers/embedding.py

@@ -1,7 +1,8 @@
-""" Periodic Boundary Embedding modulus. """
+""" Embedding modulus. """
 
 import torch
 from pina.utils import check_consistency
+from typing import Union, Sequence
 
 
 class PeriodicBoundaryEmbedding(torch.nn.Module):
@@ -100,7 +101,7 @@ class PeriodicBoundaryEmbedding(torch.nn.Module):
         Forward pass to compute the periodic boundary conditions embedding.
 
         :param torch.Tensor x: Input tensor.
-        :return: Fourier embeddings of the input.
+        :return: Periodic embedding of the input.
         :rtype: torch.Tensor
         """
         omega = torch.stack(
@@ -155,3 +156,112 @@ class PeriodicBoundaryEmbedding(torch.nn.Module):
         The period of the periodic function to approximate.
         """
         return self._period
+
+
+
+class FourierFeatureEmbedding(torch.nn.Module):
+    def __init__(self,
+                 input_dimension : int,
+                 output_dimension : int,
+                 sigmas : Union[float, int, Sequence[float], Sequence[int]],
+                 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 
+        supports multiscale feature embedding, creating embeddings for
+        each scale specified by the sigmas.
+
+        The :obj:`FourierFeatureEmbedding` augments the input
+        by the following formula (3.10 of original paper):
+
+        .. math::
+            \mathbf{x} \rightarrow \tilde{\mathbf{x}} = \left[
+            \cos\left( \mathbf{B} \mathbf{x} \right),
+            \sin\left( \mathbf{B} \mathbf{x} \right)\right],
+
+        where :math:`\mathbf{B}_{ij} \sim \mathcal{N}(0, \sigma^2)`.
+
+        In case multiple ``sigmas`` are passed, the resulting embeddings
+        are concateneted:
+
+        .. math::
+            \mathbf{x} \rightarrow \tilde{\mathbf{x}} = \left[
+            \cos\left( \mathbf{B}^1 \mathbf{x} \right),
+            \sin\left( \mathbf{B}^1 \mathbf{x} \right),
+            \cos\left( \mathbf{B}^2 \mathbf{x} \right),
+            \sin\left( \mathbf{B}^3 \mathbf{x} \right),
+            \dots,
+            \cos\left( \mathbf{B}^M \mathbf{x} \right),
+            \sin\left( \mathbf{B}^M \mathbf{x} \right)\right],
+
+        where :math:`\mathbf{B}^k_{ij} \sim \mathcal{N}(0, \sigma_k^2) \quad
+        k \in (1, \dots, M)`.
+
+        .. seealso::
+            **Original reference**:
+            Wang, Sifan, Hanwen Wang, and Paris Perdikaris. *On the eigenvector
+            bias of Fourier feature networks: From regression to solving
+            multi-scale PDEs with physics-informed neural networks.*
+            Computer Methods in Applied Mechanics and
+            Engineering 384 (2021): 113938.
+            DOI: `10.1016/j.cma.2021.113938.
+            <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 sigmas: The standard deviation(s) used for the Fourier embedding.
+            This can be a single float or integer, or a sequence of floats
+            or integers. If a sequence is provided, the embedding will be
+            computed for each sigma separately, enabling multiscale embeddings.
+        :type sigmas: Union[float, int, Sequence[float], Sequence[int]]
+        :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))
+        if isinstance(sigmas, (int, float)):
+            sigmas = [sigmas]
+        check_consistency(output_dimension, int)
+        check_consistency(input_dimension, int)
+
+        if embedding_output_dimension is None:
+            embedding_output_dimension = output_dimension
+        check_consistency(embedding_output_dimension, int)
+
+        # assign
+        self.sigmas = sigmas
+
+        # create non-trainable matrices
+        self._matrices = [
+            torch.rand(
+                size = (input_dimension,
+                        embedding_output_dimension),
+                requires_grad = False) * sigma for sigma in sigmas
+                      ]
+        
+        # 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):
+        """
+        Forward pass to compute the fourier embedding.
+
+        :param torch.Tensor x: Input tensor.
+        :return: Fourier embeddings of the input.
+        :rtype: torch.Tensor
+        """
+        # compute random matrix multiplication
+        out = torch.cat([torch.mm(x, m) for m in self._matrices], dim=-1)
+        # compute cos/sin emebedding
+        out = torch.cat([torch.cos(out), torch.sin(out)], dim=-1)
+        # return linear layer mapping
+        return self._linear(out)

tests/test_layers/test_embedding.py

@@ -1,8 +1,7 @@
 import torch
 import pytest
 
-from pina.model.layers import PeriodicBoundaryEmbedding
+from pina.model.layers import PeriodicBoundaryEmbedding, FourierFeatureEmbedding
-from pina import LabelTensor
 
 # test tolerance
 tol = 1e-6
@@ -23,7 +22,7 @@ def grad(u, x):
                                create_graph=True, allow_unused=True,
                                retain_graph=True)[0]
 
-def test_constructor():
+def test_constructor_PeriodicBoundaryEmbedding():
     PeriodicBoundaryEmbedding(input_dimension=1, periods=2)
     PeriodicBoundaryEmbedding(input_dimension=1, periods={'x': 3, 'y' : 4})
     PeriodicBoundaryEmbedding(input_dimension=1, periods={0: 3, 1 : 4})
@@ -32,14 +31,16 @@ def test_constructor():
         PeriodicBoundaryEmbedding()
     with pytest.raises(ValueError):
         PeriodicBoundaryEmbedding(input_dimension=1., periods=1)
-        PeriodicBoundaryEmbedding(input_dimension=1, periods=1, output_dimension=1.)
+        PeriodicBoundaryEmbedding(input_dimension=1, periods=1,
+                                  output_dimension=1.)
         PeriodicBoundaryEmbedding(input_dimension=1, periods={'x':'x'})
         PeriodicBoundaryEmbedding(input_dimension=1, periods={0:'x'})
 
 
 @pytest.mark.parametrize("period", [1, 4, 10])
 @pytest.mark.parametrize("input_dimension", [1, 2, 3])
-def test_forward_same_period(input_dimension, period):
+def test_forward_backward_same_period_PeriodicBoundaryEmbedding(input_dimension,
+                                                                period):
     func = torch.nn.Sequential(
         PeriodicBoundaryEmbedding(input_dimension=input_dimension,
                      output_dimension=60, periods=period),
@@ -58,46 +59,46 @@ def test_forward_same_period(input_dimension, period):
     # output
     f = func(x)
     assert check_same_columns(f)
+    # compute backward
+    loss = f.mean()
+    loss.backward()
 
+def test_constructor_FourierFeatureEmbedding():
+    FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
+                            sigmas=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(ValueError):
+        FourierFeatureEmbedding(input_dimension='x', output_dimension=20,
+                                sigmas=1)
+        FourierFeatureEmbedding(input_dimension=1, output_dimension='x',
+                                sigmas=1)
+        FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
+                                sigmas='x')
+        FourierFeatureEmbedding(input_dimension=1, output_dimension=20,
+                                sigmas=1, embedding_output_dimension='x')
 
-
+@pytest.mark.parametrize("output_dimension", [1, 2, 2])
-# def test_forward_same_period_labels():
+@pytest.mark.parametrize("input_dimension", [1, 2, 3])
-#     func = torch.nn.Sequential(
+@pytest.mark.parametrize("sigmas", [1, [0.01, 0.1, 1]])
-#         PeriodicBoundaryEmbedding(input_dimension=2,
+@pytest.mark.parametrize("embedding_output_dimension", [1, 2, 3])
-#                      output_dimension=60, periods={'x':1, 'y':2}),
+def test_forward_backward_FourierFeatureEmbedding(input_dimension,
-#         torch.nn.Tanh(),
+                                         output_dimension,
-#         torch.nn.Linear(60, 60),
+                                         sigmas,
-#         torch.nn.Tanh(),
+                                         embedding_output_dimension):
-#         torch.nn.Linear(60, 1)
+    func = FourierFeatureEmbedding(input_dimension, output_dimension,
-#     )
+                                   sigmas, embedding_output_dimension)
-#     # coordinates
+    # coordinates
-#     tensor = torch.tensor([[0., 0.], [0., 2.], [1., 0.], [1., 2.]])
+    x = torch.rand((10, input_dimension), requires_grad=True)
-#     with pytest.raises(RuntimeError):
+    # output
-#         func(tensor)
+    f = func(x)
-#     tensor = tensor.as_subclass(LabelTensor)
+    assert f.shape[-1] == output_dimension
-#     tensor.labels = ['x', 'y']
+    # compute backward
-#     tensor.requires_grad = True
+    loss = f.mean()
-#     # output
+    loss.backward()
-#     f = func(tensor)
-#     assert check_same_columns(f)
-
-# def test_forward_same_period_index():
-#     func = torch.nn.Sequential(
-#         PeriodicBoundaryEmbedding(input_dimension=2,
-#                      output_dimension=60, periods={0:1, 1:2}),
-#         torch.nn.Tanh(),
-#         torch.nn.Linear(60, 60),
-#         torch.nn.Tanh(),
-#         torch.nn.Linear(60, 1)
-#     )
-#     # coordinates
-#     tensor = torch.tensor([[0., 0.], [0., 2.], [1., 0.], [1., 2.]])
-#     tensor.requires_grad = True
-#     # output
-#     f = func(tensor)
-#     assert check_same_columns(f)
-#     tensor = tensor.as_subclass(LabelTensor)
-#     tensor.labels = ['x', 'y']
-#     # output
-#     f = func(tensor)
-#     assert check_same_columns(f)