Some improvement and additions to MO kernels (#354)

* Fix type stability of independent kernel

* Add convenience functions

* Add/fix tests

* Suggested change for kernelmatrix computation

* Add results to test script

* Switching to new kernelmatrix for independent kernel, tests

* Change type stability

* New kernelmatrix for IntrinsicCoregion kernel, tests

* JuliaFormatter

* Improve code reuse, Traits

* Use FillArray

* Remove traits, move to explicit function

* Relax matrixkernel types

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

* Adjust MOInput specifications

* Change supertype, move matrixkernel

* JuliaFormatter

* Remove forced typesymmetry, and improve fallback

* Improve documentation, add matrixkernel to other kernels

* Comment out specialized kernelmatrix!

* Add kernelmatrix! back in with version check

* Formatter

* Improve doc phrasing

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

* Change Union name

* Substantially improve tests

* Add ed clarity for lazy kronecker

* Improve tests for lmm and slfm

* Formatter

* Update src/mokernels/mokernel.jl

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

* Remove comments

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

* Add check for lmm kernel, tests

* Fix constructor mistake

* Remove kwarg dispatch

* Remove matrixkernel and lazy kron

* Remove matrixkernel export

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

* Update src/matrix/kernelkroneckermat.jl

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

* Update src/mokernels/independent.jl

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

* Update src/mokernels/independent.jl

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

* Update src/mokernels/mokernel.jl

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

* helper function name, formatter

* Remove explicit in-place tests

* Remove temp file and bump patch

* Change helper function name

* Fixed rename oversight

* Fix missing Union

* Silly mistake

* Make arguments more consistent

* Forgot to uncomment

Co-authored-by: st-- <[email protected]>
Co-authored-by: willtebbutt <[email protected]>

Author: 3 people

Project.toml

@@ -1,6 +1,6 @@
 name = "KernelFunctions"
 uuid = "ec8451be-7e33-11e9-00cf-bbf324bd1392"
-version = "0.10.15"
+version = "0.10.16"
 
 [deps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"

src/mokernels/independent.jl

@@ -24,17 +24,45 @@ struct IndependentMOKernel{Tkernel<:Kernel} <: MOKernel
 end
 
 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 _kernelmatrix_kron_helper(::MOInputIsotopicByFeatures, Kfeatures, B)
+    return kron(Kfeatures, B)
 end
 
-function kernelmatrix(k::IndependentMOKernel, x::MOInput, y::MOInput)
+function _kernelmatrix_kron_helper(::MOInputIsotopicByOutputs, Kfeatures, B)
+    return kron(B, Kfeatures)
+end
+
+function kernelmatrix(
+    k::IndependentMOKernel, x::MOI, y::MOI
+) where {MOI<:IsotopicMOInputsUnion}
     @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))
+    Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
+    mtype = eltype(Kfeatures)
+    return _kernelmatrix_kron_helper(x, Kfeatures, Eye{mtype}(x.out_dim))
+end
+
+if VERSION >= v"1.6"
+    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByFeatures, Kfeatures, B)
+        return kron!(K, Kfeatures, B)
+    end
+
+    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByOutputs, Kfeatures, B)
+        return kron!(K, B, Kfeatures)
+    end
+
+    function kernelmatrix!(
+        K::AbstractMatrix, k::IndependentMOKernel, x::MOI, y::MOI
+    ) where {MOI<:IsotopicMOInputsUnion}
+        @assert x.out_dim == y.out_dim
+        Ktmp = kernelmatrix(k.kernel, x.x, y.x)
+        mtype = eltype(Ktmp)
+        return _kernelmatrix_kron_helper!(
+            K, x, Ktmp, Matrix{mtype}(I, x.out_dim, x.out_dim)
+        )
+    end
 end
 
 function Base.show(io::IO, k::IndependentMOKernel)

src/mokernels/intrinsiccoregion.jl

