Merge pull request #357 from jlchan/jc/VectorOfArray_multidim

Add support for `VectorOfArray(AbstractArray{<:AbstractArray})`

Author: ChrisRackauckas

src/vector_of_array.jl

@@ -27,6 +27,12 @@ Additionally, the `convert(Array,VA::AbstractVectorOfArray)` function is provide
 the `VectorOfArray` into a matrix/tensor. Also, `vecarr_to_vectors(VA::AbstractVectorOfArray)`
 returns a vector of the series for each component, that is, `A[i,:]` for each `i`.
 A plot recipe is provided, which plots the `A[i,:]` series.
+
+There is also support for `VectorOfArray` with constructed from multi-dimensional arrays
+```julia
+VectorOfArray(u::AbstractArray{AT}) where {T, N, AT <: AbstractArray{T, N}}
+```
+where `IndexStyle(typeof(u)) isa IndexLinear`. 
 """
 mutable struct VectorOfArray{T, N, A} <: AbstractVectorOfArray{T, N, A}
     u::A # A <: AbstractVector{<: AbstractArray{T, N - 1}}
@@ -150,6 +156,13 @@ function VectorOfArray(vec::AbstractVector{VT}) where {T, N, VT <: AbstractArray
     VectorOfArray{T, N + 1, typeof(vec)}(vec)
 end
 
+# allow multi-dimensional arrays as long as they're linearly indexed 
+function VectorOfArray(array::AbstractArray{AT}) where {T, N, AT <: AbstractArray{T, N}}
+    @assert IndexStyle(typeof(array)) isa IndexLinear
+
+    return VectorOfArray{T, N + 1, typeof(array)}(array)
+end
+
 function DiffEqArray(vec::AbstractVector{T},
         ts::AbstractVector,
         ::NTuple{N, Int},

test/basic_indexing.jl

@@ -232,3 +232,24 @@ mulX .= sqrt.(abs.(testva .* testvb))
 a = ArrayPartition(1:5, 1:6)
 a[1:8]
 a[[1, 3, 8]]
+
+####################################################################
+# test when VectorOfArray is constructed from a linearly indexed 
+# multidimensional array of arrays
+####################################################################
+
+u_matrix = VectorOfArray(fill([1, 2], 2, 3))
+u_vector = VectorOfArray(vec(u_matrix.u))
+
+# test broadcasting 
+function foo!(u)
+    @. u += 2 * u * abs(u)
+    return u
+end
+foo!(u_matrix)
+foo!(u_vector)
+@test u_matrix ≈ u_vector
+
+# test efficiency 
+num_allocs = @allocations foo!(u_matrix)
+@test num_allocs == 0