Merge branch 'master-dev'

Author: theogf

Project.toml

@@ -5,6 +5,8 @@ version = "0.1.0"
 [deps]
 Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
@@ -14,6 +16,7 @@ FiniteDifferences = ">= 0.7.2"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
 test = ["FiniteDifferences", "Random", "Test"]

dev/debugAD.jl

@@ -0,0 +1,50 @@
+using KernelFunctions
+using Zygote, ForwardDiff, Tracker
+using Test, LinearAlgebra
+
+dims = [10,5]
+A = rand(dims...)
+B = rand(dims...)
+K = [zeros(dims[1],dims[1]),zeros(dims[2],dims[2])]
+l = 0.1
+vl = l*ones(dims[1])
+testfunction(k,A,B) = det(kernelmatrix(k,A,B))
+testfunction(k,A) = sum(kernelmatrix(k,A))
+k = MaternKernel(vl)
+KernelFunctions.kappa(k,3)
+testfunction(SquaredExponentialKernel(vl),A)
+testfunction(MaternKernel(vl),A)
+@which kernelmatrix(MaternKernel(vl),A,B)
+#For debugging
+@info "Running Zygote gradients"
+Zygote.refresh()
+## Zygote
+Zygote.gradient(x->testfunction(SquaredExponentialKernel(x),A),vl)
+Zygote.gradient(x->testfunction(MaternKernel(x),A),vl)
+Zygote.gradient(x->testfunction(SquaredExponentialKernel(x),A,B),vl)[1]
+Zygote.gradient(x->testfunction(MaternKernel(x),A,B),vl)[1]
+Zygote.gradient(x->testfunction(SquaredExponentialKernel(x),A,B),l)
+Zygote.gradient(x->testfunction(MaternKernel(x),A,B),l)
+Zygote.gradient(x->testfunction(SquaredExponentialKernel(x),A),l)
+Zygote.gradient(x->testfunction(MaternKernel(x),A),l)
+Zygote.gradient(x->testfunction(MaternKernel(x),A),l)
+Zygote.gradient(x->kernelmatrix(MaternKernel(x,1.0),A)[1],l)
+@info "Running Tracker gradients"
+## Tracker
+# Tracker.gradient(x->testfunction(SquaredExponentialKernel(vl),x,B),A)
+# Tracker.gradient(x->testfunction(SquaredExponentialKernel(l),x[:,:]),A)
+# # Tracker.gradient(x->testfunction(SquaredExponentialKernel(x),A,B),vl)
+# Tracker.gradient(x->testfunction(SquaredExponentialKernel(x),A),vl)
+# Tracker.gradient(x->testfunction(SquaredExponentialKernel(x),A,B),l)
+# Tracker.gradient(x->testfunction(SquaredExponentialKernel(x),A),l)
+
+@info "Running ForwardDiff gradients"
+## ForwardDiff
+ForwardDiff.gradient(x->testfunction(SquaredExponentialKernel(x),A,B),vl) #✓
+ForwardDiff.gradient(x->testfunction(MaternKernel(x),A,B),vl) #✓
+ForwardDiff.gradient(x->testfunction(SquaredExponentialKernel(x),A),vl) #✓
+ForwardDiff.gradient(x->testfunction(MaternKernel(x),A),vl) #✓
+ForwardDiff.gradient(x->testfunction(SquaredExponentialKernel(x[1]),A,B),[l])
+ForwardDiff.gradient(x->testfunction(MaternKernel(x[1]),A,B),[l])
+ForwardDiff.gradient(x->testfunction(SquaredExponentialKernel(x[1]),A),[l])
+ForwardDiff.gradient(x->testfunction(MaternKernel(x[1]),A),[l])

dev/matrixvsvectors.jl

@@ -0,0 +1,28 @@
+using KernelFunctions
+using Stheno
+using Stheno: pw
+using BenchmarkTools
+using Zygote
+
+# Ds = [1,2,5,10,20,50,100,200,500,1000]
+Ds = [1,10,100,1000]
+timestheno = zeros(Float64,length(Ds)); memstheno = similar(timestheno)
+timekf = similar(timestheno); memkf = similar(timestheno)
+@progress for (i,D) in enumerate(Ds)
+
+    A = randn(D,1000)
+    B = randn(D,1001)
+
+    # Standardised eq kernel with length-scale 0.1.
+    medkf = median(@benchmark KernelFunctions.kernelmatrix(SquaredExponentialKernel(0.01),$A,$B,obsdim=2))
+    timekf[i] = medkf.time/1e6; memkf[i] = medkf.memory/2^20
+    medstheno = median(@benchmark pw(eq(; l=0.1), ColsAreObs($A), ColsAreObs($B)))
+    timestheno[i] = medstheno.time/1e6; memstheno[i] = medstheno.memory/2^20
+end
+
+using Plots
+ptime = plot(Ds,timestheno,lab="Stheno",xaxis=:log,xlabel="D",ylabel="t [ms]",title="Time")
+plot!(Ds,timekf,lab="KernelFunctions")
+pmem = plot(Ds,memstheno,lab="Stheno",xaxis=:log,xlabel="D",ylabel="Mem [MB]",title="Memory Usage")
+plot!(Ds,memkf,lab="KernelFunctions")
+plot(ptime,pmem)

