Merge branch 'master' into add-gibbskernel-to-prior-example

Author: Cyberface

Project.toml

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

docs/src/api.md

@@ -84,7 +84,6 @@ To find out more about the background, read this [review of kernels for vector-v
 
 KernelFunctions also provides miscellaneous utility functions.
 ```@docs
-kernelpdmat
 nystrom
 NystromFact
 ```
@@ -96,4 +95,11 @@ To keep the dependencies of KernelFunctions lean, some functionality is only ava
 [*https://github.com/MichielStock/Kronecker.jl*](https://github.com/MichielStock/Kronecker.jl)
 ```@docs
 kronecker_kernelmatrix
+kernelkronmat
 ```
+
+### PDMats.jl
+[*https://github.com/JuliaStats/PDMats.jl*](https://github.com/JuliaStats/PDMats.jl)
+```@docs
+kernelpdmat
+```

src/KernelFunctions.jl

@@ -58,7 +58,10 @@ using IrrationalConstants: logtwo, twoπ, invsqrt2
 using LogExpFunctions: softplus
 using StatsBase
 using TensorCore
-using ZygoteRules: ZygoteRules
+using ZygoteRules: ZygoteRules, AContext, literal_getproperty, literal_getfield
+
+# Hack to work around Zygote type inference problems.
+const Distances_pairwise = Distances.pairwise
 
 abstract type Kernel end
 abstract type SimpleKernel <: Kernel end

src/basekernels/fbm.jl

@@ -43,6 +43,8 @@ end
 _fbm(modX, modY, modXY, h) = (modX^h + modY^h - modXY^h) / 2
 
 _mod(x::AbstractVector{<:Real}) = abs2.(x)
+_mod(x::AbstractVector{<:AbstractVector{<:Real}}) = sum.(abs2, x)
+# two lines above could be combined into the second (dispatching on general AbstractVectors), but this (somewhat) more performant
 _mod(x::ColVecs) = vec(sum(abs2, x.X; dims=1))
 _mod(x::RowVecs) = vec(sum(abs2, x.X; dims=2))
 

src/basekernels/rational.jl

@@ -5,11 +5,11 @@ Rational kernel with shape parameter `α` and given `metric`.
 
 # Definition
 
-For inputs ``x, x'``, the rational kernel with shape parameter
-``\\alpha > 0`` is defined as
+For inputs ``x, x'`` and metric ``d(\\cdot, \\cdot)``, the rational kernel with shape parameter ``\\alpha > 0`` is defined as
 ```math
-k(x, x'; \\alpha) = \\bigg(1 + \\frac{\\|x - x'\\|}{\\alpha}\\bigg)^{-\\alpha}.
+k(x, x'; \\alpha) = \\bigg(1 + \\frac{d(x, x')}{\\alpha}\\bigg)^{-\\alpha}.
 ```
+By default, ``d`` is the Euclidean metric ``d(x, x') = \\|x - x'\\|_2``.
 
 The [`ExponentialKernel`](@ref) is recovered in the limit as ``\\alpha \\to \\infty``.
 

src/distances/delta.jl

@@ -1,19 +1,18 @@
 # Delta is not following the PreMetric rules since d(x, x) == 1
 struct Delta <: Distances.UnionPreMetric end
 
-@inline function Distances._evaluate(::Delta, a::AbstractVector, b::AbstractVector)
+(dist::Delta)(a::Number, b::Number) = a == b
+Base.@propagate_inbounds function (dist::Delta)(
+    a::AbstractArray{<:Number}, b::AbstractArray{<:Number}
+)
     @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)).",
+                "first array has length $(length(a)) which does not match the length of the second, $(length(b)).",
             ),
         )
     end
     return a == b
 end
 
 Distances.result_type(::Delta, Ta::Type, Tb::Type) = Bool
-
-@inline (dist::Delta)(a::AbstractArray, b::AbstractArray) = Distances._evaluate(dist, a, b)
-@inline (dist::Delta)(a::Number, b::Number) = a == b

src/distances/pairwise.jl

@@ -7,17 +7,17 @@ end
 pairwise(d::PreMetric, X::AbstractVector) = pairwise(d, X, X)
 
 function pairwise!(out::AbstractMatrix, d::PreMetric, X::AbstractVector, Y::AbstractVector)
