Added Gauss-Legendre Quadrature Fourier features

.github/workflows/quality-check.yaml

@@ -21,7 +21,7 @@ jobs:
     runs-on: ubuntu-latest
     strategy:
       matrix:
-          python-version: [3.6, 3.7, 3.8]
+          python-version: [3.7, 3.8]
           tensorflow: ["~=2.5.0"]
     name: Python-${{ matrix.python-version }} tensorflow${{ matrix.tensorflow }}
     env:

Makefile

@@ -20,6 +20,9 @@ LINT_FILE_IGNORES = "$(LIB_NAME)/__init__.py:F401,F403 \
                      $(LIB_NAME)/initializers/__init__.py:F401 \
                      $(LIB_NAME)/layers/__init__.py:F401 \
                      $(LIB_NAME)/layers/basis_functions/__init__.py:F401 \
+                     $(LIB_NAME)/layers/basis_functions/fourier_features/__init__.py:F401 \
+                     $(LIB_NAME)/layers/basis_functions/fourier_features/random/__init__.py:F401 \
+                     $(LIB_NAME)/layers/basis_functions/fourier_features/quadrature/__init__.py:F401 \
                      $(LIB_NAME)/models/__init__.py:F401 \
                      $(LIB_NAME)/optimization/__init__.py:F401 \
                      $(LIB_NAME)/sampling/__init__.py:F401 \

gpflux/layers/basis_functions/fourier_features/__init__.py

@@ -18,16 +18,15 @@
 :class:`gpflux.sampling.KernelWithFeatureDecomposition`
 """
 
-from gpflux.layers.basis_functions.fourier_features.quadrature import QuadratureFourierFeatures
+from gpflux.layers.basis_functions.fourier_features.quadrature import (
+    GaussHermiteQuadratureFourierFeatures,
+    GaussLegendreQuadratureFourierFeatures,
+    QuadratureFourierFeatures,
+    SimpsonQuadratureFourierFeatures,
+)
 from gpflux.layers.basis_functions.fourier_features.random import (
     OrthogonalRandomFeatures,
+    QuasiRandomFourierFeatures,
     RandomFourierFeatures,
     RandomFourierFeaturesCosine,
 )
-
-__all__ = [
-    "QuadratureFourierFeatures",
-    "OrthogonalRandomFeatures",
-    "RandomFourierFeatures",
-    "RandomFourierFeaturesCosine",
-]

gpflux/layers/basis_functions/fourier_features/base.py

@@ -16,7 +16,7 @@
 """ Shared functionality for stationary kernel basis functions. """
 
 from abc import ABC, abstractmethod
-from typing import Mapping
+from typing import Mapping, Tuple, Type
 
 import tensorflow as tf
 
@@ -27,13 +27,19 @@
 
 
 class FourierFeaturesBase(ABC, tf.keras.layers.Layer):
+
+    SUPPORTED_KERNELS: Tuple[Type[gpflow.kernels.Stationary], ...]
+
     def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
         """
         :param kernel: kernel to approximate using a set of Fourier bases.
         :param n_components: number of components (e.g. Monte Carlo samples,
             quadrature nodes, etc.) used to numerically approximate the kernel.
         """
         super(FourierFeaturesBase, self).__init__(**kwargs)
+        assert isinstance(
+            kernel, self.SUPPORTED_KERNELS
+        ), f"Only the following kernels are supported: {self.SUPPORTED_KERNELS}"
         self.kernel = kernel
         self.n_components = n_components
         if kwargs.get("input_dim", None):

gpflux/layers/basis_functions/fourier_features/quadrature/__init__.py

@@ -15,7 +15,10 @@
 #
 """ A kernel's features and coefficients using quadrature Fourier features (QFF). """
 from gpflux.layers.basis_functions.fourier_features.quadrature.gaussian import (
+    GaussHermiteQuadratureFourierFeatures,
+    GaussLegendreQuadratureFourierFeatures,
     QuadratureFourierFeatures,
 )
-
-__all__ = ["QuadratureFourierFeatures"]
+from gpflux.layers.basis_functions.fourier_features.quadrature.newton_cotes import (
+    SimpsonQuadratureFourierFeatures,
+)

gpflux/layers/basis_functions/fourier_features/quadrature/base.py

