Added ColVecs

Author: theogf

src/transform/transform.jl

@@ -19,7 +19,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 = 1))
 
 ### TODO Maybe defining adjoints could help but so far it's not working
 

src/utils.jl

@@ -10,6 +10,30 @@ macro check_args(K, param, cond, desc=string(cond))
 end
 
 
+"""
+    ColVecs{T, TX<:AbstractMatrix}
+
+A lightweight box for an `AbstractMatrix` to make it behave like a vector of vectors.
+"""
+struct ColVecs{T, TX<:AbstractMatrix{T}} <: AbstractVector{Vector{T}}
+    X::TX
+    ColVecs(X::TX) where {T, TX<:AbstractMatrix{T}} = new{T, TX}(X)
+end
+
+Base.:(==)(D1::ColVecs, D2::ColVecs) = D1.X == D2.X
+Base.size(D::ColVecs) = (size(D.X, 2),)
+Base.length(D::ColVecs) = size(D.X, 2)
+Base.getindex(D::ColVecs, n::Int) = D.X[:, n]
+Base.getindex(D::ColVecs, n::CartesianIndex{1}) = getindex(D, n[1])
+Base.getindex(D::ColVecs, n) = ColVecs(D.X[:, n])
+Base.view(D::ColVecs, n::Int) = view(D.X, :, n)
+Base.view(D::ColVecs, n) = ColVecs(view(D.X, :, n))
+Base.eltype(D::ColVecs{T}) where T = Vector{T}
+Base.zero(D::ColVecs) = ColVecs(zero(D.X))
+Base.iterate(D::ColVecs) = (view(D.X, :, 1), 2)
+Base.iterate(D::ColVecs, state) = state > length(D) ? nothing : (view(D.X, :, state), state + 1)
+
+
 # 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,33 @@
 @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 Vector
+        @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 view(DX, 4) isa AbstractVector
+        @test view(DX, 4) == view(X, :, 4)
+        @test view(DX, 1:2:4) isa ColVecs
+        @test view(DX, 1:2:4) == ColVecs(view(X, :, 1:2:4))
+        @test eltype(DX) == Vector{Float64}
+        @test eachindex(DX) == 1:N
+
+        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