Debugged the Matern Kernel 3/2

Author: theogf

Project.toml

@@ -5,6 +5,8 @@ version = "0.1.0"
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]

dev/debugAD.jl

@@ -1,14 +1,14 @@
 using KernelFunctions
 using Zygote, ForwardDiff, Tracker
-using Test
+using Test, LinearAlgebra
 
 dims = [10,5]
 A = rand(dims...)
 B = rand(dims...)
 K = [zeros(dims[1],dims[1]),zeros(dims[2],dims[2])]
 l = 0.1
 vl = l*ones(dims[1])
-testfunction(k,A,B) = sum(kernelmatrix(k,A,B))
+testfunction(k,A,B) = det(kernelmatrix(k,A,B))
 testfunction(k,A) = sum(kernelmatrix(k,A))
 k = MaternKernel(vl)
 KernelFunctions.kappa(k,3)
@@ -27,6 +27,8 @@ 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)
+Zygote.gradient(x->testfunction(MaternKernel(x),A),l)
+Zygote.gradient(x->kernelmatrix(MaternKernel(x,1.0),A)[1],l)
 @info "Running Tracker gradients"
 ## Tracker
 # Tracker.gradient(x->testfunction(SquaredExponentialKernel(vl),x,B),A)

src/KernelFunctions.jl

@@ -8,6 +8,7 @@ export Transform, ScaleTransform
 using Distances, LinearAlgebra
 using Zygote: @adjoint
 using SpecialFunctions: lgamma, besselk
+using StatsFuns: logtwo
 
 const defaultobs = 2
 abstract type Kernel{T,Tr} end

src/kernelmatrix.jl

@@ -94,9 +94,9 @@ function kernelmatrix(
     # Tₖ = typeof(zero(eltype(X))*zero(T))
     # m = size(X,obsdim)
     #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))
+    # 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
 end
 

src/kernels/matern.jl

@@ -28,7 +28,7 @@ struct MaternKernel{T,Tr<:Transform} <: Kernel{T,Tr}
     metric::SemiMetric
     ν::Real
     function MaternKernel{T,Tr}(transform::Tr,ν::Real) where {T,Tr<:Transform}
-        return new{T,Tr}(transform,SqEuclidean(),ν)
+        return new{T,Tr}(transform,Euclidean(),ν)
     end
 end
 
@@ -71,25 +71,25 @@ function MaternKernel(t::T₁,ν::T₂=1.5) where {T₁<:Transform,T₂<:Real}
     end
 end
 
-@inline kappa(κ::MaternKernel, d::Real) where {T} = exp((1.0-κ.ν)*log2 - lgamma(κ.ν) - κ.ν*log(sqrt(2κ.ν*d²)))*besselk(κ.ν,sqrt(2κ.ν*d²))
+@inline kappa(κ::MaternKernel, d::Real) where {T} = exp((1.0-κ.ν)*logtwo - 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())
+        return new{T,Tr}(transform,Euclidean())
     end
 end
 
-@inline kappa(κ::Matern3_2Kernel, d²::T) where {T<:Real} = (1+sqrt(3*d²))*exp(-sqrt(3*d²))
+@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())
+        return new{T,Tr}(transform,Euclidean())
     end
 end
 
-@inline kappa(κ::Matern5_2Kernel, d²::Real) where {T} = (1+sqrt(5*d²)+5*d²/3)*exp(-sqrt(5*d²))
+@inline kappa(κ::Matern5_2Kernel, d::Real) where {T} = (1+sqrt(5)*d+5*d^2/3)*exp(-sqrt(5)*d)

test/testAD.jl

@@ -1,6 +1,6 @@
 using KernelFunctions
 using Zygote, ForwardDiff, Tracker
-using Test
+using Test, LinearAlgebra
 
 dims = [10,5]
 
@@ -10,8 +10,8 @@ K = [zeros(dims[1],dims[1]),zeros(dims[2],dims[2])]
 kernels = [SquaredExponentialKernel,MaternKernel]
 l = 2.0
 vl = l*ones(dims[1])
-testfunction(k,A,B) = sum(kernelmatrix(k,A,B))
-testfunction(k,A) = sum(kernelmatrix(k,A))
+testfunction(k,A,B) = det(kernelmatrix(k,A,B))
+testfunction(k,A) = det(kernelmatrix(k,A))
 
 ##Eventually store real results in file
 @testset "Zygote Automatic Differentiation test" begin