Adding naive LMM model (#304)

Author: thomasgudjonwright

Project.toml

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

docs/src/kernels.md

@@ -134,4 +134,5 @@ MOKernel
 IndependentMOKernel
 LatentFactorMOKernel
 IntrinsicCoregionMOKernel
+LinearMixingModelKernel
 ```

src/KernelFunctions.jl

@@ -36,7 +36,8 @@ export spectral_mixture_kernel, spectral_mixture_product_kernel
 export ColVecs, RowVecs
 
 export MOInput
-export IndependentMOKernel, LatentFactorMOKernel, IntrinsicCoregionMOKernel
+export IndependentMOKernel,
+    LatentFactorMOKernel, IntrinsicCoregionMOKernel, LinearMixingModelKernel
 
 # Reexports
 export tensor, ⊗, compose
@@ -106,6 +107,7 @@ include(joinpath("mokernels", "moinput.jl"))
 include(joinpath("mokernels", "independent.jl"))
 include(joinpath("mokernels", "slfm.jl"))
 include(joinpath("mokernels", "intrinsiccoregion.jl"))
+include(joinpath("mokernels", "lmm.jl"))
 
 include("chainrules.jl")
 include("zygoterules.jl")

src/mokernels/lmm.jl

@@ -0,0 +1,45 @@
+@doc raw"""
+    LinearMixingModelKernel(g, e::MOKernel, A::AbstractMatrix)
+
+Kernel associated with the linear mixing model.
+
+# Definition
+
+For inputs ``x, x'`` and output dimensions ``p_x, p_{x'}'``, the kernel is defined as[^BPTHST]
+```math
+k\big((x, p_x), (x, p_{x'})\big) = H_{:,p_{x}}K(x, x')H_{:,p_{x'}}
+```
+where ``K(x, x') = Diag(k_1(x, x'), ..., k_m(x, x'))`` with zero off-diagonal entries.
+``H_{:,p_{x}}`` is the ``p_x``-th column (`p_x`-th output) of ``H \in \mathbb{R}^{m \times p}``
+representing ``m`` basis vectors for the ``p`` dimensional output space of ``f``.
+``k_1, \ldots, k_m`` are ``m`` kernels, one for each latent process, ``H`` is a
+mixing matrix of ``m`` basis vectors spanning the output space.
+
+[^BPTHST]: Wessel P. Bruinsma, Eric Perim, Will Tebbutt, J. Scott Hosking, Arno Solin, Richard E. Turner (2020). [Scalable Exact Inference in Multi-Output Gaussian Processes](https://arxiv.org/pdf/1911.06287.pdf).
+"""
+struct LinearMixingModelKernel{Tk<:AbstractVector{<:Kernel},Th<:AbstractMatrix} <: MOKernel
+    K::Tk
+    H::Th
+end
+
+function LinearMixingModelKernel(k::Kernel, H::AbstractMatrix)
+    return LinearMixingModelKernel(Fill(k, size(H, 1)), H)
+end
+
+function (κ::LinearMixingModelKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,Int})
+    (px > size(κ.H, 2) || py > size(κ.H, 2) || px < 1 || py < 1) &&
+        error("`px` and `py` must be within the range of the number of outputs")
+    return sum(κ.H[i, px] * κ.K[i](x, y) * κ.H[i, py] for i in 1:length(κ.K))
+end
+
+function Base.show(io::IO, k::LinearMixingModelKernel)
+    return print(io, "Linear Mixing Model Multi-Output Kernel")
+end
+
+function Base.show(io::IO, mime::MIME"text/plain", k::LinearMixingModelKernel)
+    print(io, "Linear Mixing Model Multi-Output Kernel. Kernels:")
+    for k in k.K
+        print(io, "\n\t")
+        show(io, mime, k)
+    end
+end

test/mokernels/lmm.jl

@@ -0,0 +1,48 @@
+@testset "lmm" begin
+    rng = MersenneTwister(123)
+    FDM = FiniteDifferences.central_fdm(5, 1)
+    N = 10
+    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
+    )
+    @test k isa LinearMixingModelKernel
+    @test k isa MOKernel
+    @test k isa Kernel
+    @test k(x1[1], x2[1]) isa Real
+
+    @test string(k) == "Linear Mixing Model Multi-Output Kernel"
+    @test repr("text/plain", k) == (
+        "Linear Mixing Model Multi-Output Kernel. Kernels:\n" *
+        "\tMatern 3/2 Kernel (metric = Euclidean(0.0))\n" *
+        "\tSquared Exponential Kernel (metric = Euclidean(0.0))\n" *
+        "\tFractional Brownian Motion Kernel (h = 0.5)\n" *
+        "\tMatern 3/2 Kernel (metric = Euclidean(0.0))"
+    )
+
+    k = LinearMixingModelKernel(SEKernel(), H)
+
+    @test k isa LinearMixingModelKernel
+    @test k isa MOKernel
+    @test k isa Kernel
+    @test length(k.K) == 4
+    for kernel in k.K
+        @test isa(kernel, SEKernel)
+    end
+
+    @test string(k) == "Linear Mixing Model Multi-Output Kernel"
+    @test repr("text/plain", k) == (
+        "Linear Mixing Model Multi-Output Kernel. Kernels:\n" *
+        "\tSquared Exponential Kernel (metric = Euclidean(0.0))\n" *
+        "\tSquared Exponential Kernel (metric = Euclidean(0.0))\n" *
+        "\tSquared Exponential Kernel (metric = Euclidean(0.0))\n" *
+        "\tSquared Exponential Kernel (metric = Euclidean(0.0))"
+    )
+
+    test_ADs(k)
+end

test/runtests.jl

@@ -140,6 +140,7 @@ include("test_utils.jl")
         include(joinpath("mokernels", "independent.jl"))
         include(joinpath("mokernels", "slfm.jl"))
         include(joinpath("mokernels", "intrinsiccoregion.jl"))
+        include(joinpath("mokernels", "lmm.jl"))
     end
     @info "Ran tests on Multi-Output Kernels"