@@ -34,10 +34,32 @@ 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 kernelmatrix(
+    k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
+) where {MOI<:IsotopicMOInputsUnion}
+    @assert x.out_dim == y.out_dim
+    Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
+    return _kernelmatrix_kron_helper(x, Kfeatures, k.B)
+end
+
+if VERSION >= v"1.6"
+    function kernelmatrix!(
+        K::AbstractMatrix, k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
+    ) where {MOI<:IsotopicMOInputsUnion}
+        @assert x.out_dim == y.out_dim
+        Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
+        return _kernelmatrix_kron_helper!(K, x, Kfeatures, k.B)
+    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
 
@@ -20,6 +21,10 @@ mixing matrix of ``m`` basis vectors spanning the output space.
 struct LinearMixingModelKernel{Tk<:AbstractVector{<:Kernel},Th<:AbstractMatrix} <: MOKernel
     K::Tk
     H::Th
+    function LinearMixingModelKernel(Tk::AbstractVector{<:Kernel}, H::AbstractMatrix)
+        @assert length(Tk) == size(H, 1) "Number of kernels and number of rows in H must match"
+        return new{typeof(Tk),typeof(H)}(Tk, H)
+    end
 end
 
 function LinearMixingModelKernel(k::Kernel, H::AbstractMatrix)

src/mokernels/moinput.jl

@@ -58,7 +58,7 @@ struct MOInputIsotopicByOutputs{S,T<:AbstractVector{S}} <: AbstractVector{Tuple{
     out_dim::Integer
 end
 
-const IsotopicMOInputs = Union{MOInputIsotopicByFeatures,MOInputIsotopicByOutputs}
+const IsotopicMOInputsUnion = Union{MOInputIsotopicByFeatures,MOInputIsotopicByOutputs}
 
 function Base.getindex(inp::MOInputIsotopicByOutputs, ind::Integer)
     @boundscheck checkbounds(inp, ind)
@@ -74,7 +74,7 @@ function Base.getindex(inp::MOInputIsotopicByFeatures, ind::Integer)
     return feature, output_index
 end
 
-Base.size(inp::IsotopicMOInputs) = (inp.out_dim * length(inp.x),)
+Base.size(inp::IsotopicMOInputsUnion) = (inp.out_dim * length(inp.x),)
 
 function Base.vcat(x::MOInputIsotopicByFeatures, y::MOInputIsotopicByFeatures)
     x.out_dim == y.out_dim || throw(DimensionMismatch("out_dim mismatch"))

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.
 

test/mokernels/independent.jl

@@ -1,20 +1,30 @@
 @testset "independent" begin
-    x = MOInput([rand(5) for _ in 1:4], 3)
-    y = MOInput([rand(5) for _ in 1:4], 3)
+    outdim = 3
+    x = KernelFunctions.MOInputIsotopicByOutputs([rand(5) for _ in 1:4], outdim)
+    y = KernelFunctions.MOInputIsotopicByOutputs([rand(5) for _ in 1:4], outdim)
+    z = KernelFunctions.MOInputIsotopicByOutputs([rand(5) for _ in 1:2], outdim)
+
+    xIF = KernelFunctions.MOInputIsotopicByFeatures(x.x, outdim)
+    yIF = KernelFunctions.MOInputIsotopicByFeatures(y.x, outdim)
+    zIF = KernelFunctions.MOInputIsotopicByFeatures(z.x, outdim)
 
     k = IndependentMOKernel(GaussianKernel())
     @test k isa IndependentMOKernel
     @test k isa MOKernel
     @test k isa Kernel
     @test k.kernel isa Kernel
-    @test k(x[2], y[2]) isa Real
 
     @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
-    @test k(x1[1], x1[1]) isa Real
-    @test kernelmatrix(k, x1) isa Matrix
+    ## accuracy
+    KernelFunctions.TestUtils.test_interface(k, x, y, z)
+    KernelFunctions.TestUtils.test_interface(k, xIF, yIF, zIF)
+
+    # type stability (maybe move to test_interface?)
+    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" *

test/mokernels/intrinsiccoregion.jl

@@ -2,20 +2,39 @@
     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)]
