Merge pull request #30 from IsakFalk/isakfalk/nystrom

Nystrom Approximation

Author: theogf

Project.toml

@@ -8,6 +8,7 @@ Distances = "b4f34e82-e78d-54a5-968a-f98e89d6e8f7"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 

src/KernelFunctions.jl

@@ -14,11 +14,14 @@ export KernelSum, KernelProduct
 
 export Transform, SelectTransform, ChainTransform, ScaleTransform, LowRankTransform, IdentityTransform, FunctionTransform
 
+export NystromFact, nystrom
+
 using Compat
 using Distances, LinearAlgebra
 using SpecialFunctions: logabsgamma, besselk
 using ZygoteRules: @adjoint
 using StatsFuns: logtwo
+using StatsBase
 using PDMats: PDMat
 
 const defaultobs = 2
@@ -41,6 +44,7 @@ include("matrix/kernelmatrix.jl")
 include("matrix/kernelpdmat.jl")
 include("kernels/kernelsum.jl")
 include("kernels/kernelproduct.jl")
+include("approximations/nystrom.jl")
 
 include("generic.jl")
 

src/approximations/nystrom.jl

@@ -0,0 +1,102 @@
+# Following the algorithm by William and Seeger, 2001
+# Cs is equivalent to X_mm and C to X_mn
+
+function sampleindex(X::AbstractMatrix, r::Real; obsdim::Integer=defaultobs)
+    0 < r <= 1 || throw(ArgumentError("Sample rate `r` must be in range (0,1]"))
+    n = size(X, obsdim)
+    m = ceil(Int, n*r)
+    S = StatsBase.sample(1:n, m; replace=false, ordered=true)
+    return S
+end
+
+function nystrom_sample(k::Kernel, X::AbstractMatrix, S::Vector{<:Integer}; obsdim::Integer=defaultobs)
+    obsdim ∈ [1, 2] || throw(ArgumentError("`obsdim` should be 1 or 2 (see docs of kernelmatrix))"))
+    Xₘ = obsdim == 1 ? X[S, :] : X[:, S]
+    C = k(Xₘ, X; obsdim=obsdim)
+    Cs = C[:, S]
+    return (C, Cs)
+end
+
+function nystrom_pinv!(Cs::Matrix{T}, tol::T=eps(T)*size(Cs,1)) where {T<:Real}
+    # Compute eigendecomposition of sampled component of K
+    QΛQᵀ = LinearAlgebra.eigen!(LinearAlgebra.Symmetric(Cs))
+
+    # Solve for D = Λ^(-1/2) (pseudo inverse - use tolerance from before factorization)
+    D = QΛQᵀ.values
+    λ_tol = maximum(D)*tol
+
+    for i in eachindex(D)
+        @inbounds D[i] = abs(D[i]) <= λ_tol ? zero(T) : one(T)/sqrt(D[i])
+    end
+
+    # Scale eigenvectors by D
+    Q = QΛQᵀ.vectors
+    QD = LinearAlgebra.rmul!(Q, LinearAlgebra.Diagonal(D))  # Scales column i of Q by D[i]
+
+    # W := (QD)(QD)ᵀ = (QΛQᵀ)^(-1)  (pseudo inverse)
+    W = QD*QD'
+
+    # Symmetrize W
+    return LinearAlgebra.copytri!(W, 'U')
+end
+
+@doc raw"""
+    NystromFact
+
+Type for storing a Nystrom factorization. The factorization contains two fields: `W` and
+`C`, two matrices satisfying:
+```math
+\mathbf{K} \approx \mathbf{C}^{\intercal}\mathbf{W}\mathbf{C}
+```
+"""
+struct NystromFact{T<:Real}
+    W::Matrix{T}
+    C::Matrix{T}
+end
+
+function NystromFact(W::Matrix{<:Real}, C::Matrix{<:Real})
+    T = Base.promote_eltypeof(W, C)
+    return NystromFact(convert(Matrix{T}, W), convert(Matrix{T}, C))
+end
+
+@doc raw"""
+    nystrom(k::Kernel, X::Matrix, S::Vector; obsdim::Int=defaultobs)
+
+Computes a factorization of Nystrom approximation of the square kernel matrix of data
+matrix `X` with respect to kernel `k`. Returns a `NystromFact` struct which stores a
+Nystrom factorization satisfying:
+```math
+\mathbf{K} \approx \mathbf{C}^{\intercal}\mathbf{W}\mathbf{C}
+```
+"""
+function nystrom(k::Kernel, X::AbstractMatrix, S::Vector{<:Integer}; obsdim::Int=defaultobs)
+    C, Cs = nystrom_sample(k, X, S; obsdim=obsdim)
+    W = nystrom_pinv!(Cs)
+    return NystromFact(W, C)
+end
+
+@doc raw"""
+    nystrom(k::Kernel, X::Matrix, r::Real; obsdim::Int=defaultobs)
+
+Computes a factorization of Nystrom approximation of the square kernel matrix of data
+matrix `X` with respect to kernel `k` using a sample ratio of `r`.
+Returns a `NystromFact` struct which stores a Nystrom factorization satisfying:
+```math
+\mathbf{K} \approx \mathbf{C}^{\intercal}\mathbf{W}\mathbf{C}
+```
+"""
+function nystrom(k::Kernel, X::AbstractMatrix, r::Real; obsdim::Int=defaultobs)
+    S = sampleindex(X, r; obsdim=obsdim)
+    return nystrom(k, X, S; obsdim=obsdim)
+end
+
+"""
+    nystrom(CᵀWC::NystromFact)
+
+Compute the approximate kernel matrix based on the Nystrom factorization.
+"""
+function kernelmatrix(CᵀWC::NystromFact{<:Real})
+    W = CᵀWC.W
+    C = CᵀWC.C
+    return C'*W*C
+end

test/runtests.jl

@@ -5,6 +5,7 @@ using Random
 
 @testset "KernelFunctions" begin
 include("test_kernelmatrix.jl")
+include("test_approximations.jl")
 include("test_constructors.jl")
 # include("test_AD.jl")
 include("test_transform.jl")

test/test_approximations.jl

@@ -0,0 +1,15 @@
+using Distances, LinearAlgebra
+using Test
+using KernelFunctions
+
+dims = [10,5]
+X = rand(dims...)
+k = SqExponentialKernel()
+@testset "Kernel Matrix Approximations" begin
+    @testset "Nystrom" begin
+        for obsdim in [1, 2]
+            @test kernelmatrix(k, X; obsdim=obsdim) ≈ kernelmatrix(nystrom(k, X, 1.0; obsdim=obsdim))
+            @test kernelmatrix(k, X; obsdim=obsdim) ≈ kernelmatrix(nystrom(k, X, collect(1:dims[obsdim]); obsdim=obsdim))
+        end
+    end
+end