Merge pull request #2 from gabrielgellner/vector_of_array

Add code that gives an nice Array interface to a vector of arrays

Author: ChrisRackauckas

README.md

@@ -8,6 +8,34 @@
 RecursiveArrayTools.jl is a set of tools for dealing with recursive arrays like
 arrays of arrays. The current functionality includes:
 
+### Types
+
+```julia
+ArrayPartition(x::AbstractArray...)
+```
+
+An `ArrayPartition` `A` is an array which is made up of different arrays `A.x`.
+These index like a single array, but each subarray may have a different type.
+However, broadcast is overloaded to loop in an efficient manner, meaning that
+`A .+= 2.+B` is type-stable in its computations, even if `A.x[i]` and `A.x[j]`
+do not match types. A full array interface is included for completeness, which
+allows this array type to be used in place of a standard array in places where
+such a type stable broadcast may be needed. One example is in heterogeneous
+differential equations for [DifferentialEquations.jl](https://github.com/JuliaDiffEq/DifferentialEquations.jl).
+
+An `ArrayPartition` acts like a single array. `A[i]` indexes through the first
+array, then the second, etc. all linearly. But `A.x` is where the arrays are stored.
+Thus for
+
+```julia
+using RecursiveArrayTools
+A = ArrayPartition(y,z)
+```
+
+We would have `A.x[1]==y` and `A.x[2]==z`. Broadcasting like `f.(A)` is efficient.
+
+### Functions
+
 ```julia
 recursivecopy!(b::Array{T,N},a::Array{T,N})
 ```
@@ -33,3 +61,16 @@ copyat_or_push!{T}(a::AbstractVector{T},i::Int,x)
 
 If `i<length(x)`, it's simply a `recursivecopy!` to the `i`th element. Otherwise it will
 `push!` a `deepcopy`.
+
+```julia
+recursive_one(a)
+```
+
+Calls `one` on the bottom container to get the "true element one type"
+
+```julia
+mean{T<:AbstractArray}(vecvec::Vector{T})
+mean{T<:AbstractArray}(matarr::Matrix{T},region=0)
+```
+
+Generalized mean functions for vectors of arrays and matrix of arrays.

src/RecursiveArrayTools.jl

@@ -7,8 +7,11 @@ module RecursiveArrayTools
   import Base: mean
 
   include("utils.jl")
+  include("vector_of_array.jl")
   include("array_partition.jl")
 
+  export VectorOfArray, vecarr_to_arr
+
   export recursivecopy!, vecvecapply, copyat_or_push!, vecvec_to_mat, recursive_one,mean
 
   export ArrayPartition

src/array_partition.jl

@@ -11,14 +11,45 @@ function ArrayPartition{T}(x,::Type{Val{T}}=Val{false})
 end
 Base.similar(A::ArrayPartition) = ArrayPartition((similar(x) for x in A.x)...)
 Base.similar(A::ArrayPartition,dims::Tuple) = ArrayPartition((similar(x,dim) for (x,dim) in zip(A.x,dims))...)
+Base.similar(A::ArrayPartition,T,dims::Tuple) = ArrayPartition((similar(x,T,dim) for (x,dim) in zip(A.x,dims))...)
 Base.copy(A::ArrayPartition) = Base.similar(A)
 Base.zeros(A::ArrayPartition) = ArrayPartition((zeros(x) for x in A.x)...)
 
+# Special to work with units
+function Base.ones(A::ArrayPartition)
+  B = similar(A::ArrayPartition)
+  for i in eachindex(A.x)
+    B.x[i] .= eltype(A.x[i])(one(first(A.x[i])))
+  end
+  B
+end
+
+Base.:+(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .+ y for (x,y) in zip(A.x,B.x))...)
+Base.:+(A::Number, B::ArrayPartition) = ArrayPartition((A .+ x for x in B.x)...)
+Base.:+(A::ArrayPartition, B::Number) = ArrayPartition((B .+ x for x in A.x)...)
+Base.:-(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .- y for (x,y) in zip(A.x,B.x))...)
+Base.:-(A::Number, B::ArrayPartition) = ArrayPartition((A .- x for x in B.x)...)
+Base.:-(A::ArrayPartition, B::Number) = ArrayPartition((x .- B for x in A.x)...)
 Base.:*(A::Number, B::ArrayPartition) = ArrayPartition((A .* x for x in B.x)...)
 Base.:*(A::ArrayPartition, B::Number) = ArrayPartition((x .* B for x in A.x)...)
 Base.:/(A::ArrayPartition, B::Number) = ArrayPartition((x ./ B for x in A.x)...)
 Base.:\(A::Number, B::ArrayPartition) = ArrayPartition((x ./ A for x in B.x)...)
 
