Implement Kernel Embeddings (#36)

* Change function names to be more clear about what they do

* Fixed tests

* Patch bump

Project.toml

@@ -1,7 +1,7 @@
 name = "BayesianQuadrature"
 uuid = "609f5bd8-aef1-42b2-b90e-083e3346dba9"
 authors = ["Theo Galy-Fajou <theo.galyfajou@gmail.com> and contributors"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 AbstractGPs = "99985d1d-32ba-4be9-9821-2ec096f28918"

src/BayesianQuadrature.jl

@@ -13,6 +13,8 @@ export PriorSampling
 export BayesModel
 export prior, integrand, logintegrand, logprior, logjoint
 export BQ # Short version for calling BayesianQuadrature
+export KernelEmbedding
+export kernel_mean, kernel_variance
 
 const BQ = BayesianQuadrature
 
@@ -42,7 +44,7 @@ abstract type AbstractBQModel{Tp,Ti} <: AbstractMCMC.AbstractModel end
 
 include("bayesquads/abstractbq.jl")
 include("samplers/abstractbqsampler.jl")
-include("kernelmeans/kernels.jl")
+include("kernelembeddings/kernelembedding.jl")
 include("interface.jl")
 include("models.jl")
 include("utils.jl")

src/bayesquads/abstractbq.jl

@@ -27,7 +27,7 @@ end
 
 function get_kernel_params(k::TransformedKernel; kwargs...)
     check_transform(k.transform)
-    return get_kernel_params(k.kernel; kwargs..., l=param(k.transform))
+    return get_kernel_params(k.kernel; kwargs..., l=transform_param(k.transform))
 end
 
 function check_transform(transform)
@@ -44,7 +44,3 @@ function kernel(b::AbstractBQ)
         return b.σ * (b.kernel ∘ ScaleTransform(inv.(b.l)))
     end
 end
-
-Λ(l::Real) = abs2(l) * I
-Λ(l::AbstractVector) = Diagonal(abs2.(l))
-Λ(l::LowerTriangular) = l * l'

src/bayesquads/bayesquad.jl

@@ -34,8 +34,9 @@ function quadrature(
     isempty(samples) && error("The collection of samples is empty")
     y = integrand(model).(samples)
     K = kernelpdmat(kernel(bquad), samples)
-    z = calc_z(samples, p_0(model), bquad)
-    C = calc_C(p_0(model), bquad)
+    ke = KernelEmbedding(bquad, p_0(model))
+    z = kernel_mean(ke, samples)
+    C = kernel_variance(ke)
     var = evaluate_var(z, K, C)
     if var < 0
         if var > -1e-5

src/bayesquads/logbayesquad.jl

@@ -41,22 +41,23 @@ function quadrature(
     f = exp.(logf) # Compute integrand on samples
 
     x_c = sample_candidates(bquad, samples, bquad.n_candidates) # Sample candidates around the samples
+    ke = KernelEmbedding(bquad, p_0(model))
 
     gp = create_gp(bquad, samples)
     f_c_0 = mean.(predict(gp, f, x_c)) # Predict integrand on x_c
     logf_c_0 = mean.(predict(gp, logf, x_c)) # Predict log-integrand on x_c
     Δ_c = exp.(logf_c_0) - f_c_0 # Compute difference of predictions
 
-    z = calc_z(samples, p_0(model), bquad) # Compute mean for the basic BQ
+    z = kernel_mean(ke, samples) # Compute mean for the basic BQ
     K = kernelpdmat(kernel(bquad), samples) # and the kernel matrix
 
-    z_c = calc_z(x_c, p_0(model), bquad) # Compute mean for the ΔlogBQ
+    z_c = kernel_mean(ke, x_c) # Compute mean for the ΔlogBQ
     K_c = kernelpdmat(kernel(bquad), x_c) # and the kernel matrix for the candidates
 
     m_evidence = evaluate_mean(z, K, f) # Compute m(Z|samples)
     m_correction = evaluate_mean(z_c, K_c, Δ_c) # 
 
-    C = calc_C(p_0(model), bquad) # Compute the C component for the variance
+    C = kernel_variance(ke) # Compute the kernel variance
 
     var_evidence = evaluate_var(z, K, C)
     var_correction = evaluate_var(z_c, K_c, C)

src/kernelembeddings/kernelembedding.jl

@@ -0,0 +1,62 @@
+abstract type AbstractKernelEmbedding end
+
+measure(ke::AbstractKernelEmbedding) = ke.measure
+
+struct KernelEmbedding{Tk<:Kernel,Tm,Tσ<:Real,Tl} <: AbstractKernelEmbedding
+    kernel::Tk # Kernel function
+    σ::Tσ # Kernel variance
+    l::Tl # Kernel lengthscale
+    measure::Tm # Measure
+end
+
+function KernelEmbedding(bquad::AbstractBQ, prior)
+    return KernelEmbedding(bquad.kernel, bquad.σ, bquad.l, prior)
+end
+
+scale(ke::KernelEmbedding) = ke.σ
+
+@doc raw"""
+    kernel_mean(ke::KernelEmbedding{Kernel,Measure}, samples::AbstractVector)
+
+Compute the kernel mean of the kernel embedding `ke` for each one of 
+the `samples` $$x_i$$:
+```math
+    z_i = \int k(x, x_i)d\mu(x)
+```
+"""
+kernel_mean
+
+function kernel_mean(ke::KernelEmbedding{<:SqExponentialKernel,<:AbstractMvNormal}, samples::AbstractVector)
+    z = samples .- Ref(mean(measure(ke)))
+    B = Λ(ke.l)
+    return scale(ke) / sqrt(det(inv(B) * cov(measure(ke)) + I)) *
+           exp.(- PDMats.invquad.(Ref(PDMat(B + cov(measure(ke)))), z) / 2)
+end
+
+@doc raw"""
+    kernel_variance(ke::KernelEmbedding{Kernel,Measure})
+
+Compute the kernel variance of the given kernel embedding:
+```math
+    C = \int\int k(x,x')d\mu(x)d\mu(x')
+```
+"""
+kernel_variance
+
+
+function kernel_variance(ke::KernelEmbedding{<:SqExponentialKernel,<:AbstractMvNormal})
+    B = Λ(ke.l)
+    return scale(ke) / sqrt(det(2 * inv(B) * cov(measure(ke)) + I))
+end
+
+
+# Turn the lengthscale into a Diagonal matrix of noise
+Λ(l::Real) = abs2(l) * I
+Λ(l::AbstractVector) = Diagonal(abs2.(l))
+Λ(l::LowerTriangular) = Cholesky(l, :L, 1)
+Λ(l::AbstractMatrix) = l * l'
+
+# Turns a transform into a lengthscale
+transform_param(t::ScaleTransform) = inv(first(t.s))
+transform_param(t::ARDTransform) = inv.(t.v)
+transform_param(t::LinearTransform) = inv(t.A)

test/bayesquads/bayesquad.jl

@@ -1,11 +1,5 @@
 @testset "bayesquad" begin
-    s = 2.0
-    l = [1.0, 2.0]
-    L = LowerTriangular(rand(2, 2))
     k = SqExponentialKernel()
-    @test BQ.Λ(s) ≈ s^2 * I
-    @test BQ.Λ(l) ≈ Diagonal(abs2.(l))
-    @test BQ.Λ(L) ≈ L * L'
 
     σ = 4.0
     @test BQ.scale(BayesQuad(σ * k)) == σ

test/kernelembeddings/kernelembedding.jl

@@ -0,0 +1,22 @@
+@testset "kernelembeddings" begin 
+    rng = MersenneTwister(42)
+    N = 3
+    measure = MvNormal(ones(3), ones(3))
+    k = SqExponentialKernel()
+    l = 2.0
+    σ = 0.5
+    ke = KernelEmbedding(k, σ, l, measure)
+    kernel = σ * with_lengthscale(k, l)
+    @test BQ.scale(ke) == σ
+
+    sample = [rand(rng, 3)]
+    @test kernel_mean(ke, sample) ≈ [mean(kernel.(sample, eachcol(rand(rng, measure, 10000))))] atol=1e-2
+    @test kernel_variance(ke) ≈ mean(kernel.(eachcol(rand(rng, measure, 10000)), eachcol(rand(rng, measure, 10000)))) atol=1e-2
+
+    s = 2.0
+    l = [1.0, 2.0]
+    L = LowerTriangular(rand(2, 2))
+    @test BQ.Λ(s) ≈ s^2 * I
+    @test BQ.Λ(l) ≈ Diagonal(abs2.(l))
+    @test Matrix(BQ.Λ(L)) ≈ L * L'
+end

test/runtests.jl

@@ -25,9 +25,8 @@ include("testing_tools.jl")
         include(joinpath("samplers", "priorsampling.jl"))
     end
 
-    @info "Testing kernel means"
-    @testset "Kernel Means" begin
-        include(joinpath("kernelmeans", "kernels.jl"))
-        include(joinpath("kernelmeans", "sekernel.jl"))
+    @info "Testing kernel embeddings"
+    @testset "Kernel Embeddings" begin
+        include(joinpath("kernelembeddings", "kernelembedding.jl"))
     end
 end