@@ -0,0 +1,98 @@
+#
+# Copyright (c) 2021 The GPflux Contributors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+
+from abc import ABC, abstractmethod
+
+import numpy as np
+import tensorflow as tf
+
+from gpflow.base import TensorType
+
+from gpflux.layers.basis_functions.fourier_features.base import FourierFeaturesBase
+from gpflux.layers.basis_functions.fourier_features.utils import _bases_concat
+from gpflux.types import ShapeType
+
+
+class QuadratureFourierFeaturesBase(FourierFeaturesBase):
+    def _compute_output_dim(self, input_shape: ShapeType) -> int:
+        input_dim = input_shape[-1]
+        return 2 * self.n_components ** input_dim
+
+    def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
+        """
+        Compute basis functions.
+
+        :return: A tensor with the shape ``(N, 2L^D)``.
+        """
+        return _bases_concat(inputs, self.abscissa)
+
+    def _compute_constant(self) -> tf.Tensor:
+        """
+        Compute normalizing constant for basis functions.
+
+        :return: A tensor with the shape ``(2L^D,)``
+        """
+        return tf.tile(tf.sqrt(self.kernel.variance * self.factors), multiples=[2])
+
+
+class Transform(ABC):
+    r"""
+    This class encapsulates functions :math:`g(x) = u` and :math:`h(x)` such that
+    .. math::
+        \int_{g(a)}^{g(b)} z(u) f(u) du
+        = \int_a^b w(x) h(x) f(g(x)) dx
+    for some integrand :math:`f(u)` and weight function :math:`z(u)`.
+    """
+
+    @abstractmethod
+    def __call__(self, x: TensorType) -> tf.Tensor:
+        pass
+
+    @abstractmethod
+    def multiplier(self, x: TensorType) -> tf.Tensor:
+        pass
+
+
+class TanTransform(Transform):
+    r"""
+    This class encapsulates functions :math:`g(x) = u` and :math:`h(x) = du/dx`
+    such that
+    .. math::
+        \int_{-\infty}^{\infty} f(u) du
+        = \int_{-1}^{1} f(g(x)) h(x) dx
+    """
+    CONST = 0.5 * np.pi
+
+    def __call__(self, x: TensorType) -> tf.Tensor:
+        return tf.tan(TanTransform.CONST * x)
+
+    def multiplier(self, x: TensorType) -> tf.Tensor:
+        return TanTransform.CONST / tf.square(tf.cos(TanTransform.CONST * x))
+
+
+class NormalWeightTransform(Transform):
+    r"""
+    This class encapsulates functions :math:`g(x) = u` and :math:`h(x)` such that
+    .. math::
+        \int_{-\infty}^{\infty} \mathcal{N}(u|0,1) f(u) du
+        = \int_{-infty}^{infty} e^{-x^2} f(g(x)) h(x) dx
+    """
+
+    def __call__(self, x: TensorType) -> tf.Tensor:
+        return tf.sqrt(2.0) * x
+
+    def multiplier(self, x: TensorType) -> tf.Tensor:
+        return tf.rsqrt(np.pi)

gpflux/layers/basis_functions/fourier_features/quadrature/gaussian.py

@@ -19,39 +19,39 @@
 """
 
 import warnings
-from typing import Mapping, Tuple, Type
+from typing import Mapping
 
+import numpy as np
 import tensorflow as tf
+from scipy.stats import multivariate_normal, multivariate_t
 
 import gpflow
-from gpflow.base import TensorType
-from gpflow.quadrature.gauss_hermite import ndgh_points_and_weights
 
-from gpflux.layers.basis_functions.fourier_features.base import FourierFeaturesBase
-from gpflux.layers.basis_functions.fourier_features.utils import _bases_concat
-from gpflux.types import ShapeType
-
-"""
-Kernels supported by :class:`QuadratureFourierFeatures`.
+# from gpflow.config import default_float
+from gpflow.quadrature.gauss_hermite import ndgh_points_and_weights, repeat_as_list, reshape_Z_dZ
 
-Currently we only support the :class:`gpflow.kernels.SquaredExponential` kernel.
-For Matern kernels please use :class:`RandomFourierFeatures`
-or :class:`RandomFourierFeaturesCosine`.
-"""
-QFF_SUPPORTED_KERNELS: Tuple[Type[gpflow.kernels.Stationary], ...] = (
-    gpflow.kernels.SquaredExponential,
+from gpflux.layers.basis_functions.fourier_features.quadrature.base import (
+    QuadratureFourierFeaturesBase,
+    TanTransform,
 )
+from gpflux.layers.basis_functions.fourier_features.utils import _matern_dof
+from gpflux.types import ShapeType
 
 
-class QuadratureFourierFeatures(FourierFeaturesBase):
+class GaussianQuadratureFourierFeatures(QuadratureFourierFeaturesBase):
     def __init__(self, kernel: gpflow.kernels.Kernel, n_components: int, **kwargs: Mapping):
-        assert isinstance(kernel, QFF_SUPPORTED_KERNELS), "Unsupported Kernel"
+        super(GaussianQuadratureFourierFeatures, self).__init__(kernel, n_components, **kwargs)
         if tf.reduce_any(tf.less(kernel.lengthscales, 1e-1)):
             warnings.warn(
-                "Quadrature Fourier feature approximation of kernels "
-                "with small lengthscale lead to unexpected behaviors!"
+                "Fourier feature approximation of kernels with small "
+                "lengthscales using Gaussian quadrature can have "
+                "unexpected behaviors!"
             )
-        super(QuadratureFourierFeatures, self).__init__(kernel, n_components, **kwargs)
+
+
+class GaussHermiteQuadratureFourierFeatures(GaussianQuadratureFourierFeatures):
+
+    SUPPORTED_KERNELS = (gpflow.kernels.SquaredExponential,)
 
     def build(self, input_shape: ShapeType) -> None:
         """
