[WIP] Use Tullio for pairwise distances

Project.toml

@@ -18,6 +18,7 @@ SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 TensorCore = "62fd8b95-f654-4bbd-a8a5-9c27f68ccd50"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Tullio = "bc48ee85-29a4-5162-ae0b-a64e1601d4bc"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [compat]

docs/src/kernels.md

@@ -100,7 +100,7 @@ GammaRationalKernel
 ### Spectral Mixture Kernels
 
 ```@docs
-spectral_mixture_kernel
+SpectralMixtureKernel
 spectral_mixture_product_kernel
 ```
 

src/KernelFunctions.jl

@@ -15,6 +15,7 @@ export LinearKernel, PolynomialKernel
 export RationalKernel, RationalQuadraticKernel, GammaRationalKernel
 export PiecewisePolynomialKernel
 export PeriodicKernel, NeuralNetworkKernel
+export SpectralMixtureKernel
 export KernelSum, KernelProduct, KernelTensorProduct
 export TransformedKernel, ScaledKernel, NormalizedKernel
 export GibbsKernel
@@ -33,7 +34,7 @@ export with_lengthscale
 export NystromFact, nystrom
 
 export gaborkernel
-export spectral_mixture_kernel, spectral_mixture_product_kernel
+export spectral_mixture_product_kernel
 
 export ColVecs, RowVecs
 
@@ -58,6 +59,7 @@ using IrrationalConstants: logtwo, twoπ, invsqrt2
 using LogExpFunctions: softplus
 using StatsBase
 using TensorCore
+using Tullio
 using ZygoteRules: ZygoteRules, AContext, literal_getproperty, literal_getfield
 
 # Hack to work around Zygote type inference problems.
@@ -67,6 +69,13 @@ abstract type Kernel end
 abstract type SimpleKernel <: Kernel end
 
 include("utils.jl")
+
+const VecOfVecs = Union{ColVecs,RowVecs}
+
+# A general binary op type not respecting Distances metric rules
+abstract type AbstractBinaryOp end
+const BinaryOp = Union{AbstractBinaryOp,Distances.PreMetric}
+
 include("distances/pairwise.jl")
 include("distances/dotproduct.jl")
 include("distances/delta.jl")
@@ -120,6 +129,8 @@ include("mokernels/lmm.jl")
 include("chainrules.jl")
 include("zygoterules.jl")
 
+include("deprecated.jl")
+
 include("test_utils.jl")
 
 function __init__()

src/basekernels/sm.jl

@@ -1,24 +1,39 @@
-"""
-    spectral_mixture_kernel(
+@doc raw"""
+    SpectralMixtureKernel(
         h::Kernel=SqExponentialKernel(),
-        αs::AbstractVector{<:Real},
-        γs::AbstractMatrix{<:Real},
-        ωs::AbstractMatrix{<:Real},
+        α::AbstractVector{<:Real},
+        γ::AbstractMatrix{<:Real},
+        ω::AbstractMatrix{<:Real},
+    )
+    SpectralMixtureKernel(
+        h::Kernel=SqExponentialKernel(),
+        α::AbstractVector{<:Real},
+        γ::AbstractVector{<:AbstractVecOrMat{<:Real}},
+        ω::AbstractVector{<:AbstractVecOrMat{<:Real}},
     )
 
-where αs are the weights of dimension (A, ), γs is the covariance matrix of
-dimension (D, A) and ωs are the mean vectors and is of dimension (D, A).
-Here, D is input dimension and A is the number of spectral components.
-
-`h` is the kernel, which defaults to [`SqExponentialKernel`](@ref) if not specified.
+Generalised Spectral Mixture kernel function as described in [1] in equation 6.
+This family of functions is dense in the family of stationary real-valued kernels with respect to the pointwise convergence.[1]
 
-Generalised Spectral Mixture kernel function. This family of functions is  dense
-in the family of stationary real-valued kernels with respect to the pointwise convergence.[1]
+## Definition
 
