# Source code for probnum.randprocs.covfuncs._matern

"""Matérn covariance function."""

import fractions
import functools
from typing import Optional, Tuple

import numpy as np
import scipy.special

from probnum.typing import ArrayLike, ScalarLike, ScalarType, ShapeLike
import probnum.utils as _utils

from ._covariance_function import CovarianceFunction, IsotropicMixin

_USE_KEOPS = True
try:
from pykeops.numpy import LazyTensor
except ImportError:  # pragma: no cover
_USE_KEOPS = False

class Matern(CovarianceFunction, IsotropicMixin):
r"""Matérn covariance function.

Covariance function defined by

.. math::
:nowrap:

k_\nu(x_0, x_1)
:=
\frac{2^{1 - \nu}}{\Gamma(\nu)}
\left( \sqrt{2 \nu} \lVert x_0 - x_1 \rVert_{\Lambda^{-1}} \right)^\nu
K_\nu \left( \sqrt{2 \nu} \lVert x_0 - x_1 \rVert_{\Lambda^{-1}} \right),

where :math:K_\nu is a modified Bessel function of the second kind and

.. math::
\lVert x_0 - x_1 \rVert_{\Lambda^{-1}}^2
:=
\sum_{i = 1}^d \frac{(x_{0,i} - x_{1,i})^2}{l_i}.

The Matérn covariance function generalizes the :class:~probnum.randprocs.covfuncs.\
ExpQuad covariance function via its additional parameter :math:\nu controlling
the smoothness of the functions in the associated RKHS.
For :math:\nu \rightarrow \infty, the Matérn covariance function converges to the
:class:~probnum.randprocs.covfuncs.ExpQuad covariance function.
A Gaussian process with Matérn covariance function is :math:\lceil \nu \rceil - 1
times differentiable.

If :math:\nu is a half-integer, i.e. :math:\nu = p + \frac{1}{2} for some
nonnegative integer :math:p, then the expression for the covariance function
simplifies to a product of an exponential and a polynomial

.. math::
:nowrap:

k_{\nu = p + \frac{1}{2}}(x_0, x_1)
=
\exp \left( -\sqrt{2 \nu} \lVert x_0 - x_1 \rVert_{\Lambda^{-1}} \right)
\frac{p!}{(2p)!}
\sum_{i = 0}^p \frac{(p + i)!}{i!(p - i)!} 2^{p - i}
\left( \sqrt{2 \nu} \lVert x_0 - x_1 \rVert_{\Lambda^{-1}} \right)^{p - i}.

Parameters
----------
input_shape
Shape of the covariance function's inputs.
nu
Hyperparameter :math:\nu controlling differentiability.
lengthscales
Lengthscales :math:l_i along the different input dimensions of the covariance
function.
Describes the input scales on which the process varies.
The lengthscales will be broadcast to the input shape of the covariance
function.

--------
ProductMatern : Tensor product of 1D Matérn covariance functions.

Examples
--------
>>> import numpy as np
>>> from probnum.randprocs.covfuncs import Matern
>>> K = Matern((), nu=2.5, lengthscales=0.1)
>>> xs = np.linspace(0, 1, 3)
>>> K.matrix(xs)
array([[1.00000000e+00, 7.50933789e-04, 3.69569622e-08],
[7.50933789e-04, 1.00000000e+00, 7.50933789e-04],
[3.69569622e-08, 7.50933789e-04, 1.00000000e+00]])
"""

def __init__(
self,
input_shape: ShapeLike,
nu: ScalarLike = 1.5,
*,
lengthscales: Optional[ArrayLike] = None,
):
super().__init__(input_shape=input_shape)

# Smoothness parameter ν
nu = _utils.as_numpy_scalar(nu)

if nu <= 0:
raise ValueError(f"Hyperparameter nu={nu} must be positive.")

self._nu = nu

# 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"All lengthscales l={lengthscales} must be positive.")

lengthscales, self.input_shape
)

self._lengthscales = lengthscales

@property
def nu(self) -> ScalarType:
r"""Smoothness parameter :math:\nu."""
return self._nu

