Use Functors

Author: devmotion

Project.toml

@@ -5,6 +5,7 @@ version = "0.6.0"
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
+Functors = "d9f16b24-f501-4c13-a1f2-28368ffc5196"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
@@ -16,6 +17,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 [compat]
 Compat = "2.2, 3"
 Distances = "0.9"
+Functors = "0.1"
 Requires = "1.0.1"
 SpecialFunctions = "0.8, 0.9, 0.10"
 StatsBase = "0.32, 0.33"

src/KernelFunctions.jl

@@ -37,6 +37,7 @@ export IndependentMOKernel
 using Compat
 using Requires
 using Distances, LinearAlgebra
+using Functors
 using SpecialFunctions: loggamma, besselk, polygamma
 using ZygoteRules: @adjoint, pullback
 using StatsFuns: logtwo
@@ -79,7 +80,6 @@ include("zygote_adjoints.jl")
 function __init__()
     @require Kronecker="2c470bb0-bcc8-11e8-3dad-c9649493f05e" include("matrix/kernelkroneckermat.jl")
     @require PDMats="90014a1f-27ba-587c-ab20-58faa44d9150" include("matrix/kernelpdmat.jl")
-    @require Flux="587475ba-b771-5e3f-ad9e-33799f191a9c" include("trainable.jl")
 end
 
 end

src/basekernels/constant.jl

@@ -15,7 +15,6 @@ metric(::ZeroKernel) = Delta()
 
 Base.show(io::IO, ::ZeroKernel) = print(io, "Zero Kernel")
 
-
 """
     WhiteKernel()
 
@@ -55,6 +54,8 @@ struct ConstantKernel{Tc<:Real} <: SimpleKernel
     end
 end
 
+@functor ConstantKernel
+
 kappa(κ::ConstantKernel,x::Real) = first(κ.c)*one(x)
 
 metric(::ConstantKernel) = Delta()

src/basekernels/exponential.jl

@@ -63,6 +63,8 @@ struct GammaExponentialKernel{Tγ<:Real} <: SimpleKernel
     end
 end
 
+@functor GammaExponentialKernel
+
 kappa(κ::GammaExponentialKernel, d²::Real) = exp(-d²^first(κ.γ))
 
 metric(::GammaExponentialKernel) = SqEuclidean()

src/basekernels/fbm.jl

@@ -17,6 +17,8 @@ struct FBMKernel{T<:Real} <: Kernel
     end
 end
 
+@functor FBMKernel
+
 function (κ::FBMKernel)(x::AbstractVector{<:Real}, y::AbstractVector{<:Real})
     modX = sum(abs2, x)
     modY = sum(abs2, y)

src/basekernels/gabor.jl

@@ -15,6 +15,8 @@ struct GaborKernel{K<:Kernel} <: Kernel
     end
 end
 
+@functor GaborKernel
+
 (κ::GaborKernel)(x, y) = κ.kernel(x ,y)
 
 function _gabor(; ell = nothing, p = nothing)

src/basekernels/maha.jl

@@ -16,6 +16,8 @@ struct MahalanobisKernel{T<:Real, A<:AbstractMatrix{T}} <: SimpleKernel
     end
 end
 
+@functor MahalanobisKernel
+
 kappa(κ::MahalanobisKernel, d::T) where {T<:Real} = exp(-d)
 
 metric(κ::MahalanobisKernel) = SqMahalanobis(κ.P)

src/basekernels/matern.jl

@@ -15,6 +15,8 @@ struct MaternKernel{Tν<:Real} <: SimpleKernel
     end
 end
 
+@functor MaternKernel
+
 @inline function kappa(κ::MaternKernel, d::Real)
     result = _matern(first(κ.ν), d)
     return ifelse(iszero(d), one(result), result)

src/basekernels/periodic.jl

@@ -20,6 +20,8 @@ PeriodicKernel(dims::Int) = PeriodicKernel(Float64, dims)
 
 PeriodicKernel(T::DataType, dims::Int = 1) = PeriodicKernel(r = ones(T, dims))
 
+@functor PeriodicKernel
+
 metric(κ::PeriodicKernel) = Sinus(κ.r)
 
 kappa(κ::PeriodicKernel, d::Real) = exp(- 0.5d)

src/basekernels/piecewisepolynomial.jl

@@ -25,6 +25,15 @@ function PiecewisePolynomialKernel(;v::Integer=0, maha::AbstractMatrix{<:Real})
     return PiecewisePolynomialKernel{v}(maha)
 end
 
+# Have to reconstruct the type parameter
+# See also https://github.com/FluxML/Functors.jl/issues/3#issuecomment-626747663
+function Functors.functor(::Type{<:PiecewisePolynomialKernel{V}}, x) where V
+    function reconstruct_kernel(xs)
+        return PiecewisePolynomialKernel{V}(xs.maha)
+    end
+    return (maha = x.maha,), reconstruct_kernel
+end
+
 _f(κ::PiecewisePolynomialKernel{0}, r, j) = 1
 _f(κ::PiecewisePolynomialKernel{1}, r, j) = 1 + (j + 1) * r
 _f(κ::PiecewisePolynomialKernel{2}, r, j) = 1 + (j + 2) * r + (j^2 + 4 * j + 3) / 3 * r.^2