src/KernelFunctions.jl

@@ -1,21 +1,29 @@
 module KernelFunctions
 
-export kernelmatrix, kernelmatrix!, kappa
-export Kernel, SquaredExponentialKernel
+export kernelmatrix, kernelmatrix!, kerneldiagmatrix, kerneldiagmatrix!, kappa
+export Kernel, SquaredExponentialKernel, MaternKernel, Matern32Kernel, Matern52Kernel
+
+export Transform, ScaleTransform
 
 using Distances, LinearAlgebra
+using Zygote: @adjoint
+using SpecialFunctions: lgamma, besselk
+using StatsFuns: logtwo
 
 const defaultobs = 2
-abstract type Kernel{T<:Real} end
+abstract type Kernel{T,Tr} end
 
 include("zygote_rules.jl")
 include("utils.jl")
-include("common.jl")
+include("transform/transform.jl")
 include("kernelmatrix.jl")
 
-kernels = ["squaredexponential"]
+kernels = ["squaredexponential","matern"]
 for k in kernels
     include(joinpath("kernels",k*".jl"))
 end
 
+include("generic.jl")
+
+
 end

src/common.jl

@@ -0,0 +1,15 @@
+
+@inline metric(κ::Kernel) = κ.metric
+kernels =
+for k in [:SquaredExponentialKernel,:MaternKernel,:Matern32Kernel,:Matern52Kernel]
+    eval(quote
+        @inline (κ::$k)(d::Real) = kappa(κ,d)
+        @inline (κ::$k)(x::AbstractVector{T},y::AbstractVector{T}) where {T} = kernel(κ,evaluate(κ.(metric),x,y))
+        @inline (κ::$k)(x::AbstractMatrix{T},y::AbstractMatrix{T},obsdim::Integer=defaultobs) where {T} = kernelmatrix(κ,x,y,obsdim=obsdim)
+    end)
+end
+### Transform generics
+
+@inline transform(κ::Kernel) = κ.transform
+@inline transform(κ::Kernel,x::AbstractVecOrMat) = transform(κ.transform,x)
+@inline transform(κ::Kernel,x::AbstractVecOrMat,obsdim::Int) = transform(κ.transform,x,obsdim)

src/generic.jl

@@ -1,26 +1,20 @@
-
-function _kappamatrix!(κ::Kernel{T}, P::AbstractMatrix{T₁}) where {T<:Real,T₁<:Real}
-    for i in eachindex(P)
-        @inbounds P[i] = kappa(κ, P[i])
-    end
-    P
-end
-
-function _symmetric_kappamatrix!(
+"""
+```
+    kernelmatrix!(K::Matrix, κ::Kernel, X::Matrix; obsdim::Integer=2, symmetrize::Bool=true)
+```
+In-place version of `kernelmatrix` where pre-allocated matrix `K` will be overwritten with the kernel matrix.
+"""
+function kernelmatrix!(
+        K::Matrix{T₁},
         κ::Kernel{T},
-        P::AbstractMatrix{T₁},
-        symmetrize::Bool
-    ) where {T<:Real,T₁<:Real}
-    if !((n = size(P,1)) == size(P,2))
-        throw(DimensionMismatch("Pairwise matrix must be square."))
-    end
-    for j = 1:n, i = (1:j)
-        @inbounds P[i,j] = kappa(κ, P[i,j])
-    end
-    symmetrize ? LinearAlgebra.copytri!(P, 'U') : P
+        X::AbstractMatrix{T₂};
+        obsdim::Int = defaultobs,
+        symmetrize::Bool = true
+        ) where {T,T₁<:Real,T₂<:Real}
+        @assert check_dims(K,X,X,obsdim) "Dimensions of the target array are not consistent with X and Y"
+        map!(x->kappa(κ,x),K,pairwise(metric(κ),transform(κ,X,obsdim),dims=obsdim))
 end
 
