Add lazy kronecker product for matrix kernels, if Kronecker.jl is loaded

src/KernelFunctions.jl

@@ -105,8 +105,8 @@ include(joinpath("kernels", "neuralkernelnetwork.jl"))
 include(joinpath("approximations", "nystrom.jl"))
 include("generic.jl")
 
-include(joinpath("mokernels", "mokernel.jl"))
 include(joinpath("mokernels", "moinput.jl"))
+include(joinpath("mokernels", "mokernel.jl"))
 include(joinpath("mokernels", "independent.jl"))
 include(joinpath("mokernels", "slfm.jl"))
 include(joinpath("mokernels", "intrinsiccoregion.jl"))

src/matrix/kernelkroneckermat.jl

@@ -25,3 +25,11 @@ end
     k(x,x') = ∏ᵢᴰ k(xᵢ,x'ᵢ)
 """
 @inline iskroncompatible(κ::Kernel) = false # Default return for kernels
+
+function _kernelmatrix_kron_helper(::MOInputIsotopicByFeatures, Kfeatures, Koutputs)
+    return Kronecker.kronecker(Kfeatures, Koutputs)
+end
+
+function _kernelmatrix_kron_helper(::MOInputIsotopicByOutputs, Kfeatures, Koutputs)
+    return Kronecker.kronecker(Koutputs, Kfeatures)
+end

src/mokernels/independent.jl

@@ -27,43 +27,9 @@ function (κ::IndependentMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{Any,I
     return κ.kernel(x, y) * (px == py)
 end
 
-function _kernelmatrix_kron_helper(::MOInputIsotopicByFeatures, Kfeatures, B)
-    return kron(Kfeatures, B)
-end
-
-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
-    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
+_mo_output_covariance(k::IndependentMOKernel, out_dim) = Eye{Bool}(out_dim)
 
-    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
+# for specialized kernelmatrix implementation see intrinsiccoregion.jl
 
 function Base.show(io::IO, k::IndependentMOKernel)
     return print(io, string("Independent Multi-Output Kernel\n\t", string(k.kernel)))

src/mokernels/intrinsiccoregion.jl

@@ -42,21 +42,33 @@ function (k::IntrinsicCoregionMOKernel)((x, px)::Tuple{Any,Int}, (y, py)::Tuple{
     return k.B[px, py] * k.kernel(x, y)
 end
 
+function _mo_output_covariance(k::IntrinsicCoregionMOKernel, out_dim)
+    @assert size(k.B) == (out_dim, out_dim)
+    return k.B
+end
+
 function kernelmatrix(
-    k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
-) where {MOI<:IsotopicMOInputsUnion}
+    k::Union{IndependentMOKernel,IntrinsicCoregionMOKernel},
+    x::IsotopicMOInputsUnion,
+    y::IsotopicMOInputsUnion,
+)
     @assert x.out_dim == y.out_dim
     Kfeatures = kernelmatrix(k.kernel, x.x, y.x)
-    return _kernelmatrix_kron_helper(x, Kfeatures, k.B)
+    Koutputs = _mo_output_covariance(k, x.out_dim)
+    return _kernelmatrix_kron_helper(x, Kfeatures, Koutputs)
 end
 
 if VERSION >= v"1.6"
     function kernelmatrix!(
-        K::AbstractMatrix, k::IntrinsicCoregionMOKernel, x::MOI, y::MOI
-    ) where {MOI<:IsotopicMOInputsUnion}
+        K::AbstractMatrix,
+        k::Union{IndependentMOKernel,IntrinsicCoregionMOKernel},
+        x::IsotopicMOInputsUnion,
+        y::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)
+        Koutputs = _mo_output_covariance(k, x.out_dim)
+        return _kernelmatrix_kron_helper!(K, x, Kfeatures, Koutputs)
     end
 end
 

src/mokernels/mokernel.jl

@@ -4,3 +4,21 @@
 Abstract type for kernels with multiple outpus.
 """
 abstract type MOKernel <: Kernel end
+
+function _kernelmatrix_kron_helper(::MOInputIsotopicByFeatures, Kfeatures, Koutputs)
+    return kron(Kfeatures, Koutputs)
+end
+
+function _kernelmatrix_kron_helper(::MOInputIsotopicByOutputs, Kfeatures, Koutputs)
+    return kron(Koutputs, Kfeatures)
+end
+
+if VERSION >= v"1.6"
+    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByFeatures, Kfeatures, Koutputs)
+        return kron!(K, Kfeatures, Koutputs)
+    end
+
+    function _kernelmatrix_kron_helper!(K, ::MOInputIsotopicByOutputs, Kfeatures, Koutputs)
+        return kron!(K, Koutputs, Kfeatures)
+    end
+end