@functools.cached_property
def p(self) -> Optional[int]:
r"""Degree :math:p of the polynomial part of a Matérn covariance function with
half-integer smoothness parameter :math:\nu = p + \frac{1}{2}. If :math:\nu
is not a half-integer, this is set to :data:None.

Sample paths of a Gaussian process with this covariance function are
:math:p-times continuously differentiable."""
nu_minus_half = self._nu - 0.5

# Half-integer values can be represented exactly in IEEE 754 floating point
# numbers
if nu_minus_half == int(nu_minus_half):
return int(nu_minus_half)

return None

@property
def is_half_integer(self) -> bool:
r"""Indicates whether :math:\nu is a half-integer."""
return self.p is not None

@property
def lengthscales(self) -> np.ndarray:
r"""Input lengthscales along the different input dimensions."""
return self._lengthscales

@property
def lengthscale(self) -> np.ndarray:  # pylint: disable=missing-raises-doc
"""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(2 * self._nu) / self._lengthscales

def _evaluate(self, x0: np.ndarray, x1: Optional[np.ndarray]) -> np.ndarray:
scaled_dists = self._euclidean_distances(
x0, x1, scale_factors=self._scale_factors
)

if self.is_half_integer:
# Evaluate the polynomial part using Horner's method
coeffs = Matern._half_integer_coefficients_floating(self.p)

res = coeffs[self.p]

for i in range(self.p - 1, -1, -1):
res *= scaled_dists
res += coeffs[i]

# Exponential part
res *= np.exp(-scaled_dists)

return res

return _matern_bessel(scaled_dists, nu=self._nu)

def _keops_lazy_tensor(
self, x0: np.ndarray, x1: Optional[np.ndarray]
) -> "LazyTensor":
if not _USE_KEOPS:  # pragma: no cover
raise ImportError()

scaled_dists = self._euclidean_distances_keops(
x0, x1, scale_factors=self._scale_factors
)

if self.is_half_integer:
# Evaluate the polynomial part using Horner's method
coeffs = Matern._half_integer_coefficients_floating(self.p)

res = coeffs[self.p]

for i in range(self.p - 1, -1, -1):
res *= scaled_dists
res += coeffs[i]

# Exponential part
res *= (-scaled_dists).exp()

return res

# TODO: Add KeOps implementation for non-half-integer case
raise NotImplementedError()

[docs] @staticmethod @functools.lru_cache(maxsize=None) def half_integer_coefficients(p: int) -> Tuple[fractions.Fraction]: r"""Computes the rational coefficients :math:c_i of the polynomial part of a Matérn covariance function with half-integer smoothness parameter :math:\nu = \ p + \frac{1}{2}. We leverage the recursion .. math:: :nowrap: \begin{align} c_{p - i} & := \frac{p!}{(2p)!} \frac{(p + i)!}{i!(p - i)!} 2^{p - i} \\ & = \frac{2 (i + 1)}{(p + i + 1) (p - i)} c_{p - (i + 1)}, \end{align} where :math:c_0 = c_{p - p} = 1. Parameters ---------- p: Degree :math:p of the polynomial part. Returns ------- coefficients: A tuple containing the exact rational coefficients of the polynomial part, where the entry at index :math:i contains the coefficient corresponding to the monomial with degree :math:i. """ coeffs = [fractions.Fraction(1, 1)] for i in range(p - 1, -1, -1): coeffs.append(coeffs[-1] * 2 * (i + 1) / (p + i + 1) / (p - i)) return tuple(coeffs)
@staticmethod @functools.lru_cache(maxsize=None) def _half_integer_coefficients_floating(p: int) -> np.ndarray: return np.asarray(Matern.half_integer_coefficients(p), dtype=np.float64) def _matern_bessel(scaled_dists: np.ndarray, nu: ScalarType) -> np.ndarray: # The modified Bessel function K_nu is not defined for z=0 scaled_dists = np.maximum(scaled_dists, np.finfo(scaled_dists.dtype).eps) return ( 2 ** (1.0 - nu) / scipy.special.gamma(nu) * scaled_dists**nu * scipy.special.kv(nu, scaled_dists) )