Added matern kernel

Author: theogf

dev/debugAD.jl

@@ -3,28 +3,33 @@ using Zygote, ForwardDiff, Tracker
 using Test
 
 dims = [10,5]
-
 A = rand(dims...)
 B = rand(dims...)
 K = [zeros(dims[1],dims[1]),zeros(dims[2],dims[2])]
-kernels = [SquaredExponentialKernel]
-l = 2.0
+l = 0.1
 vl = l*ones(dims[1])
 testfunction(k,A,B) = sum(kernelmatrix(k,A,B))
 testfunction(k,A) = sum(kernelmatrix(k,A))
-
+k = MaternKernel(vl)
+KernelFunctions.kappa(k,3)
 testfunction(SquaredExponentialKernel(vl),A)
+testfunction(MaternKernel(vl),A)
+@which kernelmatrix(MaternKernel(vl),A,B)
 #For debugging
 @info "Running Zygote gradients"
 Zygote.refresh()
 ## Zygote
 Zygote.gradient(x->testfunction(SquaredExponentialKernel(x),A),vl)
+Zygote.gradient(x->testfunction(MaternKernel(x),A),vl)
 Zygote.gradient(x->testfunction(SquaredExponentialKernel(x),A,B),vl)[1]
+Zygote.gradient(x->testfunction(MaternKernel(x),A,B),vl)[1]
 Zygote.gradient(x->testfunction(SquaredExponentialKernel(x),A,B),l)
+Zygote.gradient(x->testfunction(MaternKernel(x),A,B),l)
 Zygote.gradient(x->testfunction(SquaredExponentialKernel(x),A),l)
+Zygote.gradient(x->testfunction(MaternKernel(x),A),l)
 @info "Running Tracker gradients"
 ## Tracker
-Tracker.gradient(x->testfunction(SquaredExponentialKernel(vl),x,B),A)
+# Tracker.gradient(x->testfunction(SquaredExponentialKernel(vl),x,B),A)
 # Tracker.gradient(x->testfunction(SquaredExponentialKernel(l),x[:,:]),A)
 # # Tracker.gradient(x->testfunction(SquaredExponentialKernel(x),A,B),vl)
 # Tracker.gradient(x->testfunction(SquaredExponentialKernel(x),A),vl)
@@ -34,6 +39,10 @@ Tracker.gradient(x->testfunction(SquaredExponentialKernel(vl),x,B),A)
 @info "Running ForwardDiff gradients"
 ## ForwardDiff
 ForwardDiff.gradient(x->testfunction(SquaredExponentialKernel(x),A,B),vl) #✓
+ForwardDiff.gradient(x->testfunction(MaternKernel(x),A,B),vl) #✓
 ForwardDiff.gradient(x->testfunction(SquaredExponentialKernel(x),A),vl) #✓
+ForwardDiff.gradient(x->testfunction(MaternKernel(x),A),vl) #✓
 ForwardDiff.gradient(x->testfunction(SquaredExponentialKernel(x[1]),A,B),[l])
+ForwardDiff.gradient(x->testfunction(MaternKernel(x[1]),A,B),[l])
 ForwardDiff.gradient(x->testfunction(SquaredExponentialKernel(x[1]),A),[l])
+ForwardDiff.gradient(x->testfunction(MaternKernel(x[1]),A),[l])

src/KernelFunctions.jl

@@ -1,12 +1,13 @@
 module KernelFunctions
 
 export kernelmatrix, kernelmatrix!, kappa
-export Kernel, SquaredExponentialKernel
+export Kernel, SquaredExponentialKernel, MaternKernel, Matern3_2Kernel, Matern5_2Kernel
 
 export Transform, ScaleTransform
 
 using Distances, LinearAlgebra
 using Zygote: @adjoint
+using SpecialFunctions: lgamma, besselk
 
 const defaultobs = 2
 abstract type Kernel{T,Tr} end
@@ -16,7 +17,7 @@ include("common.jl")
 include("transform/transform.jl")
 include("kernelmatrix.jl")
 
-kernels = ["squaredexponential"]
+kernels = ["squaredexponential","matern"]
 for k in kernels
     include(joinpath("kernels",k*".jl"))
 end

src/kernelmatrix.jl