@@ -60,33 +60,72 @@ def build(self, input_shape: ShapeType) -> None:
         <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#build>`_.
         """
         input_dim = input_shape[-1]
-        abscissa_value, omegas_value = ndgh_points_and_weights(
+        # (L^D, D), (L^D, 1)
+        abscissa_value, factors_value = ndgh_points_and_weights(
             dim=input_dim, n_gh=self.n_components
         )
-        omegas_value = tf.squeeze(omegas_value, axis=-1)
+        factors_value = tf.squeeze(factors_value, axis=-1)  # (L^D,)
 
-        # Quadrature node points
-        self.abscissa = tf.Variable(initial_value=abscissa_value, trainable=False)  # (M^D, D)
-        # Gauss-Hermite weights
-        self.factors = tf.Variable(initial_value=omegas_value, trainable=False)  # (M^D,)
-        super(QuadratureFourierFeatures, self).build(input_shape)
+        # Gauss-Christoffel nodes (L^D, D)
+        self.abscissa = tf.Variable(initial_value=abscissa_value, trainable=False)
+        # Gauss-Christoffel weights (L^D,)
+        self.factors = tf.Variable(initial_value=factors_value, trainable=False)
+        super(GaussHermiteQuadratureFourierFeatures, self).build(input_shape)
 
-    def _compute_output_dim(self, input_shape: ShapeType) -> int:
-        input_dim = input_shape[-1]
-        return 2 * self.n_components ** input_dim
 
-    def _compute_bases(self, inputs: TensorType) -> tf.Tensor:
-        """
-        Compute basis functions.
+class GaussLegendreQuadratureFourierFeatures(GaussianQuadratureFourierFeatures):
 
-        :return: A tensor with the shape ``[N, 2M^D]``.
-        """
-        return _bases_concat(inputs, self.abscissa)
+    SUPPORTED_KERNELS = (
+        gpflow.kernels.SquaredExponential,
+        gpflow.kernels.Matern12,
+        gpflow.kernels.Matern32,
+        gpflow.kernels.Matern52,
+    )
 
-    def _compute_constant(self) -> tf.Tensor:
+    def build(self, input_shape: ShapeType) -> None:
         """
-        Compute normalizing constant for basis functions.
-
-        :return: A tensor with the shape ``[2M^D,]``
+        Creates the variables of the layer.
+        See `tf.keras.layers.Layer.build()
+        <https://www.tensorflow.org/api_docs/python/tf/keras/layers/Layer#build>`_.
         """
-        return tf.tile(tf.sqrt(self.kernel.variance * self.factors), multiples=[2])
+        input_dim = input_shape[-1]
+
+        if isinstance(self.kernel, gpflow.kernels.SquaredExponential):
+            dist = multivariate_normal(mean=np.zeros(input_dim))
+        else:
+            nu = _matern_dof(self.kernel)  # degrees of freedom
+            dist = multivariate_t(loc=np.zeros(input_dim), df=nu)
+
+        # raw 1-dimensional quadrature nodes and weights (L,) (L,)
+        abscissa_value_flat, factors_value_flat = np.polynomial.legendre.leggauss(
+            deg=self.n_components
+        )
+
+        # transformed 1-dimensional quadrature nodes and weights
+        transform = TanTransform()
+        factors_value_flat *= transform.multiplier(abscissa_value_flat)  # (L,)
+        abscissa_value_flat = transform(abscissa_value_flat)  # (L,)
+
+        # transformed D-dimensional quadrature nodes and weights
+        abscissa_value_rep = repeat_as_list(abscissa_value_flat, n=input_dim)  # (L, ..., L)
+        factors_value_rep = repeat_as_list(factors_value_flat, n=input_dim)  # (L, ..., L)
+        # (L^D, D), (L^D, 1)
+        abscissa_value, factors_value = reshape_Z_dZ(abscissa_value_rep, factors_value_rep)
+
+        factors_value = tf.squeeze(factors_value, axis=-1)  # (L^D,)
+        factors_value *= dist.pdf(abscissa_value)  # (L^D,)
+
+        # Gauss-Christoffel nodes (L^D, D)
+        self.abscissa = tf.Variable(initial_value=abscissa_value, trainable=False)
+        # Gauss-Christoffel weights (L^D,)
+        self.factors = tf.Variable(initial_value=factors_value, trainable=False)
+
+        super(GaussLegendreQuadratureFourierFeatures, self).build(input_shape)
+
+
+class QuadratureFourierFeatures(GaussHermiteQuadratureFourierFeatures):
+    """
+    Alias for `GaussHermiteQuadratureFourierFeatures`.
+    """
+
+    pass