Some improvement and additions to MO kernels

docs/src/api.md

@@ -74,6 +74,11 @@ 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.jl provides a convenience function that computes the resulting kernel when applied to a pair of inputs
+```@docs
+matrixkernel
+```
+
 ## 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
 export IndependentMOKernel,
     LatentFactorMOKernel, IntrinsicCoregionMOKernel, LinearMixingModelKernel
+export matrixkernel
 
 # Reexports
 export tensor, ⊗, compose

src/matrix/kernelkroneckermat.jl

@@ -25,3 +25,35 @@ end
     k(x,x') = ∏ᵢᴰ k(xᵢ,x'ᵢ)
 """
 @inline iskroncompatible(κ::Kernel) = false # Default return for kernels
+
+@inline ismatrixkroncompatible(κ::MOKernel) = false # Default return for kernels
+@inline ismatrixkroncompatible(κ::IndependentMOKernel) = true
+@inline ismatrixkroncompatible(κ::IntrinsicCoregionMOKernel) = true
+
+function _kroneckerkernelmatrix(Ktmp, B, ::MOInputIsotopicByFeatures)
+    return Kronecker.kronecker(Ktmp, B)
+end
+
+function _kroneckerkernelmatrix(K, Ktmp, B, ::MOInputIsotopicByOutputs)
+    return Kronecker.kronecker(B, Ktmp)
+end
+
+function kernelkronmat(k::IndependentMOKernel, x::MOI, y::MOI) where {MOI<:MOInputsUnion}
+    @assert x.out_dim == y.out_dim
+    Ktmp = kernelmatrix(k.kernel, x.x, y.x)
+    mtype = eltype(Ktmp)
+    return _kroneckerkernelmatrix(Ktmp, Eye{mtype}(x.out_dim), x)
+end
+
+function kernelkronmat(
+    k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
+) where {MOI<:MOInputsUnion}
+    @assert x.out_dim == y.out_dim
+    Ktmp = kernelmatrix(k.kernel, x.x, y.x)
+    return _kroneckerkernelmatrix(Ktmp, k.B, x)
+end
+
+function kernelkronmat(k::MOK, x::MOI) where {MOI<:MOInputsUnion,MOK<:MOKernel}
+    @assert iskroncompatible(κ) "The chosen kernel is not compatible for Kronecker matrices"
+    return kernelkronmat(k, x, x)
+end

src/mokernels/independent.jl

@@ -23,18 +23,44 @@ struct IndependentMOKernel{Tkernel<:Kernel} <: MOKernel
     kernel::Tkernel
 end
 
+# kernel function should be symmetric
+# would really like (κ::IndependentMOKernel)((x, px)::Tuple{T,Int}, (y, py)::Tuple{T,Int}) where T, but seems to cause autodiff problems
 function (κ::IndependentMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,Int})
-    if px == py
-        return κ.kernel(x, y)
-    else
-        return 0.0
-    end
+    return κ.kernel(x, y) * (px == py)
+end
+
+function _kronkernelmatrix(Ktmp, B, ::MOInputIsotopicByFeatures)
+    return kron(Ktmp, B)
 end
 
-function kernelmatrix(k::IndependentMOKernel, x::MOInput, y::MOInput)
+function _kronkernelmatrix(Ktmp, B, ::MOInputIsotopicByOutputs)
+    return kron(B, Ktmp)
+end
+
+function kernelmatrix(k::IndependentMOKernel, x::MOI, y::MOI) where {MOI<:MOInputsUnion}
     @assert x.out_dim == y.out_dim
-    temp = k.kernel.(x.x, permutedims(y.x))
-    return cat((temp for _ in 1:(y.out_dim))...; dims=(1, 2))
+    Ktmp = kernelmatrix(k.kernel, x.x, y.x)
+    mtype = eltype(Ktmp)
+    return _kronkernelmatrix(Ktmp, Eye{mtype}(x.out_dim), x)
+end
+
+if VERSION >= v"1.6"
+    function _kronkernelmatrix!(K, Ktmp, B, ::MOInputIsotopicByFeatures)
+        return kron!(K, Ktmp, B)
+    end
+
+    function _kronkernelmatrix!(K, Ktmp, B, ::MOInputIsotopicByOutputs)
+        return kron!(K, B, Ktmp)
+    end
+
+    function kernelmatrix!(
+        K::AbstractMatrix, k::IndependentMOKernel, x::MOI, y::MOI
+    ) where {MOI<:MOInputsUnion}
+        @assert x.out_dim == y.out_dim
+        Ktmp = kernelmatrix(k.kernel, x.x, y.x)
+        mtype = eltype(Ktmp)
+        return _kronkernelmatrix!(K, Ktmp, Matrix{mtype}(I, x.out_dim, x.out_dim), x)
+    end
 end
 
 function Base.show(io::IO, k::IndependentMOKernel)

src/mokernels/intrinsiccoregion.jl

@@ -34,10 +34,36 @@ function IntrinsicCoregionMOKernel(; kernel::Kernel, B::AbstractMatrix)
     return IntrinsicCoregionMOKernel{typeof(kernel),typeof(B)}(kernel, B)
 end
 
+function IntrinsicCoregionMOKernel(kernel::Kernel, B::AbstractMatrix)
+    return IntrinsicCoregionMOKernel{typeof(kernel),typeof(B)}(kernel, B)
+end
+
 function (k::IntrinsicCoregionMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,Int})
     return k.B[px, py] * k.kernel(x, y)
 end
 
+function matrixkernel(k::IntrinsicCoregionMOKernel, x, y)
+    return matrixkernel(k, x, y; outputsize=size(k.B, 1))
+end
+
+function kernelmatrix(
+    k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
+) where {MOI<:MOInputsUnion}
+    @assert x.out_dim == y.out_dim
+    Ktmp = kernelmatrix(k.kernel, x.x, y.x)
+    return _kronkernelmatrix(Ktmp, k.B, x)
+end
+
+if VERSION >= v"1.6"
+    function kernelmatrix!(
+        K::AbstractMatrix, k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
+    ) where {MOI<:MOInputsUnion}
+        @assert x.out_dim == y.out_dim
+        Ktmp = kernelmatrix(k.kernel, x.x, y.x)
+        return _kronkernelmatrix!(K, Ktmp, k.B, x)
+    end
+end
+
 function Base.show(io::IO, k::IntrinsicCoregionMOKernel)
     return print(
         io, "Intrinsic Coregion Kernel: ", k.kernel, " with ", size(k.B, 1), " outputs"

src/mokernels/lmm.jl

@@ -1,7 +1,8 @@
 @doc raw"""
