Add datatype for multi-output GP input and add independent multi-output kernel (#138)

* Add datatype for multi-output GP input

* Add IndependentKernel

* Address code review

* Address code review

* Address code review

* Update Tests

* Redefine Indendent kernel and add kernelmatrix function

* Use block diagonal outputs and fix kernel and kix style issues

* Remove exports for base functions.

* Remove MOKernel type and make MOInput a subtype of AbstractVector

* Make Delta metric return promoted type of the input arrays

* Promote and convert in single step

* Add more tests

Author: sharanry

src/KernelFunctions.jl

@@ -30,6 +30,8 @@ export NystromFact, nystrom
 
 export spectral_mixture_kernel, spectral_mixture_product_kernel
 
+export MOInput
+export IndependentMOKernel
 
 using Compat
 using Requires
@@ -69,6 +71,9 @@ include("kernels/tensorproduct.jl")
 include("approximations/nystrom.jl")
 include("generic.jl")
 
+include("mokernels/moinput.jl")
+include("mokernels/independent.jl")
+
 include("zygote_adjoints.jl")
 
 function __init__()

src/distances/delta.jl

@@ -1,11 +1,11 @@
 struct Delta <: Distances.PreMetric
 end
 
-@inline function Distances._evaluate(::Delta, a::AbstractVector, b::AbstractVector) where {T}
+@inline function Distances._evaluate(::Delta, a::AbstractVector{Ta}, b::AbstractVector{Tb}) where {Ta, Tb}
     @boundscheck if length(a) != length(b)
         throw(DimensionMismatch("first array has length $(length(a)) which does not match the length of the second, $(length(b))."))
     end
-    return a == b
+    return convert(promote_type(Ta, Tb), a == b)
 end
 
 Distances.result_type(::Delta, Ta::Type, Tb::Type) = promote_type(Ta, Tb)

src/mokernels/independent.jl

@@ -0,0 +1,26 @@
+"""
+    IndependentMOKernel(k::Kernel) <: Kernel
+
+A Multi-Output kernel which assumes each output is independent of the other.
+"""
+struct IndependentMOKernel{Tkernel<:Kernel} <: Kernel
+    kernel::Tkernel
+end
+
+function (κ::IndependentMOKernel)(x::Tuple{Vector, Int}, y::Tuple{Vector, Int})
+    if last(x) == last(y)
+        return κ.kernel(first(x), first(y))
+    else
+        return 0.0
+    end
+end
+
+function kernelmatrix(k::IndependentMOKernel, x::MOInput, y::MOInput)
+    @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))
+end
+
+function Base.show(io::IO, k::IndependentMOKernel)
+    print(io, string("Independent Multi-Output Kernel\n\t", string(k.kernel)))
+end

src/mokernels/moinput.jl

@@ -0,0 +1,31 @@
+"""
+    MOInput(x::AbstractVector, out_dim::Integer)
+
+A data type to accomodate modelling multi-dimensional output data.
+"""
+struct MOInput{T<:AbstractVector} <: AbstractVector{Tuple{Any,Int}}
+    x::T
+    out_dim::Integer
+end
+
+Base.length(inp::MOInput) = inp.out_dim * length(inp.x)
+
+Base.size(inp::MOInput, d) = d::Integer == 1 ? inp.out_dim * size(inp.x, 1) : 1 
+Base.size(inp::MOInput) = (inp.out_dim * size(inp.x, 1),)
+
+Base.lastindex(inp::MOInput) = length(inp)
+Base.firstindex(inp::MOInput) = 1
+
+function Base.getindex(inp::MOInput, ind::Integer)
+    if ind > 0
+        out_dim = ind ÷ length(inp.x) + 1
+        ind = ind % length(inp.x)
+        if ind==0 ind = length(inp.x); out_dim-=1 end
+        return (inp.x[ind], out_dim::Int)
+    else
+        throw(BoundsError(string("Trying to access at ", ind)))
+    end
+end
+
+Base.iterate(inp::MOInput) = (inp[1], 1)
+Base.iterate(inp::MOInput, state) = (state<length(inp)) ? (inp[state + 1], state + 1) : nothing

src/utils.jl

@@ -21,6 +21,7 @@ function vec_of_vecs(X::AbstractMatrix; obsdim::Int = 2)
 end
 
 dim(x::AbstractVector{<:Real}) = 1
+dim(x::AbstractVector{Tuple{Any,Int}}) = 1
 
 """
     ColVecs(X::AbstractMatrix)

test/matrix/kernelmatrix.jl

@@ -130,4 +130,13 @@ KernelFunctions.kappa(::ToySimpleKernel, d) = exp(-d / 2)
             @test kerneldiagmatrix(k, x) ≈ kerneldiagmatrix!(tmp_diag, k, X; obsdim=obsdim)
         end
     end
+
+    @testset "Multi Output Kernels" begin
+        x = MOInput([rand(5) for _ in 1:4], 3)
+        y = MOInput([rand(5) for _ in 1:4], 3)
+
+        k = IndependentMOKernel(GaussianKernel())
+        @test kernelmatrix(k, x, y) == k.(collect(x), permutedims(collect(y)))
+        @test kernelmatrix(k, x, x) == kernelmatrix(k, x)
+    end
 end

test/mokernels/independent.jl

@@ -0,0 +1,14 @@
+@testset "independent" begin
+    x = MOInput([rand(5) for _ in 1:4], 3)
+    y = MOInput([rand(5) for _ in 1:4], 3)
+
+    k = IndependentMOKernel(GaussianKernel())
+    @test k isa IndependentMOKernel
+    @test k isa Kernel
+    @test k.kernel isa KernelFunctions.BaseKernel
+    @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)
+    @test string(k) == "Independent Multi-Output Kernel\n\tSquared Exponential Kernel"
+end

test/mokernels/moinput.jl

@@ -0,0 +1,21 @@
+@testset "moinput" begin
+    
+    x =  [rand(5) for _ in 1:4]
+    mgpi = MOInput(x, 3)
+    
+    @test length(mgpi) == 12
+    @test size(mgpi) == (12,)
+    @test size(mgpi, 1) == 12
+    @test size(mgpi, 2) == 1
+    @test lastindex(mgpi) == 12
+    @test firstindex(mgpi) == 1
+    @test iterate(mgpi) == (mgpi[1], 1)
+    @test iterate(mgpi, 2) == (mgpi[3], 3)
+    @test_throws BoundsError mgpi[0]
+
+    @test mgpi[2] == (x[2], 1)
+    @test mgpi[5] == (x[1], 2)
+    @test mgpi[7] == (x[3], 2)
+    @test all([(x_, i) for i in 1:3 for x_ in x ] .== mgpi)
+
+end

test/runtests.jl

@@ -105,6 +105,12 @@ using KernelFunctions: metric, kappa, ColVecs, RowVecs
     end
     @info "Ran tests on matrix"
 
+    @testset "multi_output" begin
+        include(joinpath("mokernels", "moinput.jl"))
+        include(joinpath("mokernels", "independent.jl"))
+    end
+    @info "Ran tests on Multi-Output Kernels"
+
     @testset "approximations" begin
         include(joinpath("approximations", "nystrom.jl"))
     end