@@ -83,32 +83,29 @@ end
     kernelmatrix(κ::Kernel, X::Matrix ; obsdim::Int=2, symmetrize::Bool=true)
 ```
 Calculate the kernel matrix of `X` with respect to kernel `κ`.
+# USED
 """
 function kernelmatrix(
-        κ::Kernel{T,<:Transform{A}},
+        κ::Kernel{T,<:Transform},
         X::AbstractMatrix;
         obsdim::Int = defaultobs,
         symmetrize::Bool = true
     ) where {T,A}
     # Tₖ = typeof(zero(eltype(X))*zero(T))
     # m = size(X,obsdim)
-    K = map(x->kappa(κ,x),pairwise(metric(κ),transform(κ,X,obsdim),dims=obsdim))
-    # K = Matrix{Tₖ}(undef,m,m)
-    # for i in 1:m
-    #     tx = transform(κ,@view X[i,:])
-    #     for j in 1:i
-    #         K[i,j] = kappa(κ,kernel(κ,tx,transform(@view X[j,:])))
-    #     end
-    # end
+    #WARNING TEMP FIX
+    X̂ = transform(κ,X,obsdim)
+    K = map(x->kappa(κ,x),pairwise(metric(κ),X̂,X̂,dims=obsdim))
+    # K = map(x->kappa(κ,x),pairwise(metric(κ),transform(κ,X,obsdim),dims=obsdim))
     return K
-    # return kernelmatrix!(Matrix{Tₖ}(undef,m,m),κ,X,obsdim=obsdim,symmetrize=symmetrize)
 end
 
 """
 ```
     kernelmatrix(κ::Kernel, X::Matrix, Y::Matrix; obsdim::Int=2)
 ```
 Calculate the base matrix of `X` and `Y` with respect to kernel `κ`.
+# USED
 """
 function kernelmatrix(
         κ::Kernel{T},

src/kernels/matern.jl

@@ -0,0 +1,95 @@
+"""
+    MaternKernel([[ρ=1],ν=3/2])
+
+The matern kernel is an isotropic Mercer kernel given by the formula:
+
+```
+    κ(x,y) = 2^{1-ν}/Γ(ν)*(√(2ν)‖x-y‖/ρ)^ν K_ν(√(2ν)‖x-y‖/ρ)
+```
+
+For `ν=n+1/2, n=0,1,2,...` it can be simplified (it will be converted automatically).
+`ρ` is a lengthscale parameter.
+
+# Examples
+
+```jldoctest; setup = :(using KernelFunctions)
+julia> MaternKernel()
+Matern3_2Kernel{Float64,Float64}(1.0)
+
+julia> MaternKernel(2.0f0,3.0)
+MaternKernel{Float32,Float32}(2.0,3.0)
+
+julia> MaternKernel([2.0,3.0],5/2)
+Matern5_2Kernel{Float64,Array{Float64}}([2.0,3.0])
+```
+"""
+struct MaternKernel{T,Tr<:Transform} <: Kernel{T,Tr}
+    transform::Tr
+    metric::SemiMetric
+    ν::Real
+    function MaternKernel{T,Tr}(transform::Tr,ν::Real) where {T,Tr<:Transform}
+        return new{T,Tr}(transform,SqEuclidean(),ν)
+    end
+end
+
+function MaternKernel(ρ::T₁=1.0,ν::T₂=1.5) where {T₁<:Real,T₂<:Real}
+    @check_args(MaternKernel, ν, ν > zero(T₂), "ν > 0")
+    if ν == 0.5
+        ExponentialKernel{T₁,ScaleTransform{T₁}}(ScaleTransform(ρ))
+    elseif ν == 1.5
+        Matern3_2Kernel{T₁,ScaleTransform{T₁}}(ScaleTransform(ρ))
+    elseif ν == 2.5
+        Matern5_2Kernel{T₁,ScaleTransform{T₁}}(ScaleTransform(ρ))
+    else
+        MaternKernel{T₁,ScaleTransform{T₁}}(ScaleTransform(ρ),ν)
+    end
+end
+
+function MaternKernel(ρ::A,ν::T=1.5) where {A<:AbstractVector{<:Real},T<:Real}
+    @check_args(MaternKernel, ν, ν > zero(T), "ν > 0")
+    if ν == 0.5
+        ExponentialKernel{eltype(A),ScaleTransform{A}}(ScaleTransform(ρ))
+    elseif ν == 1.5
+        Matern3_2Kernel{eltype(A),ScaleTransform{A}}(ScaleTransform(ρ))
+    elseif ν == 2.5
+        Matern5_2Kernel{eltype(A),ScaleTransform{A}}(ScaleTransform(ρ))
+    else
+        MaternKernel{eltype(A),ScaleTransform{A}}(ScaleTransform(ρ),ν)
+    end
+end
+
+function MaternKernel(t::T₁,ν::T₂=1.5) where {T₁<:Transform,T₂<:Real}
+    @check_args(MaternKernel, ν, ν > zero(T₂), "ν > 0")
+    if ν == 0.5
+        ExponentialKernel{eltype(t),T₁}(ScaleTransform(ρ))
+    elseif ν == 1.5
+        Matern3_2Kernel{eltype(t),T₁}(ScaleTransform(ρ))
+    elseif ν == 2.5
+        Matern5_2Kernel{eltype(t),T₁}(ScaleTransform(ρ))
+    else
+        MaternKernel{eltype(t),T₁}(ScaleTransform(ρ),ν)
+    end
+end
+
+@inline kappa(κ::MaternKernel, d::Real) where {T} = exp((1.0-κ.ν)*log2 - lgamma(κ.ν) - κ.ν*log(sqrt(2κ.ν*d²)))*besselk(κ.ν,sqrt(2κ.ν*d²))
+
+
+struct Matern3_2Kernel{T,Tr<:Transform} <: Kernel{T,Tr}
+    transform::Tr
+    metric::SemiMetric
+    function Matern3_2Kernel{T,Tr}(transform::Tr) where {T,Tr<:Transform}
+        return new{T,Tr}(transform,SqEuclidean())
+    end
+end
+
+@inline kappa(κ::Matern3_2Kernel, d²::T) where {T<:Real} = (1+sqrt(3*d²))*exp(-sqrt(3*d²))
+
+struct Matern5_2Kernel{T,Tr<:Transform} <: Kernel{T,Tr}
+    transform::Tr
+    metric::SemiMetric
+    function Matern5_2Kernel{T,Tr}(transform::Tr) where {T,Tr<:Transform}
+        return new{T,Tr}(transform,SqEuclidean())
+    end
+end
+
+@inline kappa(κ::Matern5_2Kernel, d²::Real) where {T} = (1+sqrt(5*d²)+5*d²/3)*exp(-sqrt(5*d²))

test/kernelmatrix.jl

@@ -1,9 +1,12 @@
+using Distances
+
 dims = [10,5]
 
 A = rand(dims...)
 B = rand(dims...)
 K = [zeros(dims[1],dims[1]),zeros(dims[2],dims[2])]
 k = SquaredExponentialKernel()
+k = MaternKernel()
 
 @testset "Inplace Kernel Matrix" begin
     for obsdim in [1,2]