+    # XIF = [(rand(dims.in), rand(1:(dims.out))) for i in 1:(dims.obs)]
+    x = [rand(dims.in) for _ in 1:2]
+    XIF = KernelFunctions.MOInputIsotopicByFeatures(x, dims.out)
+    XIO = KernelFunctions.MOInputIsotopicByOutputs(x, dims.out)
+    y = [rand(dims.in) for _ in 1:2]
+    YIF = KernelFunctions.MOInputIsotopicByFeatures(y, dims.out)
+    YIO = KernelFunctions.MOInputIsotopicByOutputs(y, dims.out)
+    z = [rand(dims.in) for _ in 1:3]
+    ZIF = KernelFunctions.MOInputIsotopicByFeatures(z, dims.out)
+    ZIO = KernelFunctions.MOInputIsotopicByOutputs(z, dims.out)
 
     kernel = SqExponentialKernel()
-    icoregionkernel = IntrinsicCoregionMOKernel(; kernel=kernel, B=B)
+    icoregionkernel = IntrinsicCoregionMOKernel(kernel, B)
+
+    icoregionkernel2 = IntrinsicCoregionMOKernel(; kernel=kernel, B=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 icoregionkernel(XIF[1], XIF[1]) ≈
+          B[XIF[1][2], XIF[1][2]] * kernel(XIF[1][1], XIF[1][1])
+    @test icoregionkernel(XIF[1], XIF[end]) ≈
+          B[XIF[1][2], XIF[end][2]] * kernel(XIF[1][1], XIF[end][1])
+
+    # kernelmatrix
+    KernelFunctions.TestUtils.test_interface(icoregionkernel, XIF, YIF, ZIF)
+
+    KernelFunctions.TestUtils.test_interface(icoregionkernel, XIO, YIO, ZIO)
 
     KernelFunctions.TestUtils.test_interface(
         icoregionkernel, Vector{Tuple{Float64,Int}}; dim_out=dims.out

test/mokernels/lmm.jl

@@ -1,20 +1,30 @@
 @testset "lmm" begin
     rng = MersenneTwister(123)
     FDM = FiniteDifferences.central_fdm(5, 1)
-    N = 10
+    N = 6
     in_dim = 3
-    out_dim = 6
-    x1 = MOInput([rand(rng, in_dim) for _ in 1:N], out_dim)
-    x2 = MOInput([rand(rng, in_dim) for _ in 1:N], out_dim)
-    H = rand(4, 6)
-
-    k = LinearMixingModelKernel(
-        [Matern32Kernel(), SqExponentialKernel(), FBMKernel(), Matern32Kernel()], H
+    out_dim = 3
+    x1IO = KernelFunctions.MOInputIsotopicByOutputs(
+        [rand(rng, in_dim) for _ in 1:N], out_dim
+    )
+    x2IO = KernelFunctions.MOInputIsotopicByOutputs(
+        [rand(rng, in_dim) for _ in 1:N], out_dim
+    )
+    x3IO = KernelFunctions.MOInputIsotopicByOutputs(
+        [rand(rng, in_dim) for _ in 1:div(N, 2)], out_dim
     )
+
+    latentkernels = [Matern32Kernel(), SqExponentialKernel(), FBMKernel(), Matern32Kernel()]
+    H = rand(length(latentkernels), out_dim)
+    k = LinearMixingModelKernel(latentkernels, H)
+
+    badH = rand(length(latentkernels) - 1, out_dim)
+    @test_throws AssertionError LinearMixingModelKernel(latentkernels, badH)
+
     @test k isa LinearMixingModelKernel
     @test k isa MOKernel
     @test k isa Kernel
-    @test k(x1[1], x2[1]) isa Real
+    @test k(x1IO[1], x2IO[1]) isa Real
 
     @test string(k) == "Linear Mixing Model Multi-Output Kernel"
     @test repr("text/plain", k) == (
@@ -25,6 +35,14 @@
         "\tMatern 3/2 Kernel (metric = Euclidean(0.0))"
     )
 
+    TestUtils.test_interface(k, x1IO, x2IO, x3IO)
+
+    x1IF = KernelFunctions.MOInputIsotopicByFeatures(x1IO.x, out_dim)
+    x2IF = KernelFunctions.MOInputIsotopicByFeatures(x2IO.x, out_dim)
+    x3IF = KernelFunctions.MOInputIsotopicByFeatures(x3IO.x, out_dim)
+
+    TestUtils.test_interface(k, x1IF, x2IF, x3IF)
+
     k = LinearMixingModelKernel(SEKernel(), H)
 
     @test k isa LinearMixingModelKernel

test/mokernels/slfm.jl

@@ -1,11 +1,20 @@
 @testset "slfm" begin
     rng = MersenneTwister(123)
     FDM = FiniteDifferences.central_fdm(5, 1)
-    N = 10
-    in_dim = 5
+    N = 6
+    in_dim = 3
     out_dim = 4
-    x1 = MOInput([rand(rng, in_dim) for _ in 1:N], out_dim)
-    x2 = MOInput([rand(rng, in_dim) for _ in 1:N], out_dim)
+    x1IO = KernelFunctions.MOInputIsotopicByOutputs(
+        [rand(rng, in_dim) for _ in 1:N], out_dim
+    )
+    x2IO = KernelFunctions.MOInputIsotopicByOutputs(
+        [rand(rng, in_dim) for _ in 1:N], out_dim
+    )
+    x3IO = KernelFunctions.MOInputIsotopicByOutputs(
+        [rand(rng, in_dim) for _ in 1:div(N, 2)], out_dim
+    )
+    x1IO = MOInput([rand(rng, in_dim) for _ in 1:N], out_dim)
+    x2IO = MOInput([rand(rng, in_dim) for _ in 1:N], out_dim)
 
     k = LatentFactorMOKernel(
         [Matern32Kernel(), SqExponentialKernel(), FBMKernel()],
@@ -15,10 +24,17 @@
     @test k isa LatentFactorMOKernel
     @test k isa MOKernel
     @test k isa Kernel
-    @test k(x1[1], x2[1]) isa Real
+    @test k(x1IO[1], x2IO[1]) isa Real
+
+    @test kernelmatrix(k, x1IO, x2IO) ≈ kernelmatrix(k, collect(x1IO), collect(x2IO))
+
+    TestUtils.test_interface(k, x1IO, x2IO, x3IO)
+
+    x1IF = KernelFunctions.MOInputIsotopicByFeatures(x1IO.x, out_dim)
+    x2IF = KernelFunctions.MOInputIsotopicByFeatures(x2IO.x, out_dim)
+    x3IF = KernelFunctions.MOInputIsotopicByFeatures(x3IO.x, out_dim)
 
-    @test kernelmatrix(k, x1, x2) ≈ kernelmatrix(k, collect(x1), collect(x2))
-    @test kernelmatrix(k, x1, x1) ≈ kernelmatrix(k, x1)
+    TestUtils.test_interface(k, x1IF, x2IF, x3IF)
 
     @test string(k) == "Semi-parametric Latent Factor Multi-Output Kernel"
     @test repr("text/plain", k) == (
@@ -31,13 +47,13 @@
     )
 
     # AD test
-    function test_slfm(A::AbstractMatrix, x1, x2)
+    function test_slfm(A::AbstractMatrix, x1IO, x2IO)
         k = LatentFactorMOKernel(
             [Matern32Kernel(), SqExponentialKernel(), FBMKernel()],
             IndependentMOKernel(GaussianKernel()),
             A,
         )
-        return k((x1, 1), (x2, 1))
+        return k((x1IO, 1), (x2IO, 1))
     end
 
     k = LatentFactorMOKernel(