Merge pull request #121 from devmotion/maternad

Author: devmotion

src/KernelFunctions.jl

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

src/basekernels/matern.jl

@@ -16,12 +16,13 @@ struct MaternKernel{Tν<:Real} <: SimpleKernel
 end
 
 @inline function kappa(κ::MaternKernel, d::Real)
-    ν = first(κ.ν)
-    iszero(d) ? one(d) : _matern(ν, d)
+    result = _matern(first(κ.ν), d)
+    return ifelse(iszero(d), one(result), result)
 end
 
 function _matern(ν::Real, d::Real)
-    exp((one(d) - ν) * logtwo - loggamma(ν) + ν * log(sqrt(2ν) * d) + log(besselk(ν, sqrt(2ν) * d)))
+    y = sqrt(2ν) * d
+    return exp((one(d) - ν) * logtwo - loggamma(ν) + ν * log(y) + log(besselk(ν, y)))
 end
 
 metric(::MaternKernel) = Euclidean()

src/utils.jl

@@ -1,7 +1,5 @@
 hadamard(x, y) = x .* y
 
-loggamma(x) = first(logabsgamma(x))
-
 # Macro for checking arguments
 macro check_args(K, param, cond, desc=string(cond))
     quote

src/zygote_adjoints.jl

@@ -59,20 +59,6 @@ end
   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) = (Δ,)