+if VERSION < v"0.6-"
+  Base.:.+(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .+ y for (x,y) in zip(A.x,B.x))...)
+  Base.:.+(A::Number, B::ArrayPartition) = ArrayPartition((A .+ x for x in B.x)...)
+  Base.:.+(A::ArrayPartition, B::Number) = ArrayPartition((B .+ x for x in A.x)...)
+  Base.:.-(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .- y for (x,y) in zip(A.x,B.x))...)
+  Base.:.-(A::Number, B::ArrayPartition) = ArrayPartition((A .- x for x in B.x)...)
+  Base.:.-(A::ArrayPartition, B::Number) = ArrayPartition((x .- B for x in A.x)...)
+  Base.:.*(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x .* y for (x,y) in zip(A.x,B.x))...)
+  Base.:.*(A::Number, B::ArrayPartition) = ArrayPartition((A .* x for x in B.x)...)
+  Base.:.*(A::ArrayPartition, B::Number) = ArrayPartition((x .* B for x in A.x)...)
+  Base.:./(A::ArrayPartition, B::ArrayPartition) = ArrayPartition((x ./ y for (x,y) in zip(A.x,B.x))...)
+  Base.:./(A::ArrayPartition, B::Number) = ArrayPartition((x ./ B for x in A.x)...)
+  Base.:.\(A::Number, B::ArrayPartition) = ArrayPartition((x ./ A for x in B.x)...)
+end
+
 @inline function Base.getindex( A::ArrayPartition,i::Int)
   @boundscheck i > length(A) && throw(BoundsError("Index out of bounds"))
   @inbounds for j in 1:length(A.x)
@@ -52,9 +83,9 @@ recursive_one(A::ArrayPartition) = recursive_one(first(A.x))
 Base.zero(A::ArrayPartition) = zero(first(A.x))
 Base.first(A::ArrayPartition) = first(A.x)
 
-Base.start(A::ArrayPartition) = chain(A.x...)
-Base.next(iter::ArrayPartition,state) = next(state,state)
-Base.done(iter::ArrayPartition,state) = done(state,state)
+Base.start(A::ArrayPartition) = start(chain(A.x...))
+Base.next(A::ArrayPartition,state) = next(chain(A.x...),state)
+Base.done(A::ArrayPartition,state) = done(chain(A.x...),state)
 
 Base.length(A::ArrayPartition) = sum((length(x) for x in A.x))
 Base.size(A::ArrayPartition) = (length(A),)
@@ -72,15 +103,12 @@ add_idxs{T<:ArrayPartition}(::Type{T},expr) = :($(expr).x[i])
 end
 
 @generated function Base.broadcast(f,B::Union{Number,ArrayPartition}...)
-  exs = ((add_idxs(B[i],:(B[$i])) for i in eachindex(B))...)
   arr_idx = 0
   for (i,b) in enumerate(B)
     if b <: ArrayPartition
       arr_idx = i
       break
     end
   end
-  :(for i in eachindex(B[$arr_idx].x)
-    broadcast(f,$(exs...))
-  end)
+  :(A = similar(B[$arr_idx]); broadcast!(f,A,B...); A)
 end

src/vector_of_array.jl