-    LinearMixingModelKernel(g, e::MOKernel, A::AbstractMatrix)
+    LinearMixingModelKernel(k::Kernel, H::AbstractMatrix)
+    LinearMixingModelKernel(Tk::AbstractVector{<:Kernel},Th::AbstractMatrix)
 
-Kernel associated with the linear mixing model.
+Kernel associated with the linear mixing model, taking a vector of `m` kernels and a `m × p` matrix H for a function with `p` outputs. Also accepts a single kernel `k` for use across all `m` basis vectors. 
 
 # Definition
 
@@ -21,6 +22,7 @@ struct LinearMixingModelKernel{Tk<:AbstractVector{<:Kernel},Th<:AbstractMatrix}
     K::Tk
     H::Th
 end
+# does it maybe make sense to check that length(K) and size(H,1) are equal?
 
 function LinearMixingModelKernel(k::Kernel, H::AbstractMatrix)
     return LinearMixingModelKernel(Fill(k, size(H, 1)), H)
@@ -32,6 +34,10 @@ 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
 
+function matrixkernel(k::LinearMixingModelKernel, x, y)
+    return matrixkernel(k, x, y; outputsize=size(k.H, 2))
+end
+
 function Base.show(io::IO, k::LinearMixingModelKernel)
     return print(io, "Linear Mixing Model Multi-Output Kernel")
 end