+For inputs ``x, x′ \in \mathbb{R}^D``, the spectral mixture kernel ``\tilde{k}`` with ``K`` mixture components, mixture weights ``\alpha \in \mathbb{R}^K``, linear transformations ``\gamma_1, \ldots, \gamma_K \in \mathbb{R}^D``, and frequencies ``\omega_1, \ldots, \omega_K \in \mathbb{R}^D`` derived from a translation-invariant kernel ``k`` is defined as
 ```math
-   κ(x, y) = αs' (h(-(γs' * t)^2) .* cos(π * ωs' * t), t = x - y
+\tilde{k}(x, x'; \alpha, \gamma_1, \ldots, \gamma_K, \omega_1, \ldots, \omega_K, k) = \sum_{i=1}^K \alpha_i k(\gamma_i \odot x, \gamma_i \odot y) \cos(2\pi \omega_i^\top (x-y)).
 ```
 
+## Arguments
+- `h`: Stationary kernel (translation invariant), [`SqExponentialKernel`](@ref) by default
+- `α`: Weight vector of each mixture component (should be positive)
+- `γ`: Linear transformation of the input for `h`.
+- `ω`: Frequencies for the cosine function. (should be positive)
+
+`γ` and `ω` can be an
+- `AbstractMatrix` of dimension `D x K` where `D` is the dimension of the inputs 
+and `K` is the number of components
+- `AbstractVector` of `K` `D`-dimensional `AbstractVector`
+
+
 # References:
     [1] Generalized Spectral Kernels, by Yves-Laurent Kom Samo and Stephen J. Roberts
     [2] SM: Gaussian Process Kernels for Pattern Discovery and Extrapolation,
