Matrixkernel convenience functions and related performance improvements

docs/src/api.md

@@ -80,6 +80,13 @@ type enables specialised implementations of e.g. [`kernelmatrix`](@ref) for
 
 To find out more about the background, read this [review of kernels for vector-valued functions](https://arxiv.org/pdf/1106.6251.pdf).
 
+If you are interested in the matrix-kernel interpretation, Kernelfunction provides a convenience function that computes the resulting kernel for a pair of inputs directly. 
+```@docs
+matrixkernel
+```
+<!-- Add when exporting IsotopicOutputs -->
+<!-- One way to look at this is that applying `matrixkernel` pairwise to a list of inputs results in a block matrix, which when flattened is the same is the `kernelmatrix` computed when providing the same list of inputs as `MOInputIsotopicByOutputs`.  -->
+
 ## Generic Utilities
 
 KernelFunctions also provides miscellaneous utility functions.

src/KernelFunctions.jl

@@ -39,6 +39,7 @@ export ColVecs, RowVecs
 export MOInput, prepare_isotopic_multi_output_data, prepare_heterotopic_multi_output_data
 export IndependentMOKernel,
     LatentFactorMOKernel, IntrinsicCoregionMOKernel, LinearMixingModelKernel
+export matrixkernel
 
 # Reexports
 export tensor, ⊗, compose
@@ -105,8 +106,8 @@ include(joinpath("kernels", "neuralkernelnetwork.jl"))
 include(joinpath("approximations", "nystrom.jl"))
 include("generic.jl")
 
-include(joinpath("mokernels", "mokernel.jl"))
 include(joinpath("mokernels", "moinput.jl"))
+include(joinpath("mokernels", "mokernel.jl"))
 include(joinpath("mokernels", "independent.jl"))
 include(joinpath("mokernels", "slfm.jl"))
 include(joinpath("mokernels", "intrinsiccoregion.jl"))

src/basekernels/fbm.jl

@@ -43,6 +43,8 @@ end
 _fbm(modX, modY, modXY, h) = (modX^h + modY^h - modXY^h) / 2
 
 _mod(x::AbstractVector{<:Real}) = abs2.(x)
+_mod(x::AbstractVector{<:AbstractVector{<:Real}}) = sum.(abs2, x)
+# two lines above could be combined into the second (dispatching on general AbstractVectors), but this (somewhat) more performant
 _mod(x::ColVecs) = vec(sum(abs2, x.X; dims=1))
 _mod(x::RowVecs) = vec(sum(abs2, x.X; dims=2))
 

src/mokernels/independent.jl

@@ -27,14 +27,6 @@ function (κ::IndependentMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,I
     return κ.kernel(x, y) * (px == py)
 end
 
-function _kernelmatrix_kron_helper(::MOInputIsotopicByFeatures, Kfeatures, B)
-    return kron(Kfeatures, B)
-end
-
-function _kernelmatrix_kron_helper(::MOInputIsotopicByOutputs, Kfeatures, B)
-    return kron(B, Kfeatures)
-end
-
 function kernelmatrix(
     k::IndependentMOKernel, x::MOI, y::MOI
 ) where {MOI<:IsotopicMOInputsUnion}
@@ -44,6 +36,8 @@ function kernelmatrix(
     return _kernelmatrix_kron_helper(x, Kfeatures, Eye{mtype}(x.out_dim))
 end
 
+matrixkernel(k::IndependentMOKernel, x, y) = k.kernel(x, y) * I
+
 if VERSION >= v"1.6"
     function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByFeatures, Kfeatures, B)
         return kron!(K, Kfeatures, B)

src/mokernels/intrinsiccoregion.jl

@@ -42,6 +42,8 @@ function (k::IntrinsicCoregionMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{
     return k.B[px, py] * k.kernel(x, y)
 end
 
+matrixkernel(k::IntrinsicCoregionMOKernel, x, y) = k.kernel(x, y) * k.B
+
 function kernelmatrix(
     k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
 ) where {MOI<:IsotopicMOInputsUnion}

src/mokernels/lmm.jl

@@ -37,6 +37,49 @@ function (κ::LinearMixingModelKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{A
     return sum(κ.H[i, px] * κ.K[i](x, y) * κ.H[i, py] for i in 1:length(κ.K))
 end
 
+## current implementation
+# julia> @benchmark KernelFunctions.kernelmatrix(k, x1IO, x2IO)
+# BenchmarkTools.Trial: 3478 samples with 1 evaluation.
+#  Range (min … max):  1.362 ms …   5.498 ms  ┊ GC (min … max): 0.00% … 72.47%
+#  Time  (median):     1.396 ms               ┊ GC (median):    0.00%
+#  Time  (mean ± σ):   1.435 ms ± 358.577 μs  ┊ GC (mean ± σ):  2.28% ±  6.70%
+
+#   ▂▆█▇▄▂  ▂▁                                                  ▁
+#   ███████▆██▅▅▁▄▁▁▁▁▁▁▁▁▄▁▁▁▁▁▁▁▁▁▃▁▁▁▁▁▁▁▁▁▁▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▃ █
+#   1.36 ms      Histogram: log(frequency) by time       2.1 ms <
+
+#  Memory estimate: 410.73 KiB, allocs estimate: 23090.
+
+## proposed improvement
+# julia> @benchmark KernelFunctions.kernelmatrix2(k, x1IO, x2IO)
+# BenchmarkTools.Trial: 10000 samples with 1 evaluation.
+#  Range (min … max):  16.871 μs …   3.440 ms  ┊ GC (min … max):  0.00% … 97.80%
+#  Time  (median):     18.625 μs               ┊ GC (median):     0.00%
+#  Time  (mean ± σ):   24.734 μs ± 129.308 μs  ┊ GC (mean ± σ):  20.63% ±  3.92%
+
+#    ▄▆███▇▆▅▄▄▂▂                           ▁ ▁                  ▂
+#   ████████████████▇▆▄▅▅▄▅▅▅▅▅▆▃▄▅▃▂▂▃▃▄▆▇▇████████▅▆▅▃▅▄▆▅▆▆▆▇ █
+#   16.9 μs       Histogram: log(frequency) by time      36.4 μs <
+
+#  Memory estimate: 84.56 KiB, allocs estimate: 338.
+
+function kernelmatrix2(k::LinearMixingModelKernel, X, Y)
+    K = [kernelmatrix(ki, X.x, Y.x) for ki in k.K]
+    L = size(k.H, 2)
+    return reduce(
+        hcat, [reduce(vcat, [sum(k.H[:, i] .* (K .* k.H[:, j])) for i in 1:L]) for j in 1:L]
+    )
+end
+
+# function matrixkernel(k::LinearMixingModelKernel, x, y)
+#     return matrixkernel(k, x, y, size(k.H, 2))
+# end
+
+function matrixkernel(k::LinearMixingModelKernel, x, y)
+    K = [ki(x, y) for ki in k.K]
+    return k.H' * (K .* k.H)
+end
+
 function Base.show(io::IO, k::LinearMixingModelKernel)
     return print(io, "Linear Mixing Model Multi-Output Kernel")
 end

src/mokernels/mokernel.jl

@@ -4,3 +4,31 @@
 Abstract type for kernels with multiple outpus.
 """
 abstract type MOKernel <: Kernel end
+
+"""
+    matrixkernel(k::MOK, x, y)
+
+Convenience function to compute the matrix kernel for two inputs `x` and `y`. The `outputsize` keyword is only required for the `IndependentMOKernel` to indicated the number of outputs. 
+"""
+function matrixkernel(k::MOKernel, x, y, out_dim)
+    @assert size(x) == size(y)
+    xMO = MOInputIsotopicByFeatures([x], out_dim)
+    yMO = MOInputIsotopicByFeatures([y], out_dim)
+    return kernelmatrix(k, xMO, yMO)
+end
+
+function matrixkernel(k::MOKernel, x, y)
+    return throw(
+        ArgumentError(
+            "This kernel does not have a specific matrixkernel implementation, you can call `matrixkernel(k, x, y, out_dim)`",
+        ),
+    )
+end
+
+function _kernelmatrix_kron_helper(::MOInputIsotopicByFeatures, Kfeatures, Koutputs)
+    return kron(Kfeatures, Koutputs)
+end
+
+function _kernelmatrix_kron_helper(::MOInputIsotopicByOutputs, Kfeatures, Koutputs)
+    return kron(Koutputs, Kfeatures)
+end

src/mokernels/slfm.jl

@@ -33,7 +33,20 @@ function (κ::LatentFactorMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,
     return cov_f + κ.e((x, px), (y, py))
 end
 
-function kernelmatrix(k::LatentFactorMOKernel, x::MOInput, y::MOInput)
+# function matrixkernel(k::LatentFactorMOKernel, x, y)
+#     return matrixkernel(k, x, y, size(k.A, 1))
+# end
+
+function matrixkernel(k::LatentFactorMOKernel, x, y)
+    W = [col * col' for col in eachcol(k.A)]
+    h = [gi(x, y) for gi in k.g]
+    w_h = sum(Wi * Hi for (Wi, Hi) in zip(W, h))
+    return w_h + matrixkernel(k.e, x, y)
+end
+
+function kernelmatrix(
+    k::LatentFactorMOKernel, x::IsotopicMOInputsUnion, y::IsotopicMOInputsUnion
+)
     x.out_dim == y.out_dim || error("`x` and `y` should have the same output dimension")
     x.out_dim == size(k.A, 1) ||
         error("Kernel not compatible with the given multi-output inputs")
@@ -45,7 +58,7 @@ function kernelmatrix(k::LatentFactorMOKernel, x::MOInput, y::MOInput)
     H = [gi.(x.x, permutedims(y.x)) for gi in k.g]
 
     # Weighted latent kernel matrix ((N*out_dim) x (N*out_dim))
-    W_H = sum(kron(Wi, Hi) for (Wi, Hi) in zip(W, H))
+    W_H = sum(_kernelmatrix_kron_helper(x, Hi, Wi) for (Wi, Hi) in zip(W, H))
 
     return W_H .+ kernelmatrix(k.e, x, y)
 end

test/basekernels/fbm.jl

@@ -15,6 +15,8 @@
     @test repr(k) == "Fractional Brownian Motion Kernel (h = $(h))"
     test_ADs(FBMKernel; ADs=[:ReverseDiff])
 
+    # ToDo: needs tests for _mod
+
     # Tests failing for ForwardDiff and Zygote@0.6.
     # Related to: https://github.com/FluxML/Zygote.jl/issues/1036
     f(x, y) = x^y

test/mokernels/intrinsiccoregion.jl

@@ -31,6 +31,9 @@
     @test icoregionkernel(XIF[1], XIF[end]) ≈
           B[XIF[1][2], XIF[end][2]] * kernel(XIF[1][1], XIF[end][1])
 
+    @test matrixkernel(icoregionkernel, XIF.x[1], XIF.x[2]) ≈
+          icoregionkernel.kernel(XIF.x[1], XIF.x[2]) * icoregionkernel.B
+
     # kernelmatrix
     KernelFunctions.TestUtils.test_interface(icoregionkernel, XIF, YIF, ZIF)
 

test/mokernels/lmm.jl

@@ -43,6 +43,10 @@
 
     TestUtils.test_interface(k, x1IF, x2IF, x3IF)
 
+    a = KernelFunctions.MOInputIsotopicByOutputs([rand(rng, in_dim)], out_dim)
+    b = KernelFunctions.MOInputIsotopicByOutputs([rand(rng, in_dim)], out_dim)
+    @test matrixkernel(k, a.x[1], b.x[1]) ≈ k.(a, permutedims(b))
+
     k = LinearMixingModelKernel(SEKernel(), H)
 
     @test k isa LinearMixingModelKernel

test/mokernels/mokernel.jl

@@ -0,0 +1,14 @@
+@testset "mokernel" begin
+    struct TestMOKernel <: MOKernel end
+    @test_throws ArgumentError matrixkernel(TestMOKernel(), rand(3), rand(3))
+
+    out_dim = 3
+    A = rand(out_dim, out_dim)
+    A = A * A'
+    k = IntrinsicCoregionMOKernel(GaussianKernel(), A)
+
+    in_dim = 4
+    x = rand(in_dim)
+    y = rand(in_dim)
+    @test matrixkernel(k, x, y, 3) ≈ matrixkernel(k, x, y)
+end

test/mokernels/slfm.jl

@@ -36,6 +36,11 @@
 
     TestUtils.test_interface(k, x1IF, x2IF, x3IF)
 
+    # matrixkernel
+    a = KernelFunctions.MOInputIsotopicByOutputs([rand(rng, in_dim)], out_dim)
+    b = KernelFunctions.MOInputIsotopicByOutputs([rand(rng, in_dim)], out_dim)
+    @test matrixkernel(k, a.x[1], b.x[1]) ≈ k.(a, permutedims(b))
+
     @test string(k) == "Semi-parametric Latent Factor Multi-Output Kernel"
     @test repr("text/plain", k) == (
         "Semi-parametric Latent Factor Multi-Output Kernel\n" *

test/runtests.jl

@@ -133,6 +133,7 @@ include("test_utils.jl")
     if GROUP == "" || GROUP == "MultiOutput"
         @testset "multi_output" begin
             include(joinpath("mokernels", "moinput.jl"))
+            include(joinpath("mokernels", "mokernel.jl"))
             include(joinpath("mokernels", "independent.jl"))
             include(joinpath("mokernels", "slfm.jl"))
             include(joinpath("mokernels", "intrinsiccoregion.jl"))