Use Iterators.peel, reduce allocations, and add missing kernelmatrix_diag! definitions + tests

Project.toml

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

src/kernels/kerneltensorproduct.jl

@@ -71,10 +71,17 @@ function kernelmatrix!(K::AbstractMatrix, k::KernelTensorProduct, x::AbstractVec
     validate_inplace_dims(K, x)
     validate_domain(k, x)
 
-    kernels_and_inputs = zip(k.kernels, slices(x))
-    kernelmatrix!(K, first(kernels_and_inputs)...)
-    for (k, xi) in Iterators.drop(kernels_and_inputs, 1)
-        K .*= kernelmatrix(k, xi)
+    first_kernels, tail_kernels = Iterators.peel(k.kernels)
+    first_x, tail_x = Iterators.peel(slices(x))
+
+    # handle first kernel and input
+    kernelmatrix!(K, first_kernels, first_x)
+
+    # handle remaining kernels and inputs
+    Ktmp = similar(K)
+    for (ki, xi) in zip(tail_kernels, tail_x)
+        kernelmatrix!(Ktmp, ki, xi)
+        hadamard!(K, K, Ktmp)
     end
 
     return K
@@ -86,10 +93,18 @@ function kernelmatrix!(
     validate_inplace_dims(K, x, y)
     validate_domain(k, x)
 
-    kernels_and_inputs = zip(k.kernels, slices(x), slices(y))
-    kernelmatrix!(K, first(kernels_and_inputs)...)
-    for (k, xi, yi) in Iterators.drop(kernels_and_inputs, 1)
-        K .*= kernelmatrix(k, xi, yi)
+    first_kernels, tail_kernels = Iterators.peel(k.kernels)
+    first_x, tail_x = Iterators.peel(slices(x))
+    first_y, tail_y = Iterators.peel(slices(y))
+
+    # handle first kernel and inputs
+    kernelmatrix!(K, first_kernels, first_x, first_y)
+
+    # handle remaining kernels and inputs
+    Ktmp = similar(K)
+    for (ki, xi, yi) in zip(tail_kernels, tail_x, tail_y)
+        kernelmatrix!(Ktmp, ki, xi, yi)
+        hadamard!(K, K, Ktmp)
     end
 
     return K
@@ -99,10 +114,40 @@ function kernelmatrix_diag!(K::AbstractVector, k::KernelTensorProduct, x::Abstra
     validate_inplace_dims(K, x)
     validate_domain(k, x)
 
-    kernels_and_inputs = zip(k.kernels, slices(x))
-    kernelmatrix_diag!(K, first(kernels_and_inputs)...)
-    for (k, xi) in Iterators.drop(kernels_and_inputs, 1)
-        K .*= kernelmatrix_diag(k, xi)
+    first_kernels, tail_kernels = Iterators.peel(k.kernels)
+    first_x, tail_x = Iterators.peel(slices(x))
+
+    # handle first kernel and input
+    kernelmatrix_diag!(K, first_kernels, first_x)
+
+    # handle remaining kernels and inputs
+    Ktmp = similar(K)
+    for (ki, xi) in zip(tail_kernels, tail_x)
+        kernelmatrix_diag!(Ktmp, ki, xi)
+        hadamard!(K, K, Ktmp)
+    end
+
+    return K
+end
+
+function kernelmatrix_diag!(
+    K::AbstractVector, k::KernelTensorProduct, x::AbstractVector, y::AbstractVector
+)
+    validate_inplace_dims(K, x, y)
+    validate_domain(k, x)
+
+    first_kernels, tail_kernels = Iterators.peel(k.kernels)
+    first_x, tail_x = Iterators.peel(slices(x))
+    first_y, tail_y = Iterators.peel(slices(y))
+
+    # handle first kernel and inputs
+    kernelmatrix_diag!(K, first_kernels, first_x, first_y)
+
+    # handle remaining kernels and inputs
+    Ktmp = similar(K)
+    for (ki, xi, yi) in zip(tail_kernels, tail_x, tail_y)
+        kernelmatrix_diag!(Ktmp, ki, xi, yi)
+        hadamard!(K, K, Ktmp)
     end
 
     return K

src/kernels/scaledkernel.jl

@@ -61,6 +61,14 @@ function kernelmatrix_diag!(K::AbstractVector, κ::ScaledKernel, x::AbstractVect
     return K
 end
 
+function kernelmatrix_diag!(
+    K::AbstractVector, κ::ScaledKernel, x::AbstractVector, y::AbstractVector
+)
+    kernelmatrix_diag!(K, κ.kernel, x, y)
+    K .*= κ.σ²
+    return K
+end
+
 Base.:*(w::Real, k::Kernel) = ScaledKernel(k, w)
 
 Base.show(io::IO, κ::ScaledKernel) = printshifted(io, κ, 0)

src/test_utils.jl

@@ -87,6 +87,7 @@ function test_interface(
 
     tmp_diag = Vector{Float64}(undef, length(x0))
     @test kernelmatrix_diag!(tmp_diag, k, x0) ≈ kernelmatrix_diag(k, x0)
+    @test kernelmatrix_diag!(tmp_diag, k, x0, x1) ≈ kernelmatrix_diag(k, x0, x1)
 end
 
 function test_interface(

src/utils.jl

@@ -197,17 +197,27 @@ function validate_inplace_dims(K::AbstractMatrix, x::AbstractVector, y::Abstract
     end
 end
 
-function validate_inplace_dims(K::AbstractMatrix, x::AbstractVector)
-    return validate_inplace_dims(K, x, x)
-end
-
-function validate_inplace_dims(K::AbstractVector, x::AbstractVector)
-    if length(K) != length(x)
+function validate_inplace_dims(K::AbstractVector, x::AbstractVector, y::AbstractVector)
+    validate_inputs(x, y)
+    n = length(x)
+    if length(y) != n
+        throw(
+            DimensionMismatch(
+                "Length of input x ($n) not consistent with length of input y " *
+                "($(length(y))",
+            ),
+        )
+    end
+    if length(K) != n
         throw(
             DimensionMismatch(
-                "Length of target vector K ($(length(K))) not consistent with length of input" *
-                "vector x ($(length(x))",
+                "Length of target vector K ($(length(K))) not consistent with length of " *
+                "inputs ($n)",
             ),
         )
     end
 end
+
+function validate_inplace_dims(K::AbstractVecOrMat, x::AbstractVector)
+    return validate_inplace_dims(K, x, x)
+end