@@ -29,77 +44,129 @@ in the family of stationary real-valued kernels with respect to the pointwise co
     [4] http://www.cs.cmu.edu/~andrewgw/pattern/.
 
 """
-function spectral_mixture_kernel(
+struct SpectralMixtureKernel{
+    K<:Kernel,Tα<:AbstractVector,Tγ<:AbstractVector,Tω<:AbstractVector
+} <: Kernel
+    kernel::K
+    α::Tα
+    γ::Tγ
+    ω::Tω
+    function SpectralMixtureKernel(
+        h::Kernel,
+        α::AbstractVector{<:Real},
+        γ::AbstractVector{<:AbstractVector},
+        ω::AbstractVector{<:AbstractVector},
+    )
+        (length(α) == length(γ) == length(ω)) ||
+            throw(DimensionMismatch("The dimensions of α, γ, ans ω do not match"))
+        any(<(0), α) && throw(ArgumentError("At least one element of α is negative"))
+        any(any.(<(0), ω)) && throw(ArgumentError("At least one element of ω is negative"))
+        return new{typeof(h),typeof(α),typeof(γ),typeof(ω)}(h, α, γ, ω)
+    end
+end
+
+@functor SpectralMixtureKernel
+
+function SpectralMixtureKernel(
     h::Kernel,
-    αs::AbstractVector{<:Real},
-    γs::AbstractMatrix{<:Real},
-    ωs::AbstractMatrix{<:Real},
+    α::AbstractVector{<:Real},
+    γ::AbstractMatrix{<:Real},
+    ω::AbstractMatrix{<:Real},
 )
-    if !(size(αs, 1) == size(γs, 2) == size(ωs, 2))
-        throw(DimensionMismatch("The dimensions of αs, γs, ans ωs do not match"))
-    end
-    if size(γs) != size(ωs)
-        throw(DimensionMismatch("The dimensions of γs ans ωs do not match"))
-    end
+    size(γ) == size(ω) || throw(DimensionMismatch("γ and ω have different dimensions"))
+    return SpectralMixtureKernel(h, α, ColVecs(γ), ColVecs(ω))
+end
 
-    return sum(zip(αs, eachcol(γs), eachcol(ωs))) do (α, γ, ω)
-        a = TransformedKernel(h, LinearTransform(γ'))
-        b = TransformedKernel(CosineKernel(), LinearTransform(ω'))
-        return α * a * b
+function SpectralMixtureKernel(
+    αs::AbstractVector{<:Real}, γs::AbstractVecOrMat, ωs::AbstractVecOrMat
+)
+    return SpectralMixtureKernel(SqExponentialKernel(), αs, γs, ωs)
+end
+
+function (κ::SpectralMixtureKernel)(x, y)
+    xy = x - y
+    # use pairwise summation (https://github.com/JuliaLang/julia/pull/31020)
+    broadcasted = Broadcast.broadcasted(κ.α, κ.γ, κ.ω) do α, γ, ω
+        k = TransformedKernel(κ.kernel, ARDTransform(γ))
+        return α * k(x, y) * cospi(2 * dot(ω, xy))
     end
+    return sum(Broadcast.instantiate(broadcasted))
 end
 
-function spectral_mixture_kernel(
-    αs::AbstractVector{<:Real}, γs::AbstractMatrix{<:Real}, ωs::AbstractMatrix{<:Real}
-)
-    return spectral_mixture_kernel(SqExponentialKernel(), αs, γs, ωs)
+function Base.show(io::IO, κ::SpectralMixtureKernel)
+    return print(
+        io,
+        "SpectralMixtureKernel Kernel (kernel = ",
+        κ.kernel,
+        ", # components = ",
+        length(κ.α),
+        ")",
+    )
 end
 
-"""
+@doc raw"""
     spectral_mixture_product_kernel(
         h::Kernel=SqExponentialKernel(),
-        αs::AbstractMatrix{<:Real},
-        γs::AbstractMatrix{<:Real},
-        ωs::AbstractMatrix{<:Real},
+        α::AbstractMatrix{<:Real},
+        γ::AbstractMatrix{<:Real},
+        ω::AbstractMatrix{<:Real},
     )
 
-where αs are the weights of dimension (D, A), γs is the covariance matrix of
-dimension (D, A) and ωs are the mean vectors and is of dimension (D, A).
-Here, D is input dimension and A is the number of spectral components.
-
-Spectral Mixture Product Kernel. With enough components A, the SMP kernel
+The spectral mixture product is tensor product of spectral mixture kernel applied
+on each dimension as described in [1] in equations 13 and 14.
+With enough components, the SMP kernel
 can model any product kernel to arbitrary precision, and is flexible even
-with a small number of components [1]
+with a small number of components
 
+## Definition
 
-`h` is the kernel, which defaults to [`SqExponentialKernel`](@ref) if not specified.
+For inputs ``x, x′ \in \mathbb{R}^D``, the spectral mixture product kernel ``\tilde{k}`` with ``K`` mixture components, mixture weights ``\alpha_1, \alpha_2, \ldots, \alpha_K \in \mathbb{R}^D``, linear transformations ``\gamma_1, \ldots, \gamma_K \in \mathbb{R}^D``, and frequencies ``\omega_1, \ldots, \omega_K \in \mathbb{R}^D`` derived from a translation-invariant kernel ``k`` is defined as
 
 ```math
-   κ(x, y) = Πᵢ₌₁ᴷ Σ(αsᵢᵀ .* (h(-(γsᵢᵀ * tᵢ)²) .* cos(ωsᵢᵀ * tᵢ))), tᵢ = xᵢ - yᵢ
+   \tilde{k}(x, x'; \alpha_1, \ldots, \alpha_k, \gamma_1, \ldots, \gamma_K, \omega_1, \ldots, \omega_K, k) = \prod_{i=1}^D \sum_{k=1}^K \alpha_{ik} \cdot h(\gamma_{ik} \cdot x_i, \gamma_{ik} \cdot y_i)) \cdot \cos(2\pi \cdot \omega_{ik} \cdot (x_i - y_i))))
 ```
 
+## Arguments
+- `h`: Stationary kernel (translation invariant), [`SqExponentialKernel`](@ref) by default
+- `α`: Weight of each mixture component for each dimension
+- `γ`: Linear transformation of the input for `h`.
+- `ω`: Frequencies for the cosine function.
+
+`α`, `γ` and `ω` can be an
+- `AbstractMatrix` of dimension `D x K` where `D` is the dimension of the inputs 
+and `K` is the number of components
+- `AbstractVector` of `D` `K`-dimensional `AbstractVector`
+
+
 # References:
     [1] GPatt: Fast Multidimensional Pattern Extrapolation with GPs,
         arXiv 1310.5288, 2013, by Andrew Gordon Wilson, Elad Gilboa,
         Arye Nehorai and John P. Cunningham
 """
 function spectral_mixture_product_kernel(
     h::Kernel,
-    αs::AbstractMatrix{<:Real},
-    γs::AbstractMatrix{<:Real},
-    ωs::AbstractMatrix{<:Real},
+    α::AbstractMatrix{<:Real},
+    γ::AbstractMatrix{<:Real},
+    ω::AbstractMatrix{<:Real},
 )