src/mokernels/moinput.jl

@@ -194,3 +194,5 @@ function prepare_isotopic_multi_output_data(x::AbstractVector, y::RowVecs)
     length(x) == length(y) || throw(ArgumentError("length(x) not equal to length(y)."))
     return MOInputIsotopicByOutputs(x, size(y.X, 2)), vec(y.X)
 end
+
+MOInputsUnion = Union{MOInputIsotopicByFeatures,MOInputIsotopicByOutputs}

src/mokernels/mokernel.jl

@@ -4,3 +4,16 @@
 Abstract type for kernels with multiple outpus.
 """
 abstract type MOKernel <: Kernel end
+
+"""
+    matrixkernel(k::MOK, x, y)
+    matrixkernel(k::IndependentMOKernel, x, y(; outputsize))
+
+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::MOK, x, y; outputsize) where {MOK<:MOKernel}
+    @assert size(x) == size(y)
+    xMO = MOInputIsotopicByFeatures([x], outputsize)
+    yMO = MOInputIsotopicByFeatures([y], outputsize)
+    return kernelmatrix(k, xMO, yMO)
+end

src/mokernels/slfm.jl

@@ -1,5 +1,5 @@
 @doc raw"""
-    LatentFactorMOKernel(g, e::MOKernel, A::AbstractMatrix)
+    LatentFactorMOKernel(g::AbstractVector{<:Kernel}, e::MOKernel, A::AbstractMatrix)
 
 Kernel associated with the semiparametric latent factor model.
 
@@ -33,6 +33,10 @@ function (κ::LatentFactorMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,
     return cov_f + κ.e((x, px), (y, py))
 end
 
+function matrixkernel(k::LatentFactorMOKernel, x, y)
+    return matrixkernel(k, x, y; outputsize=size(k.A, 1))
+end
+
 function kernelmatrix(k::LatentFactorMOKernel, x::MOInput, y::MOInput)
     x.out_dim == y.out_dim || error("`x` and `y` should have the same output dimension")
     x.out_dim == size(k.A, 1) ||

temporary_script.jl

@@ -0,0 +1,78 @@
+# this script should be removed when the alternative MO kernelmatrix is accepted or rejected. 
+
+using KernelFunctions, BenchmarkTools
+using LinearAlgebra
+
+mrank = 1
+dims = (in=5, out=3)
+x = [rand(dims.in) for _ in 1:20]
+
+xMOF = KernelFunctions.MOInputIsotopicByFeatures(x, dims.out)
+xMOO = KernelFunctions.MOInputIsotopicByOutputs(x, dims.out)
+
+indk = IndependentMOKernel(GaussianKernel())
+
+Kind1 = kernelmatrix(indk, xMOF, xMOF)
+Kind2 = kernelmatrix2(indk, xMOF, xMOF)
+
+Kind1 ≈ Kind2
+# true
+
+@benchmark kernelmatrix($indk, $xMOF, $xMOF)
+# BenchmarkTools.Trial: 756 samples with 1 evaluation.
+#  Range (min … max):  6.186 ms …  11.470 ms  ┊ GC (min … max): 0.00% … 35.30%
+#  Time  (median):     6.423 ms               ┊ GC (median):    0.00%
+#  Time  (mean ± σ):   6.614 ms ± 806.792 μs  ┊ GC (mean ± σ):  2.76% ±  7.77%
+
+#    ▅▇█▇▃                                                       
+#   ▆██████▅▆▄▅▄▁▄▁▄▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▄▁▁▁▁▁▁▁▁▁▁▁▁▁▄▁▁▁▁▁▁▆██▇▇ ▇
+#   6.19 ms      Histogram: log(frequency) by time        10 ms <
+
+#  Memory estimate: 2.34 MiB, allocs estimate: 60060.
+
+@benchmark kernelmatrix2($indk, $xMOF, $xMOF)
+# BenchmarkTools.Trial: 10000 samples with 5 evaluations.
+#  Range (min … max):  6.162 μs … 341.226 μs  ┊ GC (min … max):  0.00% … 96.29%
+#  Time  (median):     6.947 μs               ┊ GC (median):     0.00%
+#  Time  (mean ± σ):   8.235 μs ±  16.802 μs  ┊ GC (mean ± σ):  12.37% ±  5.92%
+
+#    ▁▄▆▇██▇▆▅▄▃▂▁▁▁                                            ▂
+#   ▅██████████████████▇▇▆▇▅▆▆▆▅▅▆▅▅▃▄▅▅▅▁▁▃▁▃▅▅▆▆▆▇▇▇▇▇▇▇██▇▇▇ █
+#   6.16 μs      Histogram: log(frequency) by time      13.7 μs <
+
+#  Memory estimate: 34.89 KiB, allocs estimate: 7.
+
+A = randn(dims.out, mrank)
+B = A * transpose(A) + Diagonal(rand(dims.out))
+
+ickernel = IntrinsicCoregionMOKernel(GaussianKernel(), B)
+
+Kic1 = kernelmatrix(ickernel, xMOF, xMOF)
+Kic2 = kernelmatrix2(ickernel, xMOF, xMOF)
+
+Kic1 ≈ Kic2
+#true
+
+@benchmark kernelmatrix($ickernel, $xMOF, $xMOF)
+# BenchmarkTools.Trial: 1874 samples with 1 evaluation.
+#  Range (min … max):  2.522 ms …   5.424 ms  ┊ GC (min … max): 0.00% … 51.13%
+#  Time  (median):     2.601 ms               ┊ GC (median):    0.00%
+#  Time  (mean ± σ):   2.666 ms ± 369.030 μs  ┊ GC (mean ± σ):  2.01% ±  7.04%
+
+#   ▂█▆▁  ▁                                                      
+#   ███████▅▄▅▅▃▁▁▃▅▁▁▁▃▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▇▇ █
+#   2.52 ms      Histogram: log(frequency) by time       570 ms <
+
+#  Memory estimate: 985.47 KiB, allocs estimate: 39639.
+
+@benchmark kernelmatrix2($ickernel, $xMOF, $xMOF)
+# BenchmarkTools.Trial: 10000 samples with 5 evaluations.
+#  Range (min … max):  6.152 μs … 322.002 μs  ┊ GC (min … max):  0.00% … 96.14%
+#  Time  (median):     6.676 μs               ┊ GC (median):     0.00%
+#  Time  (mean ± σ):   8.100 μs ±  16.959 μs  ┊ GC (mean ± σ):  12.53% ±  5.85%
+
+#    ▄▇██▇▅▄▃▂▁▁                                        ▁▁▁▁    ▂
+#   ▆██████████████▇▇▇▆▆▇▇▇▇▆▇▆▄▅▄▃▆▁▅▄▅▆▆▅▅▄▅▄▆▇▇█▇▇▇▇██████▇▇ █
+#   6.15 μs      Histogram: log(frequency) by time      13.5 μs <
+
+#  Memory estimate: 34.77 KiB, allocs estimate: 6.

test/mokernels/independent.jl

@@ -1,6 +1,6 @@
 @testset "independent" begin
-    x = MOInput([rand(5) for _ in 1:4], 3)
-    y = MOInput([rand(5) for _ in 1:4], 3)
+    x = KernelFunctions.MOInputIsotopicByOutputs([rand(5) for _ in 1:4], 3)
+    y = KernelFunctions.MOInputIsotopicByOutputs([rand(5) for _ in 1:4], 3)
 
     k = IndependentMOKernel(GaussianKernel())
     @test k isa IndependentMOKernel
@@ -12,10 +12,32 @@
     @test kernelmatrix(k, x, y) == kernelmatrix(k, collect(x), collect(y))
     @test kernelmatrix(k, x, x) == kernelmatrix(k, x)
 
-    x1 = MOInput(rand(5), 3) # Single dim input
+    x1 = KernelFunctions.MOInputIsotopicByOutputs(rand(5), 3) # Single dim input
     @test k(x1[1], x1[1]) isa Real
     @test kernelmatrix(k, x1) isa Matrix
 
+    ## accuracy
+    @test kernelmatrix(k, x, y) ≈ k.(x, permutedims(y))
+
+    x_alt = KernelFunctions.MOInputIsotopicByFeatures(x.x, 3)
+    y_alt = KernelFunctions.MOInputIsotopicByFeatures(y.x, 3)
+    @test kernelmatrix(k, x_alt, y_alt) ≈ k.(x_alt, permutedims(y_alt))
+
+    # in-place
+    K = zeros(12, 12)
+    kernelmatrix!(K, k, x, y)
+    @test K ≈ k.(x, permutedims(y))
+
+    K = zeros(12, 12)
+    kernelmatrix!(K, k, x_alt, y_alt)
+    @test K ≈ k.(x_alt, permutedims(y_alt))
+
+    # type stability
+    x2 = MOInput(rand(Float32, 4), 2)
+    @test k(x2[1], x2[2]) isa Float32
+    @test k(x2[1], x2[1]) isa Float32
+    @test eltype(typeof(kernelmatrix(k, x2))) <: Float32
+
     @test string(k) ==
           "Independent Multi-Output Kernel\n" *
           "\tSquared Exponential Kernel (metric = Euclidean(0.0))"

test/mokernels/intrinsiccoregion.jl

@@ -2,25 +2,50 @@
     rng = MersenneTwister(123)
 
     dims = (in=3, out=2, obs=3)
-    rank = 1
+    r = 1
 
-    A = randn(dims.out, rank)
+    A = randn(dims.out, r)
     B = A * transpose(A) + Diagonal(rand(dims.out))
 
-    X = [(rand(dims.in), rand(1:(dims.out))) for i in 1:(dims.obs)]
+    # X = [(rand(dims.in), rand(1:(dims.out))) for i in 1:(dims.obs)]
+    x = [rand(dims.in) for _ in 1:2]
+    X = KernelFunctions.MOInputIsotopicByFeatures(x, dims.out)
 
     kernel = SqExponentialKernel()
     icoregionkernel = IntrinsicCoregionMOKernel(; kernel=kernel, B=B)
 
+    icoregionkernel2 = IntrinsicCoregionMOKernel(kernel, B)
+    @test icoregionkernel == icoregionkernel2
+
     @test icoregionkernel.B == B
     @test icoregionkernel.kernel == kernel
     @test icoregionkernel(X[1], X[1]) ≈ B[X[1][2], X[1][2]] * kernel(X[1][1], X[1][1])
     @test icoregionkernel(X[1], X[end]) ≈ B[X[1][2], X[end][2]] * kernel(X[1][1], X[end][1])
 
+    # test convenience function using kronecker product
+    @test matrixkernel(icoregionkernel, X.x[1], X.x[2]) ≈
+          icoregionkernel.kernel(X.x[1], X.x[2]) * icoregionkernel.B
+
+    # kernelmatrix
+    @test kernelmatrix(icoregionkernel, X) ≈ icoregionkernel.(X, permutedims(X))
+
+    X_alt = KernelFunctions.MOInputIsotopicByOutputs(x, dims.out)
+    @test kernelmatrix(icoregionkernel, X_alt) ≈ icoregionkernel.(X_alt, permutedims(X_alt))
+
     KernelFunctions.TestUtils.test_interface(
         icoregionkernel, Vector{Tuple{Float64,Int}}; dim_out=dims.out
     )
 
+    # in-place
+    kmsize = dims.out * length(x)
+    K = zeros(kmsize, kmsize)
+    kernelmatrix!(K, icoregionkernel, X, X)
+    @test K ≈ icoregionkernel.(X, permutedims(X))
+
+    K = zeros(kmsize, kmsize)
+    kernelmatrix!(K, icoregionkernel, X_alt, X_alt)
+    @test K ≈ icoregionkernel.(X_alt, permutedims(X_alt))
+
     test_ADs(icoregionkernel; dims=dims)
 
     @test string(icoregionkernel) ==