Implement ternary kerneldiagmatrix (#179)

* Implement ternary kerneldiagmatrix

* Bump patch version

Co-authored-by: wt <[email protected]>

Author: willtebbutt

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.8.2"
+version = "0.8.3"
 
 [deps]
 Compat = "34da2185-b29b-5c13-b0c7-acf172513d20"

src/matrix/kernelmatrix.jl

@@ -19,6 +19,7 @@ kernelmatrix
 
 """
     kerneldiagmatrix!(K::AbstractVector, κ::Kernel, X; obsdim::Int = 2)
+    kerneldiagmatrix!(K::AbstractVector, κ::Kernel, X, Y; obsdim::Int = 2)
 
 In place version of [`kerneldiagmatrix`](@ref)
 """
@@ -30,6 +31,11 @@ kerneldiagmatrix!
 Calculate the diagonal matrix of `X` with respect to kernel `κ`
 `obsdim = 1` means the matrix `X` has size #samples x #dimension
 `obsdim = 2` means the matrix `X` has size #dimension x #samples
+
+    kerneldiagmatrix(κ::Kernel, X, Y; obsdim::Int = 2)
+
+Calculate the diagonal of `kernelmatrix(κ, X, Y; obsdim)` efficiently. Requires that `X` and
+`Y` are the same length.
 """
 kerneldiagmatrix
 
@@ -59,8 +65,16 @@ function kerneldiagmatrix!(K::AbstractVector, κ::Kernel, x::AbstractVector)
     return map!(x -> κ(x, x), K, x)
 end
 
+function kerneldiagmatrix!(
+    K::AbstractVector, κ::Kernel, x::AbstractVector, y::AbstractVector,
+)
+    return map!(κ, x, y)
+end
+
 kerneldiagmatrix(κ::Kernel, x::AbstractVector) = map(x -> κ(x, x), x)
 
+kerneldiagmatrix(κ::Kernel, x::AbstractVector, y::AbstractVector) = map(κ, x, y)
+
 
 
 #
@@ -99,36 +113,47 @@ end
 const defaultobs = 2
 
 function kernelmatrix!(
-    K::AbstractMatrix, κ::Kernel, X::AbstractMatrix; obsdim::Int = defaultobs
+    K::AbstractMatrix, κ::Kernel, X::AbstractMatrix; obsdim::Int=defaultobs,
 )
     return kernelmatrix!(K, κ, vec_of_vecs(X; obsdim=obsdim))
 end
 
 function kernelmatrix!(
     K::AbstractMatrix, κ::Kernel, X::AbstractMatrix, Y::AbstractMatrix;
-    obsdim::Int = defaultobs
+    obsdim::Int=defaultobs,
 )
-    x = vec_of_vecs(X; obsdim=obsdim)
-    y = vec_of_vecs(Y; obsdim=obsdim)
-    return kernelmatrix!(K, κ, x, y)
+    return kernelmatrix!(K, κ, vec_of_vecs(X; obsdim=obsdim), vec_of_vecs(Y; obsdim=obsdim))
 end
 
-function kernelmatrix(κ::Kernel, X::AbstractMatrix; obsdim::Int = defaultobs)
+function kernelmatrix(κ::Kernel, X::AbstractMatrix; obsdim::Int=defaultobs)
     return kernelmatrix(κ, vec_of_vecs(X; obsdim=obsdim))
 end
 
 function kernelmatrix(κ::Kernel, X::AbstractMatrix, Y::AbstractMatrix; obsdim=defaultobs)
-    x = vec_of_vecs(X; obsdim=obsdim)
-    y = vec_of_vecs(Y; obsdim=obsdim)
-    return kernelmatrix(κ, x, y)
+    return kernelmatrix(κ, vec_of_vecs(X; obsdim=obsdim), vec_of_vecs(Y; obsdim=obsdim))
 end
 
 function kerneldiagmatrix!(
-    K::AbstractVector, κ::Kernel, X::AbstractMatrix; obsdim::Int = defaultobs
+    K::AbstractVector, κ::Kernel, X::AbstractMatrix; obsdim::Int=defaultobs
 )
     return kerneldiagmatrix!(K, κ, vec_of_vecs(X; obsdim=obsdim))
 end
 
-function kerneldiagmatrix(κ::Kernel, X::AbstractMatrix; obsdim::Int = defaultobs)
+function kerneldiagmatrix!(
+    K::AbstractVector, κ::Kernel, X::AbstractMatrix, Y::AbstractMatrix;
+    obsdim::Int = defaultobs,
+)
+    return kerneldiagmatrix!(
+        K, κ, vec_of_vecs(X; obsdim=obsdim), vec_of_vecs(Y; obsdim=obsdim),
+    )
+end
+
+function kerneldiagmatrix(κ::Kernel, X::AbstractMatrix; obsdim::Int=defaultobs)
     return kerneldiagmatrix(κ, vec_of_vecs(X; obsdim=obsdim))
 end
+
+function kerneldiagmatrix(
+    κ::Kernel, X::AbstractMatrix, Y::AbstractMatrix; obsdim::Int=defaultobs,
+)
+    return kerneldiagmatrix(κ, vec_of_vecs(X; obsdim=obsdim), vec_of_vecs(Y; obsdim=obsdim))
+end

src/test_utils.jl

@@ -34,25 +34,23 @@ function test_interface(
     x2::AbstractVector;
     atol=__ATOL,
 )
-    # TODO: uncomment the tests of ternary kerneldiagmatrix.
-
     # Ensure that we have the required inputs.
     @assert length(x0) == length(x1)
     @assert length(x0) ≠ length(x2)
 
     # Check that kerneldiagmatrix basically works.
-    # @test kerneldiagmatrix(k, x0, x1) isa AbstractVector
-    # @test length(kerneldiagmatrix(k, x0, x1)) == length(x0)
+    @test kerneldiagmatrix(k, x0, x1) isa AbstractVector
+    @test length(kerneldiagmatrix(k, x0, x1)) == length(x0)
 
     # Check that pairwise basically works.
     @test kernelmatrix(k, x0, x2) isa AbstractMatrix
     @test size(kernelmatrix(k, x0, x2)) == (length(x0), length(x2))
 
     # Check that elementwise is consistent with pairwise.
-    # @test kerneldiagmatrix(k, x0, x1) ≈ diag(kernelmatrix(k, x0, x1)) atol=atol
+    @test kerneldiagmatrix(k, x0, x1) ≈ diag(kernelmatrix(k, x0, x1)) atol=atol
 
     # Check additional binary elementwise properties for kernels.
-    # @test kerneldiagmatrix(k, x0, x1) ≈ kerneldiagmatrix(k, x1, x0)
+    @test kerneldiagmatrix(k, x0, x1) ≈ kerneldiagmatrix(k, x1, x0)
     @test kernelmatrix(k, x0, x2) ≈ kernelmatrix(k, x2, x0)' atol=atol
 
     # Check that unary elementwise basically works.
@@ -71,7 +69,7 @@ function test_interface(
     @test eigmin(Matrix(kernelmatrix(k, x0))) > -atol
 
     # Check that unary elementwise / pairwise are consistent with the binary versions.
-    # @test kerneldiagmatrix(k, x0) ≈ kerneldiagmatrix(k, x0, x0) atol=atol
+    @test kerneldiagmatrix(k, x0) ≈ kerneldiagmatrix(k, x0, x0) atol=atol
     @test kernelmatrix(k, x0) ≈ kernelmatrix(k, x0, x0) atol=atol
 
     # Check that basic kernel evaluation succeeds and is consistent with `kernelmatrix`.