Add i-times integrated Wiener Kernel (#77)

* Add WienerKernel

* Fix bugs

* Modify _wiener

* Add kernel matrix functions

* Add tests

* Modify _wiener to be separate for each i

* Style fix

* Move wiener kernel to basekernels

* Use @Assert

* Fix docstring

* Remove kappa and _kernel

* Apply suggestions from code review

Fix docstring and Base.show

Co-Authored-By: David Widmann <[email protected]>

* Remove _wiener

Co-Authored-By: David Widmann <[email protected]>

* Apply suggestions from code review

* Make docstring more readable

* Make docstring more readable

* Update Base.show

Co-Authored-By: David Widmann <[email protected]>

* Update docstring

* Update docstring

* Remove test for k0(m1, m2)

Co-Authored-By: David Widmann <[email protected]>

* Add more tests

* Fix tests

Co-Authored-By: David Widmann <[email protected]>

* Restrict line length

Co-authored-by: David Widmann <[email protected]>

Author: sharanry

src/KernelFunctions.jl

@@ -9,7 +9,7 @@ export transform
 export duplicate, set! # Helpers
 
 export Kernel
-export ConstantKernel, WhiteKernel, EyeKernel, ZeroKernel
+export ConstantKernel, WhiteKernel, EyeKernel, ZeroKernel, WienerKernel
 export CosineKernel
 export SqExponentialKernel, RBFKernel, GaussianKernel, SEKernel
 export LaplacianKernel, ExponentialKernel, GammaExponentialKernel

src/basekernels/wiener.jl

@@ -0,0 +1,78 @@
+"""
+    WienerKernel{i}()
+
+i-times integrated Wiener process kernel function.
+
+- For i=-1, this is just the white noise covariance, see [`WhiteKernel`](@ref).
+- For i= 0, this is the Wiener process covariance,
+- For i= 1, this is the integrated Wiener process covariance (velocity),
+- For i= 2, this is the twice-integrated Wiener process covariance (accel.),
+- For i= 3, this is the thrice-integrated Wiener process covariance,
+
+where ``κᵢ`` is given by
+
+```math
+    κ₋₁(x, y) =  δ(x, y)
+    κ₀(x, y)  =  min(x, y)
+```
+and for ``i >= 1``,
+```math
+    κᵢ(x, y) = 1 / aᵢ * min(x, y)^(2i + 1) + bᵢ * min(x, y)^(i + 1) * |x - y| * rᵢ(x, y),
+```
+    with the coefficients ``aᵢ``, ``bᵢ`` and the residual ``rᵢ(x, y)`` defined as follows:
+```math
+    a₁ = 3, b₁ = 1/2, r₁(x, y) = 1,
+    a₂ = 20, b₂ = 1/12, r₂(x, y) = x + y - min(x, y) / 2,
+    a₃ = 252, b₃ = 1/720, r₃(x, y) = 5 * max(x, y)² + 2 * x * z + 3 * min(x, y)²
+
+```
+
+# References:
+See the paper *Probabilistic ODE Solvers with Runge-Kutta Means* by Schober, Duvenaud and
+Hennig, NIPS, 2014, for more details.
+
+"""
+struct WienerKernel{I} <: BaseKernel
+    function WienerKernel{I}() where I
+        @assert I ∈ (-1, 0, 1, 2, 3) "Invalid parameter i=$(I). Should be -1, 0, 1, 2 or 3."
+        if I == -1
+            return WhiteKernel()
+        end
+        return new{I}()
+    end
+end
+
+function WienerKernel(;i::Integer=0)
+    return WienerKernel{i}()
+end
+
+function (::WienerKernel{0})(x, y)
+    X = sqrt(sum(abs2, x))
+    Y = sqrt(sum(abs2, y))
+    return min(X, Y)
+end
+
+function (::WienerKernel{1})(x, y)
+    X = sqrt(sum(abs2, x))
+    Y = sqrt(sum(abs2, y))
+    minXY = min(X, Y)
+    return 1 / 3 * minXY^3 + 1 / 2 * minXY^2 * euclidean(x, y)
+end
+
+function (::WienerKernel{2})(x, y)
+    X = sqrt(sum(abs2, x))
+    Y = sqrt(sum(abs2, y))
+    minXY = min(X, Y)
+    return 1 / 20 * minXY^5 + 1 / 12 * minXY^3 * euclidean(x, y) *
+        ( X + Y - 1 / 2 * minXY )
+end
+
+function (::WienerKernel{3})(x, y)
+    X = sqrt(sum(abs2, x))
+    Y = sqrt(sum(abs2, y))
+    minXY = min(X, Y)
+    return 1 / 252 * minXY^7 + 1 / 720 * minXY^4 * euclidean(x, y) *
+        ( 5 * max(X, Y)^2 + 2 * X * Y + 3 * minXY^2 )
+end
+
+Base.show(io::IO, ::WienerKernel{I}) where I = print(io, I, "-times integrated Wiener kernel")

test/basekernels/wiener.jl

@@ -0,0 +1,53 @@
+@testset "wiener" begin
+    k_1 = WienerKernel(i=-1)
+    @test typeof(k_1) <: WhiteKernel
+
+    k0 = WienerKernel()
+    @test typeof(k0) <: WienerKernel{0}
+
+    k1 = WienerKernel(i=1)
+    @test typeof(k1) <: WienerKernel{1}
+
+    k2 = WienerKernel(i=2)
+    @test typeof(k2) <: WienerKernel{2}
+
+    k3 = WienerKernel(i=3)
+    @test typeof(k3) <: WienerKernel{3}
+
+    @test_throws AssertionError WienerKernel(i=4)
+    @test_throws AssertionError WienerKernel(i=-2)
+
+    v1 = rand(4)
+    v2 = rand(4)
+
+    X = sqrt(sum(abs2, v1))
+    Y = sqrt(sum(abs2, v2))
+    minXY = min(X, Y)
+
+    @test k0(v1, v2) ≈ minXY
+    @test k1(v1, v2) ≈ 1 / 3 * minXY^3 + 1 / 2 * minXY^2 * euclidean(v1, v2)
+    @test k2(v1, v2) ≈ 1 / 20 * minXY^5 + 1 / 12 * minXY^3 * euclidean(v1, v2) *
+        ( X + Y - 1 / 2 * minXY )
+    @test k3(v1, v2) ≈ 1 / 252 * minXY^7 + 1 / 720 * minXY^4 * euclidean(v1, v2) *
+        ( 5 * max(X, Y)^2 + 2 * X * Y + 3 * minXY^2 )
+
+    # kernelmatrix tests
+    m1 = rand(3,4)
+    m2 = rand(3,4)
+    @test kernelmatrix(k0, m1, m1) ≈ kernelmatrix(k0, m1) atol=1e-5
+
+    K = zeros(4,4)
+    kernelmatrix!(K,k0,m1,m2)
+    @test K ≈ kernelmatrix(k0, m1, m2) atol=1e-5
+
+    V = zeros(4)
+    kerneldiagmatrix!(V,k0,m1)
+    @test V ≈ kerneldiagmatrix(k0,m1) atol=1e-5
+
+    x1 = rand()
+    x2 = rand()
+    @test kernelmatrix(k0, x1*ones(1,1), x2*ones(1,1))[1] ≈ k0(x1, x2) atol=1e-5
+    @test kernelmatrix(k1, x1*ones(1,1), x2*ones(1,1))[1] ≈ k1(x1, x2) atol=1e-5
+    @test kernelmatrix(k2, x1*ones(1,1), x2*ones(1,1))[1] ≈ k2(x1, x2) atol=1e-5
+    @test kernelmatrix(k3, x1*ones(1,1), x2*ones(1,1))[1] ≈ k3(x1, x2) atol=1e-5
+end

test/runtests.jl

@@ -76,6 +76,7 @@ using KernelFunctions: metric, kappa
         include(joinpath("basekernels", "polynomial.jl"))
         include(joinpath("basekernels", "piecewisepolynomial.jl"))
         include(joinpath("basekernels", "rationalquad.jl"))
+        include(joinpath("basekernels", "wiener.jl"))
     end
 
     @testset "kernels" begin