Solved tests for Matern

Author: theogf

src/KernelFunctions.jl

@@ -34,8 +34,8 @@ export spectral_mixture_kernel, spectral_mixture_product_kernel
 using Compat
 using Requires
 using Distances, LinearAlgebra
-using SpecialFunctions: logabsgamma, besselk
-using ZygoteRules: @adjoint
+using SpecialFunctions: logabsgamma, besselk, polygamma
+using ZygoteRules: @adjoint, pullback
 using StatsFuns: logtwo
 using InteractiveUtils: subtypes
 using StatsBase

src/basekernels/matern.jl

@@ -17,12 +17,11 @@ end
 
 @inline function kappa(κ::MaternKernel, d::Real)
     ν = first(κ.ν)
-    iszero(d) ? one(d) :
-    exp(
-        (one(d) - ν) * logtwo - logabsgamma(ν)[1] +
-        ν * log(sqrt(2ν) * d) +
-        log(besselk(ν, sqrt(2ν) * d))
-    )
+    iszero(d) ? one(d) : _matern(ν, d)
+end
+
+function _matern(ν::Real, d::Real)
+    exp((one(d) - ν) * logtwo - loggamma(ν) + ν * log(sqrt(2ν) * d) + log(besselk(ν, sqrt(2ν) * d)))
 end
 
 metric(::MaternKernel) = Euclidean()

src/utils.jl

@@ -1,5 +1,7 @@
 hadamard(x, y) = x .* y
 
+loggamma(x) = first(logabsgamma(x))
+
 # Macro for checking arguments
 macro check_args(K, param, cond, desc=string(cond))
     quote
@@ -124,4 +126,3 @@ function validate_dims(x::AbstractVector, y::AbstractVector)
         ))
     end
 end
-

src/zygote_adjoints.jl

@@ -4,6 +4,20 @@
   end
 end
 
+@adjoint function loggamma(x)
+    first(logabsgamma(x)) , Δ -> (Δ .* polygamma(0, x), )
+end
+
+@adjoint function kappa(κ::MaternKernel, d::Real)
+    ν = first(κ.ν)
+    val, grad = pullback(_matern, ν, d)
+    return ((iszero(d) ? one(d) : val),
+    Δ -> begin
+        ∇ = grad(Δ)
+        return ((ν = [∇[1]],), iszero(d) ? zero(d) : ∇[2])
+    end)
+end
+
 @adjoint function ColVecs(X::AbstractMatrix)
     back(Δ::NamedTuple) = (Δ.X,)
     back(Δ::AbstractMatrix) = (Δ,)