@@ -0,0 +1,44 @@
+abstract AbstractVectorOfArray{T, N} <: AbstractArray{T, N}
+
+# Based on code from M. Bauman Stackexchange answer + Gitter discussion
+type VectorOfArray{T, N, A} <: AbstractVectorOfArray{T, N}
+  u::A # A <: AbstractVector{<: AbstractArray{T, N - 1}}
+end
+
+VectorOfArray{T, N}(vec::AbstractVector{T}, dims::NTuple{N}) = VectorOfArray{eltype(T), N, typeof(vec)}(vec)
+# Assume that the first element is representative all all other elements
+VectorOfArray(vec::AbstractVector) = VectorOfArray(vec, (size(vec[1])..., length(vec)))
+
+# Interface for the linear indexing. This is just a view of the underlying nested structure
+@inline Base.endof(VA::AbstractVectorOfArray) = endof(VA.u)
+@inline Base.length(VA::AbstractVectorOfArray) = length(VA.u)
+# Linear indexing will be over the container elements, not the individual elements
+# unlike an true AbstractArray
+@inline Base.getindex{T, N}(VA::AbstractVectorOfArray{T, N}, I::Int) = VA.u[I]
+@inline Base.getindex{T, N}(VA::AbstractVectorOfArray{T, N}, I::Colon) = VA.u[I]
+@inline Base.getindex{T, N}(VA::AbstractVectorOfArray{T, N}, I::AbstractArray{Int}) = VA.u[I]
+
+# Interface for the two dimensional indexing, a more standard AbstractArray interface
+@inline Base.size(VA::AbstractVectorOfArray) = (size(VA.u[1])..., length(VA.u))
+@inline Base.getindex{T, N}(VA::AbstractVectorOfArray{T, N}, I::Vararg{Int, N}) = VA.u[I[end]][Base.front(I)...]
+
+# The iterator will be over the subarrays of the container, not the individual elements
+# unlike an true AbstractArray
+Base.start{T, N}(VA::AbstractVectorOfArray{T, N}) = 1
+Base.next{T, N}(VA::AbstractVectorOfArray{T, N}, state) = (VA[state], state + 1)
+Base.done{T, N}(VA::AbstractVectorOfArray{T, N}, state) = state >= length(VA.u) + 1
+
+# Growing the array simply adds to the container vector
+Base.copy(VA::AbstractVectorOfArray) = typeof(VA)(copy(VA.u))
+Base.sizehint!{T, N}(VA::AbstractVectorOfArray{T, N}, i) = sizehint!(VA.u, i)
+Base.push!{T, N}(VA::AbstractVectorOfArray{T, N}, new_item::AbstractVector) = push!(VA.u, new_item)
+
+function Base.append!{T, N}(VA::AbstractVectorOfArray{T, N}, new_item::AbstractVectorOfArray{T, N})
+    for item in copy(new_item)
+        push!(VA, item)
+    end
+    return VA
+end
+
+# conversion tools
+vecarr_to_arr(VA::AbstractVectorOfArray) = cat(length(size(VA)), VA.u...)

test/basic_indexing.jl