-
 """
 ```
     kernelmatrix!(K::Matrix, κ::Kernel, X::Matrix, Y::Matrix; obsdim::Integer=2)
@@ -34,24 +28,10 @@ function kernelmatrix!(
         Y::AbstractMatrix{T₃};
         obsdim::Int = defaultobs
         ) where {T,T₁,T₂,T₃}
-        #TODO Check dimension consistency
-        _kappamatrix!(κ, pairwise!(K, metric(κ), X, Y, dims=obsdim))
-end
-
-
-function kernelmatrix!(
-        K::Matrix{T₁},
-        κ::Kernel{T},
-        X::AbstractMatrix{T₂};
-        obsdim::Int = defaultobs,
-        symmetrize::Bool = true
-        ) where {T,T₁<:Real,T₂<:Real}
-        #TODO Check dimension consistency
-        _symmetric_kappamatrix!(κ,pairwise!(K, metric(κ), X, dims=obsdim), symmetrize)
+        @assert check_dims(K,X,Y,obsdim) "Dimensions of the target array are not consistent with X and Y"
+        map!(x->kappa(κ,x),K,pairwise(metric(κ),transform(κ,X,obsdim),transform(κ,Y,obsdim),dims=obsdim))
 end
 
-# Convenience Methods ======================================================================
-
 """
 ```
     kernel(κ::Kernel, x, y; obsdim=2)
@@ -60,7 +40,7 @@ Apply the kernel `κ` to ``x`` and ``y`` where ``x`` and ``y`` are vectors or sc
 some subtype of ``Real``.
 """
 function kernel(κ::Kernel{T}, x::Real, y::Real) where {T}
-    kernel(κ, T(x), T(y))
+    kernel(κ, [T(x)], [T(y)])
 end
 
 function kernel(
@@ -69,25 +49,23 @@ function kernel(
         y::AbstractArray{T₂};
         obsdim::Int = defaultobs
     ) where {T,T₁<:Real,T₂<:Real}
-    # TODO Verify dimensions
-    kappa(κ, evaluate(metric(κ),x,y))
+    @assert length(x) == length(y) "x and y don't have the same dimension!"
+    kappa(κ, evaluate(metric(κ),transform(κ,x),transform(κ,y)))
 end
 
 """
 ```
-    kernelmatrix(κ::Kernel, X::Matrix ; obsdim::Int=2, symmetrize::Bool)
+    kernelmatrix(κ::Kernel, X::Matrix ; obsdim::Int=2, symmetrize::Bool=true)
 ```
 Calculate the kernel matrix of `X` with respect to kernel `κ`.
 """
 function kernelmatrix(
-        κ::Kernel{T},
-        X::AbstractMatrix{T₁};
+        κ::Kernel{T,<:Transform},
+        X::AbstractMatrix;
         obsdim::Int = defaultobs,
         symmetrize::Bool = true
-    ) where {T,T₁<:Real}
-    Tₛ = typeof(zero(eltype(X))*zero(T))
-    m = size(X,obsdim)
-    return kernelmatrix!(Matrix{promote_float(T,T₁)}(undef,m,m),κ,X,obsdim=obsdim,symmetrize=symmetrize)
+    ) where {T}
+    K = map(x->kappa(κ,x),pairwise(metric(κ),transform(κ,X,obsdim),dims=obsdim))
 end
 
 """
@@ -102,13 +80,10 @@ function kernelmatrix(
         Y::AbstractMatrix{T₂};
         obsdim=defaultobs
     ) where {T,T₁<:Real,T₂<:Real}
-    Tₛ = typeof(zero(eltype(X))*zero(eltype(Y))*zero(T))
-    m = size(X,obsdim)
-    n = size(Y,obsdim)
-    kernelmatrix!(Matrix{Tₛ}(undef,m,n),κ,X,Y,obsdim=obsdim)
+    K = map(x->kappa(κ,x),pairwise(metric(κ),transform(κ,X,obsdim),transform(κ,Y,obsdim),dims=obsdim))
+    return K
 end
 
-
 """
 ```
     kerneldiagmatrix(κ::Kernel, X::Matrix; obsdim::Int=2)
@@ -117,8 +92,30 @@ Calculate the diagonal matrix of `X` with respect to kernel `κ`
 """
 function kerneldiagmatrix(
         κ::Kernel{T},
-        X::AbstractMatrix{T₁}
+        X::AbstractMatrix{T₁};
+        obsdim::Int = defaultobs
+        ) where {T,T₁}
+        if obsdim == 1
+            [@views kernel(κ,X[i,:],X[i,:]) for i in 1:size(X,obsdim)]
+        elseif obsdim == 2
+            [@views kernel(κ,X[i,:],X[i,:]) for i in 1:size(X,obsdim)]
+        end
+end
+
+function kerneldiagmatrix!(
+        K::AbstractVector{T₁},
+        κ::Kernel{T},
+        X::AbstractMatrix{T₂};
+        obsdim::Int = defaultobs
         ) where {T,T₁,T₂}
-        @error "Not implemented yet"
-        #TODO
+        if obsdim == 1
+            for i in eachindex(K)
+                @inbounds @views K[i] = kernel(κ, X[i,:],X[i,:])
+            end
+        else
+            for i in eachindex(K)
+                @inbounds @views K[i] = kernel(κ,X[:,i],X[:,i])
+            end
+        end
+        return K
 end