make nystrom work with AbstractVector

src/approximations/nystrom.jl

@@ -1,23 +1,29 @@
 # 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)
+function sampleindex(X::AbstractVector, r::Real)
     0 < r <= 1 || throw(ArgumentError("Sample rate `r` must be in range (0,1]"))
-    n = size(X, obsdim)
+    n = length(X)
     m = ceil(Int, n * r)
     S = StatsBase.sample(1:n, m; replace=false, ordered=true)
     return S
 end
 
+function sampleindex(X::AbstractMatrix, r::Real; obsdim::Integer=defaultobs)
+    return sampleindex(vec_of_vecs(X; obsdim), r)
+end
+
+function nystrom_sample(k::Kernel, X::AbstractVector, S::Vector{<:Integer})
+    Xₘ = X[S]
+    C = kernelmatrix(k, Xₘ, X)
+    Cs = C[:, S]
+    return (C, Cs)
+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 = kernelmatrix(k, Xₘ, X; obsdim=obsdim)
-    Cs = C[:, S]
-    return (C, Cs)
+    return nystrom_sample(k, vec_of_vecs(X; obsdim), S)
 end
 
 function nystrom_pinv!(Cs::Matrix{T}, tol::T=eps(T) * size(Cs, 1)) where {T<:Real}
@@ -63,38 +69,46 @@ function NystromFact(W::Matrix{<:Real}, C::Matrix{<:Real})
 end
 
 @doc raw"""
-    nystrom(k::Kernel, X::Matrix, S::Vector; obsdim::Int=defaultobs)
+    nystrom(k::Kernel, X::Vector, S::Vector)
 
-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:
+Computes a factorization of Nystrom approximation of the square kernel matrix
+of data vector `X` with respect to kernel `k`, using indices `S`.
+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)
+function nystrom(k::Kernel, X::AbstractVector, S::Vector{<:Integer})
+    C, Cs = nystrom_sample(k, X, S)
     W = nystrom_pinv!(Cs)
     return NystromFact(W, C)
 end
 
 @doc raw"""
-    nystrom(k::Kernel, X::Matrix, r::Real; obsdim::Int=defaultobs)
+    nystrom(k::Kernel, X::Vector, r::Real)
 
-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`.
+Computes a factorization of Nystrom approximation of the square kernel matrix
+of data vector `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::AbstractVector, r::Real)
+    S = sampleindex(X, r)
+    return nystrom(k, X, S)
+end
+
+function nystrom(k::Kernel, X::AbstractMatrix, S::Vector{<:Integer}; obsdim::Int=defaultobs)
+    return nystrom(k, vec_of_vecs(X; obsdim), S)
+end
+
 function nystrom(k::Kernel, X::AbstractMatrix, r::Real; obsdim::Int=defaultobs)
-    S = sampleindex(X, r; obsdim=obsdim)
-    return nystrom(k, X, S; obsdim=obsdim)
+    return nystrom(k, vec_of_vecs(X; obsdim), r)
 end
 
 """
-    nystrom(CᵀWC::NystromFact)
+    kernelmatrix(CᵀWC::NystromFact)
 
 Compute the approximate kernel matrix based on the Nystrom factorization.
 """

test/approximations/nystrom.jl

@@ -2,9 +2,15 @@
     dims = [10, 5]
     X = rand(dims...)
     k = SqExponentialKernel()
+    for obsdim in [1, 2]
+        Xv = vec_of_vecs(X; obsdim)
+        @assert Xv isa Union{ColVecs,RowVecs}
+        @test kernelmatrix(k, Xv) ≈ kernelmatrix(nystrom(k, Xv, 1.0))
+        @test kernelmatrix(k, Xv) ≈ kernelmatrix(nystrom(k, Xv, collect(1:dims[obsdim])))
+    end
     for obsdim in [1, 2]
         @test kernelmatrix(k, X; obsdim=obsdim) ≈
-              kernelmatrix(nystrom(k, X, 1.0; 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