@@ -0,0 +1,67 @@
+
+# Example Problem
+recs = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+testa = cat(2, recs...)
+testva = VectorOfArray(recs)
+testa[1:2, 1:2] == [1 4; 2 5]
+testva[1:2, 1:2] == [1 4; 2 5]
+testa[1:2, 1:2] == [1 4; 2 5]
+
+# # ndims == 2
+recs = [rand(8) for i in 1:10]
+testa = cat(2, recs...)
+testva = VectorOfArray(recs)
+
+# ## Linear indexing
+@test testva[1] == testa[:, 1]
+@test testva[:] == recs
+@test testva[end] == testa[:, end]
+@test testva[2:end] == [recs[i] for i = 2:length(recs)]
+
+# ## (Int, Int)
+@test testa[5, 4] == testva[5, 4]
+
+# ## (Int, Range) or (Range, Int)
+@test testa[1, 2:3] == testva[1, 2:3]
+@test testa[5:end, 1] == testva[5:end, 1]
+@test testa[:, 1] == testva[:, 1]
+@test testa[3, :] == testva[3, :]
+
+# ## (Range, Range)
+@test testa[5:end, 1:2] == testva[5:end, 1:2]
+
+# # ndims == 3
+recs = recs = [rand(10, 8) for i in 1:15]
+testa = cat(3, recs...)
+testva = VectorOfArray(recs)
+
+# ## (Int, Int, Int)
+@test testa[1, 7, 14] == testva[1, 7, 14]
+
+# ## (Int, Int, Range)
+@test testa[2, 3, 1:2] == testva[2, 3, 1:2]
+
+# ## (Int, Range, Int)
+@test testa[2, 3:4, 1] == testva[2, 3:4, 1]
+
+# ## (Int, Range, Range)
+@test testa[2, 3:4, 1:2] == testva[2, 3:4, 1:2]
+
+# ## (Range, Int, Range)
+@test testa[2:3, 1, 1:2] == testva[2:3, 1, 1:2]
+
+# ## (Range, Range, Int)
+@test testa[1:2, 2:3, 1] == testva[1:2, 2:3, 1]
+
+# ## (Range, Range, Range)
+@test testa[2:3, 2:3, 1:2] == testva[2:3, 2:3, 1:2]
+
+# ## Make sure that 1:1 like ranges are not collapsed
+@test testa[1:1, 2:3, 1:2] == testva[1:1, 2:3, 1:2]
+
+# ## Test ragged arrays work, or give errors as needed
+#TODO: I am not really sure what the behavior of this is, what does Mathematica do?
+recs = [[1, 2, 3], [3 5; 6 7], [8, 9, 10, 11]]
+testva = VectorOfArray(recs) #TODO: clearly this printed form is nonsense
+@test testva[:, 1] == recs[1]
+testva[1:2, 1:2]

test/interface_tests.jl

@@ -0,0 +1,43 @@
+
+recs = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
+testva = VectorOfArray(recs)
+
+for (i, elem) in enumerate(testva)
+    @test elem == testva[i]
+end
+
+push!(testva, [10, 11, 12])
+@test testva[:, end] == [10, 11, 12]
+testva2 = copy(testva)
+push!(testva2, [13, 14, 15])
+# make sure we copy when we pass containers
+@test size(testva) == (3, 4)
+@test testva2[:, end] == [13, 14, 15]
+
+append!(testva, testva)
+@test testva[1:2, 5:6] == [1 4; 2 5]
+
+# Test that adding a array of different dimension makes the array ragged
+push!(testva, [-1, -2, -3, -4])
+#testva #TODO: this screws up printing, try to make a fallback
+@test testva[1:2, 5:6] == [1 4; 2 5] # we just let the indexing happen if it works
+testva[4, 9] # == testva.data[9][4]
+@test_throws BoundsError testva[4:5, 5:6]
+@test testva[9] == [-1, -2, -3, -4]
+@test testva[end] == [-1, -2, -3, -4]
+
+# Currently we enforce the general shape, they can just be different lengths, ie we
+# can't do
+# Decide if this is desired, or remove this restriction
+@test_throws MethodError push!(testva, [-1 -2 -3 -4])
+@test_throws MethodError push!(testva, [-1 -2; -3 -4])
+
+# convert array from VectorOfArray
+recs = [rand(10, 7) for i = 1:8]
+testva = VectorOfArray(recs)
+testa = cat(3, recs...)
+@test vecarr_to_arr(testva) == testa
+
+recs = [[1, 2, 3], [3 5; 6 7], [8, 9, 10, 11]]
+testva = VectorOfArray(recs)
+@test_throws DimensionMismatch vecarr_to_arr(testva)

test/runtests.jl

@@ -5,4 +5,7 @@ using Base.Test
 tic()
 @time @testset "Utils Tests" begin include("utils_test.jl") end
 @time @testset "Partitions Tests" begin include("partitions_test.jl") end
+@time @testset "VecOfArr Indexing Tests" begin include("basic_indexing.jl") end
+@time @testset "VecOfArr Interface Tests" begin include("interface_tests.jl") end
 toc()
+# Test the VectorOfArray code