Adds kernel related to the Intrinsic Coregionalization Model (#263)

* Adds first version of Intrinsic Coregionalization Multi-Output Kernel.

* Changes the name of the kernel. Removes "implementation code" for the
parameterization of the coregionalization matrix.

Removes white space.

Renames file.

Updates names in module file.

Corrects check_args.

Corrects using incorrect field name in coregion covariance function.

Cleans white space.

* Adds tests for CoregionMOKernel.

* Adds coregion tests to runtests file.

* Adds docstring to coregion kernel.

* Update src/mokernels/coregion.jl

Makes parameter explicit in function definition.

Co-authored-by: David Widmann <[email protected]>

* Makes notation in docstring conformant with the rest of the package.

Co-authored-by: David Widmann <[email protected]>

* Makes code respect formating style guidelines.

* Adds Coregion kernel to docs file.

* Removes unused argument.

* Creates new methods for Finitedifferences.to_vec to deal with the
inputs of multi-output kernels.

* Creates new methods to compute the jacobian of multioutput kernels.

* Formatting.

* Reworkes to_vec so that the reimplementation of _jvp and jacobian is
not needed.

* Renames kernel.

* Changes all affected files by the kernel renaming.

* Adds test_interface method to generate inputs for MOkernels.

* Makes IntrinsicCoregionMOKernel use the standard suit of tests.

* Improves MOkernels tests to test different outputs dimensions.

* Passes RNG object to rand method.

* Adds MOKernel methods for testfunction and testdiagfunction.

* Adds test_FiniteDiff method for MOKernels.

* Changes test_FiniteDiff for MOKernel signature to have the same number
of arguments as the methods for other kernels.

* Corrects bug in tests where the created covariance matrix between outputs could
be non PSD.

* Adapts compare_gradient for MOKernels.

* Adapts test_AD for MOKernels.

* Adapts test_ADs for MOKernels.

* Corrects type annotation.

* Corrects argument splatting.

* Rewrites variables as named tuple.

* Updates arguments.

* Adds standardised AD tests for IntrinsicCoregionMOKernel.

* Adds include of finitedifferences file to runtests.

* Passes dims argument.

* Formatting.

* Formatting.

* Corrects a bug where the wrong tuple fields were used.

* Corrects test.

* Implements @theogf approach to testing gradients of MOkernels.

* With new approach this file is not needed.

* Correctly implement @theogf approach to MOkernels AD testing.

* Changes tests to use SqExponentialKernel instead of
ExponentialKernel because of the latter's AD instability.

* Deletes unsused include.

This file is not needed anymore.

Co-authored-by: David Widmann <[email protected]>

* Simplifies docstring.

Co-authored-by: David Widmann <[email protected]>

* Improves error message.

Co-authored-by: David Widmann <[email protected]>

* Deletes file because we currently pass just the Real part to gradient
functions when testing.

* Adds kernel and number of outputs when printing Icoregionkernel.

* Adds test for different inputs as arguments.

Co-authored-by: David Widmann <[email protected]>

* Stops infringing character limit.

Co-authored-by: David Widmann <[email protected]>

* Further simplify docstring.

Co-authored-by: David Widmann <[email protected]>

* Adds tests.

* Deletes unuseful tests.

* Updates print test.

* Stops using interpolations.

* Creates constructor with keyword arguments.

* Restore constructor.

* Deletes non kwargs outer constructor.

Co-authored-by: David Widmann <[email protected]>

* Formatting.

* Updates code to be in consistent with kwargs constructor.

* Updates docstring.

* Update Project.toml

* Formatting.

Co-authored-by: David Widmann <[email protected]>

Author: david-vicente

Project.toml

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

docs/src/kernels.md

@@ -133,4 +133,5 @@ NormalizedKernel
 MOKernel
 IndependentMOKernel
 LatentFactorMOKernel
+IntrinsicCoregionMOKernel
 ```

src/KernelFunctions.jl

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

src/mokernels/intrinsiccoregion.jl

@@ -0,0 +1,45 @@
+@doc raw"""
+    IntrinsicIntrinsicCoregionMOKernel(; kernel::Kernel, B::AbstractMatrix)
+
+Kernel associated with the intrinsic coregionalization model.
+
+# Definition
+
+For inputs ``x, x'`` and output dimensions ``p, p'``, the kernel is defined as[^ARL]
+```math
+k\big((x, p), (x', p'); B, \tilde{k}\big) = B_{p, p'} \tilde{k}\big(x, x'\big),
+```
+where ``B`` is a positive semidefinite matrix of size ``m \times m``, with ``m`` being the
+number of outputs, and ``\tilde{k}`` is a scalar-valued kernel shared by the latent
+processes.
+
+[^ARL]: M. Álvarez, L. Rosasco, & N. Lawrence (2012). [Kernels for Vector-Valued Functions: a Review](https://arxiv.org/pdf/1106.6251.pdf).
+"""
+struct IntrinsicCoregionMOKernel{K<:Kernel,T<:AbstractMatrix} <: MOKernel
+    kernel::K
+    B::T
+
+    function IntrinsicCoregionMOKernel{K,T}(kernel::K, B::T) where {K,T}
+        @check_args(
+            IntrinsicCoregionMOKernel,
+            B,
+            eigmin(B) >= 0,
+            "B has to be positive semi-definite"
+        )
+        return new{K,T}(kernel, B)
+    end
+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 Base.show(io::IO, k::IntrinsicCoregionMOKernel)
+    return print(
+        io, "Intrinsic Coregion Kernel: ", k.kernel, " with ", size(k.B, 1), " outputs"
+    )
+end

src/test_utils.jl

@@ -97,6 +97,18 @@ function test_interface(
     )
 end
 
+function test_interface(
+    rng::AbstractRNG, k::MOKernel, ::Type{Vector{Tuple{T,Int}}}; dim_out=1, kwargs...
+) where {T<:Real}
+    return test_interface(
+        k,
+        [(randn(rng, T), rand(rng, 1:dim_out)) for i in 1:51],
+        [(randn(rng, T), rand(rng, 1:dim_out)) for i in 1:51],
+        [(randn(rng, T), rand(rng, 1:dim_out)) for i in 1:50];
+        kwargs...,
+    )
+end
+
 function test_interface(
     rng::AbstractRNG, k::Kernel, ::Type{<:ColVecs{T}}; dim_in=2, kwargs...
 ) where {T<:Real}

test/mokernels/intrinsiccoregion.jl

@@ -0,0 +1,28 @@
+@testset "intrinsiccoregion" begin
+    rng = MersenneTwister(123)
+
+    dims = (in=3, out=2, obs=3)
+    rank = 1
+
+    A = randn(dims.out, rank)
+    B = A * transpose(A) + Diagonal(rand(dims.out))
+
+    X = [(rand(dims.in), rand(1:(dims.out))) for i in 1:(dims.obs)]
+
+    kernel = SqExponentialKernel()
+    icoregionkernel = IntrinsicCoregionMOKernel(; kernel=kernel, B=B)
+
+    @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])
+
+    KernelFunctions.TestUtils.test_interface(
+        icoregionkernel, Vector{Tuple{Float64,Int}}; dim_out=dims.out
+    )
+
+    test_ADs(icoregionkernel; dims=dims)
+
+    @test string(icoregionkernel) ==
+          string("Intrinsic Coregion Kernel: ", kernel, " with ", dims.out, " outputs")
+end

test/runtests.jl

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

test/test_utils.jl

@@ -53,6 +53,11 @@ testfunction(k, A, dim) = sum(kernelmatrix(k, A; obsdim=dim))
 testdiagfunction(k, A, dim) = sum(kernelmatrix_diag(k, A; obsdim=dim))
 testdiagfunction(k, A, B, dim) = sum(kernelmatrix_diag(k, A, B; obsdim=dim))
 
+testfunction(k::MOKernel, A, B) = sum(kernelmatrix(k, A, B))
+testfunction(k::MOKernel, A) = sum(kernelmatrix(k, A))
+testdiagfunction(k::MOKernel, A) = sum(kernelmatrix_diag(k, A))
+testdiagfunction(k::MOKernel, A, B) = sum(kernelmatrix_diag(k, A, B))
+
 function test_ADs(
     kernelfunction, args=nothing; ADs=[:Zygote, :ForwardDiff, :ReverseDiff], dims=[3, 3]
 )
@@ -64,6 +69,17 @@ function test_ADs(
     end
 end
 
+function test_ADs(
+    k::MOKernel; ADs=[:Zygote, :ForwardDiff, :ReverseDiff], dims=(in=3, out=2, obs=3)
+)
+    test_fd = test_FiniteDiff(k, dims)
+    if !test_fd.anynonpass
+        for AD in ADs
+            test_AD(AD, k, dims)
+        end
+    end
+end
+
 function test_FiniteDiff(kernelfunction, args=nothing, dims=[3, 3])
     # Init arguments :
     k = if args === nothing
@@ -128,6 +144,50 @@ function test_FiniteDiff(kernelfunction, args=nothing, dims=[3, 3])
     end
 end
 
+function test_FiniteDiff(k::MOKernel, dims=(in=3, out=2, obs=3))
+    rng = MersenneTwister(42)
+    @testset "FiniteDifferences" begin
+        ## Testing Kernel Functions
+        x = (rand(rng, dims.in), rand(rng, 1:(dims.out)))
+        y = (rand(rng, dims.in), rand(rng, 1:(dims.out)))
+
+        @test_nowarn gradient(:FiniteDiff, x[1]) do a
+            k((a, x[2]), y)
+        end
+
+        ## Testing Kernel Matrices
+
+        A = [(randn(rng, dims.in), rand(rng, 1:(dims.out))) for i in 1:(dims.obs)]
+        B = [(randn(rng, dims.in), rand(rng, 1:(dims.out))) for i in 1:(dims.obs)]
+
+        @test_nowarn gradient(:FiniteDiff, reduce(hcat, first.(A))) do a
+            A = tuple.(eachcol(a), last.(A))
+            testfunction(k, A)
+        end
+        @test_nowarn gradient(:FiniteDiff, reduce(hcat, first.(A))) do a
+            A = tuple.(eachcol(a), last.(A))
+            testfunction(k, A, B)
+        end
+        @test_nowarn gradient(:FiniteDiff, reduce(hcat, first.(B))) do b
+            B = tuple.(eachcol(b), last.(B))
+            testfunction(k, A, B)
+        end
+
+        @test_nowarn gradient(:FiniteDiff, reduce(hcat, first.(A))) do a
+            A = tuple.(eachcol(a), last.(A))
+            testdiagfunction(k, A)
+        end
+        @test_nowarn gradient(:FiniteDiff, reduce(hcat, first.(A))) do a
+            A = tuple.(eachcol(a), last.(A))
+            testdiagfunction(k, A, B)
+        end
+        @test_nowarn gradient(:FiniteDiff, reduce(hcat, first.(B))) do b
+            B = tuple.(eachcol(b), last.(B))
+            testdiagfunction(k, A, B)
+        end
+    end
+end
+
 function test_AD(AD::Symbol, kernelfunction, args=nothing, dims=[3, 3])
     @testset "$(AD)" begin
         # Test kappa function
@@ -194,3 +254,49 @@ function test_AD(AD::Symbol, kernelfunction, args=nothing, dims=[3, 3])
         end
     end
 end
+
+function test_AD(AD::Symbol, k::MOKernel, dims=(in=3, out=2, obs=3))
+    @testset "$(AD)" begin
+        rng = MersenneTwister(42)
+
+        # Testing kernel evaluations
+        x = (rand(rng, dims.in), rand(rng, 1:(dims.out)))
+        y = (rand(rng, dims.in), rand(rng, 1:(dims.out)))
+
+        compare_gradient(AD, x[1]) do a
+            k((a, x[2]), y)
+        end
+        compare_gradient(AD, y[1]) do b
+            k(x, (b, y[2]))
+        end
+
+        # Testing kernel matrices
+        A = [(randn(rng, dims.in), rand(rng, 1:(dims.out))) for i in 1:(dims.obs)]
+        B = [(randn(rng, dims.in), rand(rng, 1:(dims.out))) for i in 1:(dims.obs)]
+
+        compare_gradient(AD, reduce(hcat, first.(A))) do a
+            A = tuple.(eachcol(a), last.(A))
+            testfunction(k, A)
+        end
+        compare_gradient(AD, reduce(hcat, first.(A))) do a
+            A = tuple.(eachcol(a), last.(A))
+            testfunction(k, A, B)
+        end
+        compare_gradient(AD, reduce(hcat, first.(B))) do b
+            B = tuple.(eachcol(b), last.(B))
+            testfunction(k, A, B)
+        end
+        compare_gradient(AD, reduce(hcat, first.(A))) do a
+            A = tuple.(eachcol(a), last.(A))
+            testdiagfunction(k, A)
+        end
+        compare_gradient(AD, reduce(hcat, first.(A))) do a
+            A = tuple.(eachcol(a), last.(A))
+            testdiagfunction(k, A, B)
+        end
+        compare_gradient(AD, reduce(hcat, first.(B))) do b
+            B = tuple.(eachcol(b), last.(B))
+            testdiagfunction(k, A, B)
+        end
+    end
+end