Specialize Base.similar for VectorOfArray with multidimensional parent

src/vector_of_array.jl

@@ -156,13 +156,16 @@
     VectorOfArray{T, N + 1, typeof(vec)}(vec)
 end
 
-# allow multi-dimensional arrays as long as they're linearly indexed 
+# allow multi-dimensional arrays as long as they're linearly indexed. 
+# currently restricted to arrays whose elements are all the same type
 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
 
+Base.parent(vec::VectorOfArray) = vec.u
+
 function DiffEqArray(vec::AbstractVector{T},
         ts::AbstractVector,
         ::NTuple{N, Int},
@@ -721,6 +724,18 @@
     VectorOfArray([similar(VA[:, i], T) for i in eachindex(VA.u)])
 end
 
+# for VectorOfArray with multi-dimensional parent arrays of arrays where all elements are the same type
+@inline function Base.similar(vec::VectorOfArray{
+        T, N, AT}) where {T, N, AT <: AbstractArray{<:AbstractArray{T}}}
+    return VectorOfArray(similar(Base.parent(vec)))
+end
+
+# special-case when the multi-dimensional parent array is just an AbstractVector (call the old method)
+@inline function Base.similar(vec::VectorOfArray{
+        T, N, AT}) where {T, N, AT <: AbstractVector{<:AbstractArray{T}}}
+    return Base.similar(vec, eltype(vec))
+end
+
 # fill!
 # For DiffEqArray it ignores ts and fills only u
 function Base.fill!(VA::AbstractVectorOfArray, x)

test/basic_indexing.jl

@@ -250,6 +250,10 @@ foo!(u_matrix)
 foo!(u_vector)
 @test u_matrix ≈ u_vector
 
+# test that, for VectorOfArray with multi-dimensional parent arrays,
+# `similar` preserves the structure of the parent array
+@test typeof(parent(similar(u_matrix))) == typeof(parent(u_matrix))
+
 # test efficiency 
 num_allocs = @allocations foo!(u_matrix)
 @test num_allocs == 0