-    return broadcast!(d, out, X, Y')
+    return broadcast!(d, out, X, permutedims(Y))
 end
 
 pairwise!(out::AbstractMatrix, d::PreMetric, X::AbstractVector) = pairwise!(out, d, X, X)
 
 function pairwise(d::PreMetric, x::AbstractVector{<:Real})
-    return Distances.pairwise(d, reshape(x, :, 1); dims=1)
+    return Distances_pairwise(d, reshape(x, :, 1); dims=1)
 end
 
 function pairwise(d::PreMetric, x::AbstractVector{<:Real}, y::AbstractVector{<:Real})
-    return Distances.pairwise(d, reshape(x, :, 1), reshape(y, :, 1); dims=1)
+    return Distances_pairwise(d, reshape(x, :, 1), reshape(y, :, 1); dims=1)
 end
 
 function pairwise!(out::AbstractMatrix, d::PreMetric, x::AbstractVector{<:Real})

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/matrix/kernelkroneckermat.jl

@@ -6,27 +6,58 @@ using .Kronecker: Kronecker
 export kernelkronmat
 export kronecker_kernelmatrix
 
-function kernelkronmat(κ::Kernel, X::AbstractVector, dims::Int)
-    @assert iskroncompatible(κ) "The chosen kernel is not compatible for kroenecker matrices (see [`iskroncompatible`](@ref))"
-    k = kernelmatrix(κ, X)
-    return Kronecker.kronecker(k, dims)
+@doc raw"""
+    kernelkronmat(κ::Kernel, X::AbstractVector{<:Real}, dims::Int) -> KroneckerPower
+
+Return a `KroneckerPower` matrix on the `D`-dimensional input grid constructed by ``\otimes_{i=1}^D X``,
+where `D` is given by `dims`.
+
+!!! warning
+
+    Require `Kronecker.jl` and for `iskroncompatible(κ)` to return `true`.
+"""
+function kernelkronmat(κ::Kernel, X::AbstractVector{<:Real}, dims::Int)
+    checkkroncompatible(κ)
+    K = kernelmatrix(κ, X)
+    return Kronecker.kronecker(K, dims)
 end
 
-function kernelkronmat(
-    κ::Kernel, X::AbstractVector{<:AbstractVector}; obsdim::Int=defaultobs
-)
-    @assert iskroncompatible(κ) "The chosen kernel is not compatible for Kronecker matrices"
+@doc raw"""
+
+    kernelkronmat(κ::Kernel, X::AbstractVector{<:AbstractVector}) -> KroneckerProduct
+
+Returns a `KroneckerProduct` matrix on the grid built with the collection of vectors ``\{X_i\}_{i=1}^D``: ``\otimes_{i=1}^D X_i``.
+
+!!! warning
+
+    Requires `Kronecker.jl` and for `iskroncompatible(κ)` to return `true`.
+"""
+function kernelkronmat(κ::Kernel, X::AbstractVector{<:AbstractVector})
+    checkkroncompatible(κ)
     Ks = kernelmatrix.(κ, X)
-    return K = reduce(Kronecker.:⊗, Ks)
+    return reduce(Kronecker.:⊗, Ks)
 end
 
-"""
-    To be compatible with kroenecker constructions the kernel must satisfy
-    the property : for x,x' ∈ ℜᴰ
-    k(x,x') = ∏ᵢᴰ k(xᵢ,x'ᵢ)
+@doc raw"""
+    iskroncompatible(k::Kernel)
+
+Determine whether kernel `k` is compatible with Kronecker constructions such as [`kernelkronmat`](@ref)
+
+The function returns `false` by default. If `k` is compatible it must satisfy for all ``x, x' \in \mathbb{R}^D`:
+```math
+k(x, x') = \prod_{i=1}^D k(x_i, x'_i).
+```
 """
 @inline iskroncompatible(κ::Kernel) = false # Default return for kernels
 
+function checkkroncompatible(κ::Kernel)
+    return iskroncompatible(κ) || throw(
+        ArgumentError(
+            "The chosen kernel is not compatible for Kronecker matrices (see [`iskroncompatible`](@ref))",
+        ),
+    )
+end
+
 function _kernelmatrix_kroneckerjl_helper(
     ::Type{<:MOInputIsotopicByFeatures}, Kfeatures, Koutputs
 )

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(
@@ -133,6 +134,18 @@ function test_interface(
     )
 end
 
+function test_interface(
+    rng::AbstractRNG, k::Kernel, ::Type{<:Vector{Vector{T}}}; dim_in=2, kwargs...
+) where {T<:Real}
+    return test_interface(
+        k,
+        [randn(rng, T, dim_in) for _ in 1:1001],
+        [randn(rng, T, dim_in) for _ in 1:1001],
+        [randn(rng, T, dim_in) for _ in 1:1000];
+        kwargs...,
+    )
+end
+
 function test_interface(k::Kernel, T::Type{<:AbstractVector}; kwargs...)
     return test_interface(Random.GLOBAL_RNG, k, T; kwargs...)
 end
@@ -147,6 +160,9 @@ function test_interface(rng::AbstractRNG, k::Kernel, T::Type{<:Real}; kwargs...)
     @testset "RowVecs{$T}" begin
         test_interface(rng, k, RowVecs{T}; kwargs...)
     end
+    @testset "Vector{Vector{T}}" begin
+        test_interface(rng, k, Vector{Vector{T}}; kwargs...)
+    end
 end
 
 function test_interface(k::Kernel, T::Type{<:Real}=Float64; kwargs...)

src/utils.jl

@@ -80,13 +80,13 @@ Base.vcat(a::ColVecs, b::ColVecs) = ColVecs(hcat(a.X, b.X))
 
 dim(x::ColVecs) = size(x.X, 1)
 
-pairwise(d::PreMetric, x::ColVecs) = Distances.pairwise(d, x.X; dims=2)
-pairwise(d::PreMetric, x::ColVecs, y::ColVecs) = Distances.pairwise(d, x.X, y.X; dims=2)
+pairwise(d::PreMetric, x::ColVecs) = Distances_pairwise(d, x.X; dims=2)
+pairwise(d::PreMetric, x::ColVecs, y::ColVecs) = Distances_pairwise(d, x.X, y.X; dims=2)
 function pairwise(d::PreMetric, x::AbstractVector, y::ColVecs)
-    return Distances.pairwise(d, reduce(hcat, x), y.X; dims=2)
+    return Distances_pairwise(d, reduce(hcat, x), y.X; dims=2)
 end
 function pairwise(d::PreMetric, x::ColVecs, y::AbstractVector)
-    return Distances.pairwise(d, x.X, reduce(hcat, y); dims=2)
+    return Distances_pairwise(d, x.X, reduce(hcat, y); dims=2)
 end
 function pairwise!(out::AbstractMatrix, d::PreMetric, x::ColVecs)
     return Distances.pairwise!(out, d, x.X; dims=2)
@@ -150,13 +150,13 @@ Base.vcat(a::RowVecs, b::RowVecs) = RowVecs(vcat(a.X, b.X))
 
 dim(x::RowVecs) = size(x.X, 2)
 
-pairwise(d::PreMetric, x::RowVecs) = Distances.pairwise(d, x.X; dims=1)
-pairwise(d::PreMetric, x::RowVecs, y::RowVecs) = Distances.pairwise(d, x.X, y.X; dims=1)
+pairwise(d::PreMetric, x::RowVecs) = Distances_pairwise(d, x.X; dims=1)
+pairwise(d::PreMetric, x::RowVecs, y::RowVecs) = Distances_pairwise(d, x.X, y.X; dims=1)
 function pairwise(d::PreMetric, x::AbstractVector, y::RowVecs)
-    return Distances.pairwise(d, permutedims(reduce(hcat, x)), y.X; dims=1)
+    return Distances_pairwise(d, permutedims(reduce(hcat, x)), y.X; dims=1)
 end
 function pairwise(d::PreMetric, x::RowVecs, y::AbstractVector)
-    return Distances.pairwise(d, x.X, permutedims(reduce(hcat, y)); dims=1)
+    return Distances_pairwise(d, x.X, permutedims(reduce(hcat, y)); dims=1)
 end
 function pairwise!(out::AbstractMatrix, d::PreMetric, x::RowVecs)
     return Distances.pairwise!(out, d, x.X; dims=1)
@@ -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