-    if !(size(αs) == size(γs) == size(ωs))
-        throw(DimensionMismatch("The dimensions of αs, γs, ans ωs do not match"))
+    (size(α) == size(γ) == size(ω)) ||
+        throw(DimensionMismatch("α, γ and ω have different dimensions"))
+    return spectral_mixture_product_kernel(h, RowVecs(α), RowVecs(γ), RowVecs(ω))
+end
+
+function spectral_mixture_product_kernel(
+    h::Kernel,
+    α::AbstractVector{<:AbstractVector{<:Real}},
+    γ::AbstractVector{<:AbstractVector{<:Real}},
+    ω::AbstractVector{<:AbstractVector{<:Real}},
+)
+    return mapreduce(⊗, α, γ, ω) do αᵢ, γᵢ, ωᵢ
+        return SpectralMixtureKernel(h, αᵢ, permutedims(γᵢ), permutedims(ωᵢ))
     end
-    return KernelTensorProduct(
-        spectral_mixture_kernel(h, α, reshape(γ, 1, :), reshape(ω, 1, :)) for
-        (α, γ, ω) in zip(eachrow(αs), eachrow(γs), eachrow(ωs))
-    )
 end
 
 function spectral_mixture_product_kernel(
-    αs::AbstractMatrix{<:Real}, γs::AbstractMatrix{<:Real}, ωs::AbstractMatrix{<:Real}
+    α::AbstractVecOrMat, γ::AbstractVecOrMat, ω::AbstractVecOrMat
 )
-    return spectral_mixture_product_kernel(SqExponentialKernel(), αs, γs, ωs)
+    return spectral_mixture_product_kernel(SqExponentialKernel(), α, γ, ω)
 end

src/deprecated.jl

@@ -0,0 +1 @@
+@deprecate spectral_mixture_kernel SpectralMixtureKernel

src/distances/delta.jl

@@ -1,6 +1,6 @@
-# Delta is not following the PreMetric rules since d(x, x) == 1
-struct Delta <: Distances.UnionPreMetric end
+struct Delta <: AbstractBinaryOp end
 
+# Basic definitions
 (dist::Delta)(a::Number, b::Number) = a == b
 Base.@propagate_inbounds function (dist::Delta)(
     a::AbstractArray{<:Number}, b::AbstractArray{<:Number}
@@ -14,5 +14,3 @@ Base.@propagate_inbounds function (dist::Delta)(
     end
     return a == b
 end
-
-Distances.result_type(::Delta, Ta::Type, Tb::Type) = Bool

src/distances/dotproduct.jl

@@ -1,21 +1,22 @@
 ## DotProduct is not following the PreMetric rules since d(x, x) != 0 and d(x, y) >= 0 for all x, y
-struct DotProduct <: Distances.UnionPreMetric end
+struct DotProduct <: AbstractBinaryOp end
 
-@inline function Distances._evaluate(::DotProduct, a::AbstractVector, b::AbstractVector)
-    @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 dot(a, b)
+(::DotProduct)(a::AbstractVector, b::AbstractVector) = dot(a, b)
+
+(::DotProduct)(a::Number, b::Number) = a * b
+
+function pairwise(::DotProduct, x::ColVecs, y::ColVecs)
+    return @tullio out[i, j] := x.X[k, i] * y.X[k, j]
 end
 
-Distances.result_type(::DotProduct, Ta::Type, Tb::Type) = promote_type(Ta, Tb)
+function pairwise(::DotProduct, x::RowVecs, y::RowVecs)
+    return @tullio out[i, j] := x.X[i, k] * y.X[j, k]
+end
 
-@inline Distances.eval_op(::DotProduct, a::Real, b::Real) = a * b
-@inline function (dist::DotProduct)(a::AbstractArray, b::AbstractArray)
-    return Distances._evaluate(dist, a, b)
+function colwise(::DotProduct, x::RowVecs, y::RowVecs=x)
+    return @tullio out[i] := x.X[i, k] * y.X[i, k]
 end
-@inline (dist::DotProduct)(a::Number, b::Number) = a * b
+
+function colwise(::DotProduct, x::ColVecs, y::ColVecs=x)
+    return @tullio out[i] := x.X[k, i] * y.X[k, i]
+end