Efficient implementation of arbitrary half-integer Matérn Kernels (#771)

* Efficient implementation of all half integer Matérn kernels

* Dimensionwise lengthscales in `ExpQuad`

* Fix tests

* `isort` fixes

* Add some more tests for the Matérn

* Add an additional test for the `ExpQuad` kernel

* Add tests for `Kernel.input_size`

* Docstring improvements

src/probnum/randprocs/kernels/_exponentiated_quadratic.py

@@ -1,11 +1,11 @@
 """Exponentiated quadratic kernel."""
 
+import functools
 from typing import Optional
 
 import numpy as np
 
-from probnum.typing import ScalarLike, ShapeLike
-import probnum.utils as _utils
+from probnum.typing import ArrayLike, ShapeLike
 
 from ._kernel import IsotropicMixin, Kernel
 
@@ -38,18 +38,48 @@ class ExpQuad(Kernel, IsotropicMixin):
     --------
     >>> import numpy as np
     >>> from probnum.randprocs.kernels import ExpQuad
-    >>> K = ExpQuad(input_shape=(), lengthscale=0.1)
+    >>> K = ExpQuad((), lengthscales=0.1)
     >>> xs = np.linspace(0, 1, 3)
     >>> K.matrix(xs)
     array([[1.00000000e+00, 3.72665317e-06, 1.92874985e-22],
            [3.72665317e-06, 1.00000000e+00, 3.72665317e-06],
            [1.92874985e-22, 3.72665317e-06, 1.00000000e+00]])
     """
 
-    def __init__(self, input_shape: ShapeLike, lengthscale: ScalarLike = 1.0):
-        self.lengthscale = _utils.as_numpy_scalar(lengthscale)
+    def __init__(self, input_shape: ShapeLike, *, lengthscales: ArrayLike = 1.0):
         super().__init__(input_shape=input_shape)
 
+        # Input lengthscales
+        lengthscales = np.asarray(
+            lengthscales if lengthscales is not None else 1.0, dtype=np.double
+        )
+
+        if np.any(lengthscales <= 0):
+            raise ValueError(f"Lengthscales l={lengthscales} must be positive.")
+
+        np.broadcast_to(  # Check if the lengthscales broadcast to the input dimension
+            lengthscales, self.input_shape
+        )
+
+        self._lengthscales = lengthscales
+
+    @property
+    def lengthscales(self) -> np.ndarray:
+        r"""Input lengthscales along the different input dimensions."""
+        return self._lengthscales
+
+    @property
+    def lengthscale(self) -> np.ndarray:
+        """Deprecated."""
+        if self._lengthscales.shape == ():
+            return self._lengthscales
+
+        raise ValueError("There is more than one lengthscale.")
+
+    @functools.cached_property
+    def _scale_factors(self) -> np.ndarray:
+        return np.sqrt(0.5) / self._lengthscales
+
     def _evaluate(self, x0: np.ndarray, x1: Optional[np.ndarray] = None) -> np.ndarray:
         if x1 is None:
             return np.ones_like(  # pylint: disable=unexpected-keyword-arg
@@ -58,5 +88,7 @@ def _evaluate(self, x0: np.ndarray, x1: Optional[np.ndarray] = None) -> np.ndarr
             )
 
         return np.exp(
-            -self._squared_euclidean_distances(x0, x1) / (2.0 * self.lengthscale**2)
+            -self._squared_euclidean_distances(
+                x0, x1, scale_factors=self._scale_factors
+            )
         )

src/probnum/randprocs/kernels/_kernel.py

@@ -3,6 +3,8 @@
 from __future__ import annotations
 
 import abc
+import functools
+import operator
 from typing import Optional, Union
 
 import numpy as np
@@ -159,6 +161,11 @@ def input_ndim(self) -> int:
         """Syntactic sugar for ``len(input_shape)``."""
         return self._input_ndim
 
+    @property
+    def input_size(self) -> int:
+        """Syntactic sugar for the product over the input size."""
+        return functools.reduce(operator.mul, self.input_shape, 1)
+
     @property
     def output_shape(self) -> ShapeType:
         """Shape of single, i.e. non-batched, return values of the covariance function.
@@ -468,45 +475,59 @@ def __rmul__(self, other: BinaryOperandType) -> Kernel:
 class IsotropicMixin(abc.ABC):  # pylint: disable=too-few-public-methods
     r"""Mixin for isotropic kernels.
 
-    An isotropic kernel is a kernel which only depends on the Euclidean norm of the
-    distance between the arguments, i.e.
+    An isotropic kernel is a kernel which only depends on the norm of the difference of
+    the arguments, i.e.
 
     .. math ::
 
-        k(x_0, x_1) = k(\lVert x_0 - x_1 \rVert_2).
+        k(x_0, x_1) = k(\lVert x_0 - x_1 \rVert).
 
     Hence, all isotropic kernels are stationary.
     """
 
     def _squared_euclidean_distances(
-        self, x0: np.ndarray, x1: Optional[np.ndarray]
+        self,
+        x0: np.ndarray,
+        x1: Optional[np.ndarray],
+        *,
+        scale_factors: Optional[np.ndarray] = None,
     ) -> np.ndarray:
-        """Implementation of the squared Euclidean distance, which supports scalar
-        inputs and an optional second argument."""
+        """Implementation of the squared (modified) Euclidean distance, which supports
+        scalar inputs, an optional second argument, and separate scale factors for each
+        input dimension."""
+
         if x1 is None:
             return np.zeros_like(  # pylint: disable=unexpected-keyword-arg
                 x0,
                 shape=x0.shape[: x0.ndim - self._input_ndim],
             )
 
-        sqdiffs = (x0 - x1) ** 2
+        sqdiffs = x0 - x1
 
-        if self.input_ndim == 0:
-            return sqdiffs
+        if scale_factors is not None:
+            sqdiffs *= scale_factors
 
-        assert self.input_ndim == 1
+        sqdiffs *= sqdiffs
 
-        return np.sum(sqdiffs, axis=-1)
+        return np.sum(sqdiffs, axis=tuple(range(-self.input_ndim, 0)))
 
     def _euclidean_distances(
-        self, x0: np.ndarray, x1: Optional[np.ndarray]
+        self,
+        x0: np.ndarray,
+        x1: Optional[np.ndarray],
+        *,
+        scale_factors: Optional[np.ndarray] = None,
     ) -> np.ndarray:
-        """Implementation of the Euclidean distance, which supports scalar inputs and an
-        optional second argument."""
+        """Implementation of the (modified) Euclidean distance, which supports scalar
+        inputs, an optional second argument, and separate scale factors for each input
+        dimension."""
+
         if x1 is None:
             return np.zeros_like(  # pylint: disable=unexpected-keyword-arg
                 x0,
                 shape=x0.shape[: x0.ndim - self._input_ndim],
             )
 
-        return np.sqrt(self._squared_euclidean_distances(x0, x1))
+        return np.sqrt(
+            self._squared_euclidean_distances(x0, x1, scale_factors=scale_factors)
+        )