Add matern kernel and lean more on Distributions.jl (#16)

* Start on new abstraction

* Start rewrite test files for easier spot testing, add matern file

* Add matern kernel, rewrite things in terms of distributions

* Start implementing MvTDist

* Forgot main file

* Move dimension handing to base functions

* More steps towards MvTDist

* Use undocumented MvTDist

* Improve tests

* Fix tests for matern kernel

* Remove SpecialFunctions again

* Adjust squared exponential to new format

* Some cleanup

* Add compat entries

* Fix base tests

* Fix test script

* Small fix

* More test fixes

* Remove obsolte fallback

Author: Crown421

Project.toml

@@ -10,7 +10,10 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 
 [compat]
-julia = "1.6"
+julia = "1.8"
+Distributions = "0.25.108"
+KernelFunctions = "0.10.63"
+Reexport = "1.2.2"
 
 [extras]
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"

examples/approx-prior-sample.jl

@@ -0,0 +1 @@
+

examples/densities.jl

@@ -0,0 +1 @@
+

examples/feature-functions.jl

@@ -0,0 +1 @@
+

src/KernelSpectralDensities.jl

@@ -12,11 +12,9 @@ export SpectralDensity
 export ShiftedRFF, DoubleRFF
 export ApproximateGPSample
 
-# write tests to verify spectral density via fourier transforms
-# also add SpectralKernel (which can then be learned?)
-
 include("base.jl")
 include("expkernels.jl")
+include("matern.jl")
 include("features.jl")
 include("approx_prior.jl")
 

src/base.jl

@@ -1,8 +1,6 @@
 
 abstract type AbstractSpectralDensity end
 
-(S::AbstractSpectralDensity)(w) = error("Not implemented")
-
 """
     rand(S::AbstractSpectralDensity, [n::Int])
 
@@ -37,22 +35,63 @@ julia> k = SqExponentialKernel();
 julia> S = SpectralDensity(k, 1);
 
 julia> S(0.0)
-2.5066282746310002
+2.5066282746310007
 
 julia> S = SpectralDensity(k, 2);
 
 julia> S(zeros(2))
-6.283185307179585
+6.2831853071795845
 ```
 """
-struct SpectralDensity{K<:KernelFunctions.Kernel} <: AbstractSpectralDensity
+struct SpectralDensity{K<:KernelFunctions.Kernel,D<:Distribution} <: AbstractSpectralDensity
     kernel::K
-    dim::Int
+    # dim::Int
+    d::D
 
     function SpectralDensity(kernel::KernelFunctions.Kernel, dim::Int)
         if dim < 1
             throw(ArgumentError("Dimension must be greater than 0"))
         end
-        return new{typeof(kernel)}(kernel, dim)
+
+        sk, l = _deconstruct_kernel(kernel, dim)
+        d = _spectral_distribution(sk, l)
+
+        return new{typeof(kernel),typeof(d)}(kernel, d)
+    end
+end
+
+function (S::SpectralDensity)(w)
+    return pdf(S.d, w)
+end
+
+function rand(rng::AbstractRNG, S::SpectralDensity, n::Int...)
+    return rand(rng, S.d, n...)
+end
+
+# ToDo: This could perhaps go into a separate file
+function _deconstruct_kernel(ker::KernelFunctions.SimpleKernel, dim::Int)
+    if dim == 1
+        l = 1.0
+    else
+        l = ones(dim)
+    end
+    return ker, l
+end
+
+function _deconstruct_kernel(
+    ker::TransformedKernel{<:KernelFunctions.SimpleKernel,<:ScaleTransform}, dim::Int
+)
+    l = inv(only(ker.transform.s))
+    if dim > 1
+        l = ones(dim) * l
     end
+    return ker.kernel, l
+end
+
+function _deconstruct_kernel(ker::TransformedKernel, dim::Int)
+    return throw(MethodError(_deconstruct_kernel, (ker, dim)))
 end
+
+function _spectral_distribution(ker::KernelFunctions.Kernel, l)
+    return throw(MethodError(_spectral_distribution, (ker,)))
+end

src/expkernels.jl

@@ -3,40 +3,11 @@
 ## Squared ExponentialKernel
 
 # ToDo: Not sure about distances? Do all work?
-(S::SpectralDensity{<:SqExponentialKernel})(w) = _sqexp(w, 1)
-
-function (S::SpectralDensity{<:TransformedKernel{<:SqExponentialKernel,<:ScaleTransform}})(
-    w
-)
-    l = 1 / only(S.kernel.transform.s)
-    return _sqexp(w, l^2)
-end
-
-_sqexp(w::Real, l2::Real) = sqrt(2 * l2 * π) * exp(-2 * l2 * π^2 * w^2)
-function _sqexp(w::AbstractVector{<:Real}, l2::Real)
-    d = length(w)
-    return sqrt(2 * l2 * π)^d * exp(-2 * l2 * π^2 * dot(w, w))
-end
-
-function rand(rng::AbstractRNG, S::SpectralDensity{<:SqExponentialKernel}, n::Int...)
-    return _sqexprand(rng, S.dim, 1, n...)
-end
-
-function rand(
-    rng::AbstractRNG,
-    S::SpectralDensity{<:TransformedKernel{<:SqExponentialKernel,<:ScaleTransform}},
-    n::Int...,
-)
-    l = 1 / only(S.kernel.transform.s)
-    return _sqexprand(rng, S.dim, l, n...)
+function _spectral_distribution(ker::SqExponentialKernel, l::Real)
+    return inv(2 * π * l) * Normal()
 end
 
-function _sqexprand(rng::AbstractRNG, d::Int, l::Real, n::Int...)
-    σ = 1 / (2 * l * π)
-    if d == 1
-        return rand(rng, Normal(0, σ), n...)
-    elseif d > 1
-        σv = ones(d) * abs2(σ)
-        return rand(rng, MvNormal(Diagonal(σv)), n...)
-    end
+function _spectral_distribution(ker::SqExponentialKernel, l::AbstractVector)
+    σv = abs2.(inv.(2 * π * l))
+    return MvNormal(Diagonal(σv))
 end

src/matern.jl

@@ -0,0 +1,24 @@
+
+################################################### 
+## Matern Kernels
+
+MaternKernels = Union{MaternKernel,Matern32Kernel,Matern52Kernel}
+
+_matern_order(k::MaternKernel) = only(k.ν)
+_matern_order(::Matern32Kernel) = 3 / 2
+_matern_order(::Matern52Kernel) = 5 / 2
+
+# rewrite everything as returning a distribution (kind of as originally planned)
+# should be able to abstract/ generalize a lot of the special casing
+function _spectral_distribution(kernel::MaternKernels, l::Real)
+    ν = _matern_order(kernel)
+    return inv(2 * π * l) * TDist(2 * ν)
+end
+
+function _spectral_distribution(kernel::MaternKernels, l::AbstractVector)
+    ν = _matern_order(kernel)
+    n = length(l)
+    l = inv.(2 * π * l) .^ 2
+    D = Distributions.MvTDist(2 * ν, zeros(n), diagm(l))
+    return D
+end

test/base.jl

@@ -2,11 +2,13 @@ using KernelSpectralDensities
 using Test
 
 @testset "Fallback" begin
-    ker = ZeroKernel()
+    ker = ConstantKernel()
+
+    @test_throws MethodError SpectralDensity(ker, 1)
 
-    S = SpectralDensity(ker, 1)
+    kert = ker ∘ SelectTransform(1.0)
 
-    @test_throws ErrorException S(1.0)
+    @test_throws MethodError SpectralDensity(kert, 1)
 end
 
 @testset "Dimension check" begin

test/expkernels.jl

@@ -1,36 +1,48 @@
+if (!@isdefined RUN_TESTS) || !RUN_TESTS
+    using CairoMakie
+    show_plot = true
+    include("test_utils.jl")
+else
+    show_plot = false
+end
+
 @testset "SquaredExponential Kernel" begin
     ker = SqExponentialKernel()
     @testset "1D" begin
         @testset "Pure" begin
-            w_interval = [-2.0, 2.0]
+            # ker = SqExponentialKernel()
+            w_interval = 2.0
             t_interval = [0.0, 4.0]
 
-            f = test_spectral_density(ker, w_interval, t_interval)
+            test_spectral_density(ker, w_interval, t_interval; show_plot)
         end
 
         @testset "Scaled" begin
-            ker = with_lengthscale(SqExponentialKernel(), 0.7)
-            w_interval = [-2.0, 2.0]
+            # ker = SqExponentialKernel()
+            kert = with_lengthscale(ker, 0.7)
+            w_interval = 2.0
             t_interval = [0.0, 4.0]
 
-            f = test_spectral_density(ker, w_interval, t_interval)
+            f = test_spectral_density(kert, w_interval, t_interval; show_plot)
         end
     end
 
     @testset "2D" begin
         @testset "Pure" begin
+            # ker = SqExponentialKernel()
             w_interval = [-2.0, 2.0]
             x_interval = [-2.0, 2.0]
 
-            f = test_2Dspectral_density(ker, w_interval, x_interval)
+            f = test_2Dspectral_density(ker, w_interval, x_interval; show_plot)
         end
 
         @testset "Scaled" begin
-            ker = with_lengthscale(SqExponentialKernel(), 0.7)
+            # ker = SqExponentialKernel()
+            kert = with_lengthscale(ker, 0.7)
             w_interval = [-2.0, 2.0]
             x_interval = [-2.0, 2.0]
 
-            f = test_2Dspectral_density(ker, w_interval, x_interval)
+            f = test_2Dspectral_density(kert, w_interval, x_interval; show_plot)
         end
     end
 end