Merge pull request JuliaGaussianProcesses#84 from theogf/colvecs

Added ColVecs

Author: theogf

src/transform/transform.jl

@@ -21,7 +21,9 @@ Return exactly the input
 """
 struct IdentityTransform <: Transform end
 
-apply(t::IdentityTransform, x; obsdim::Int=defaultobs) = x
+apply(t::IdentityTransform, x; obsdim::Int = defaultobs) = x
+
+apply(t::Transform, x::ColVecs; obsdim::Int = defaultobs) = ColVecs(apply(t, x.X, obsdim = 2))
 
 ### TODO Maybe defining adjoints could help but so far it's not working
 

src/utils.jl

@@ -10,6 +10,23 @@ macro check_args(K, param, cond, desc=string(cond))
 end
 
 
+"""
+    ColVecs(X::AbstractMatrix)
+
+A lightweight box for an `AbstractMatrix` to make it behave like a vector of vectors.
+"""
+struct ColVecs{T, TX<:AbstractMatrix{T}, S} <: AbstractVector{S}
+    X::TX
+    function ColVecs(X::TX) where {T, TX<:AbstractMatrix{T}}
+        S = typeof(view(X, :, 1))
+        new{T, TX, S}(X)
+    end
+end
+
+Base.size(D::ColVecs) = (size(D.X, 2),)
+Base.getindex(D::ColVecs, i::Int) = view(D.X, :, i)
+Base.getindex(D::ColVecs, i) = ColVecs(view(D.X, :, i))
+
 # Take highest Float among possibilities
 # function promote_float(Tₖ::DataType...)
 #     if length(Tₖ) == 0

src/zygote_adjoints.jl

@@ -4,6 +4,15 @@
   end
 end
 
+@adjoint function ColVecs(X::AbstractMatrix)
+    back(Δ::NamedTuple) = (Δ.X,)
+    back(Δ::AbstractMatrix) = (Δ,)
+    function back(Δ::AbstractVector{<:AbstractVector{<:Real}})
+        throw(error("In slow method"))
+    end
+    return ColVecs(X), back
+end
+
 # @adjoint function evaluate(s::Sinus, x::AbstractVector, y::AbstractVector)
 #   d = evaluate(s, x, y)
 #   s = sum(sin.(π*(x-y)))

test/trainable.jl

@@ -1,4 +1,5 @@
 @testset "trainable" begin
+    using Flux: params
     ν = 2.0; c = 3.0; d = 2.0; γ = 2.0; α = 2.5; h = 0.5; r = rand(3)
 
     kc = ConstantKernel(c=c)

test/transform/transform.jl

@@ -3,6 +3,16 @@
     rng = MersenneTwister(123546)
     X = rand(rng, dims...)
     @testset "IdentityTransform" begin
-        @test KernelFunctions.apply(IdentityTransform(),X)==X
+        @test KernelFunctions.apply(IdentityTransform(), X) == X
+    end
+    @testset "ColVecs" begin
+        vX = KernelFunctions.ColVecs(X)
+        t = ARDTransform(rand(dims[1]))
+        @test KernelFunctions.apply(t, vX) ≈ KernelFunctions.ColVecs(KernelFunctions.apply(t, X, obsdim = 2))
+
+        Y = rand(rng, reverse(dims)...)
+        vY = KernelFunctions.ColVecs(Y')
+        t = ARDTransform(rand(dims[1]))
+        @test KernelFunctions.apply(t, vY) ≈ KernelFunctions.ColVecs(KernelFunctions.apply(t, Y, obsdim = 1)')
     end
 end

test/utils.jl

@@ -1,3 +1,30 @@
 @testset "utils" begin
+    using KernelFunctions: ColVecs
+    rng, N, D = MersenneTwister(123456), 10, 2
+    x, X = randn(rng, N), randn(rng, D, N)
 
+    # Test Matrix data sets.
+    @testset "ColVecs" begin
+        DX = ColVecs(X)
+        @test DX == DX
+        @test size(DX) == (N,)
+        @test length(DX) == N
+        @test getindex(DX, 5) isa AbstractVector
+        @test getindex(DX, 5) == X[:, 5]
+        @test getindex(DX, 1:2:6) isa ColVecs
+        @test getindex(DX, 1:2:6) == ColVecs(X[:, 1:2:6])
+        @test getindex(DX, :) == ColVecs(X)
+        @test eachindex(DX) == 1:N
+        @test first(DX) == X[:, 1]
+
+        let
+            @test Zygote.pullback(ColVecs, X)[1] == DX
+            DX, back = Zygote.pullback(ColVecs, X)
+            @test back((X=ones(size(X)),))[1] == ones(size(X))
+
+            @test Zygote.pullback(DX->DX.X, DX)[1] == X
+            X_, back = Zygote.pullback(DX->DX.X, DX)
+            @test back(ones(size(X)))[1].X == ones(size(X))
+        end
+    end
 end