matern.jl

"""
    MaternKernel(; ν::Real=1.5, metric=Euclidean())

Matérn kernel of order `ν` with respect to the `metric`.

# Definition

For inputs ``x, x'`` and metric ``d(\\cdot, \\cdot)``, the Matérn kernel of order
``\\nu > 0`` is defined as
```math
k(x,x';\\nu) = \\frac{2^{1-\\nu}}{\\Gamma(\\nu)}\\big(\\sqrt{2\\nu} d(x, x')\\big) K_\\nu\\big(\\sqrt{2\\nu} d(x, x')\\big),
```
where ``\\Gamma`` is the Gamma function and ``K_{\\nu}`` is the modified Bessel function of
the second kind of order ``\\nu``.
By default, ``d`` is the Euclidean metric ``d(x, x') = \\|x - x'\\|_2``.

A Gaussian process with a Matérn kernel is ``\\lceil \\nu \\rceil - 1``-times
differentiable in the mean-square sense.

!!! note

    Differentiation with respect to the order ν is not currently supported.

See also: [`Matern12Kernel`](@ref), [`Matern32Kernel`](@ref), [`Matern52Kernel`](@ref)
"""
struct MaternKernel{Tν<:Real,M} <: SimpleKernel
    ν::Vector{Tν}
    metric::M

    function MaternKernel(ν::Real, metric; check_args::Bool=true)
        check_args && @check_args(MaternKernel, ν, ν > zero(ν), "ν > 0")
        return new{typeof(ν),typeof(metric)}([ν], metric)
    end
end

function MaternKernel(; nu::Real=1.5, ν::Real=nu, metric=Euclidean(), check_args::Bool=true)
    return MaternKernel(ν, metric; check_args)
end

@functor MaternKernel

@inline kappa(k::MaternKernel, d::Real) = _matern(only(k.ν), d)

function _matern(ν::Real, d::Real)
    if iszero(d)
        return one(d)
    else
        y = sqrt(2ν) * d
        b = log(besselk(ν, y))
        return exp((one(d) - ν) * oftype(y, logtwo) - loggamma(ν) + ν * log(y) + b)
    end
end

metric(k::MaternKernel) = k.metric

function Base.show(io::IO, κ::MaternKernel)
    return print(io, "Matern Kernel (ν = ", only(κ.ν), ", metric = ", κ.metric, ")")
end

## Matern12Kernel = ExponentialKernel aliased in exponential.jl

"""
    Matern32Kernel(; metric=Euclidean())

Matérn kernel of order ``3/2`` with respect to the `metric`.

# Definition

For inputs ``x, x'`` and metric ``d(\\cdot, \\cdot)``, the Matérn kernel of order ``3/2`` is
 given by
```math
k(x, x') = \\big(1 + \\sqrt{3} d(x, x') \\big) \\exp\\big(- \\sqrt{3} d(x, x') \\big).
```
By default, ``d`` is the Euclidean metric ``d(x, x') = \\|x - x'\\|_2``.

See also: [`MaternKernel`](@ref)
"""
struct Matern32Kernel{M} <: SimpleKernel
    metric::M
end

Matern32Kernel(; metric=Euclidean()) = Matern32Kernel(metric)

kappa(::Matern32Kernel, d::Real) = (1 + sqrt(3) * d) * exp(-sqrt(3) * d)

metric(k::Matern32Kernel) = k.metric

function Base.show(io::IO, k::Matern32Kernel)
    return print(io, "Matern 3/2 Kernel (metric = ", k.metric, ")")
end

"""
    Matern52Kernel(; metric=Euclidean())

Matérn kernel of order ``5/2`` with respect to the `metric`.

# Definition

For inputs ``x, x'`` and metric ``d(\\cdot, \\cdot)``, the Matérn kernel of order ``5/2`` is
given by
```math
k(x, x') = \\bigg(1 + \\sqrt{5} d(x, x') + \\frac{5}{3} d(x, x')^2\\bigg)
           \\exp\\big(- \\sqrt{5} d(x, x') \\big).
```
By default, ``d`` is the Euclidean metric ``d(x, x') = \\|x - x'\\|_2``.

See also: [`MaternKernel`](@ref)
"""
struct Matern52Kernel{M} <: SimpleKernel
    metric::M
end

Matern52Kernel(; metric=Euclidean()) = Matern52Kernel(metric)

kappa(::Matern52Kernel, d::Real) = (1 + sqrt(5) * d + 5 * d^2 / 3) * exp(-sqrt(5) * d)

metric(k::Matern52Kernel) = k.metric

function Base.show(io::IO, k::Matern52Kernel)
    return print(io, "Matern 5/2 Kernel (metric = ", k.metric, ")")
end