Merge pull request SciML#543 from JoshuaLampert/fix-similar

Fix `similar` and `fill!` for mixed nested `VectorOfArray`

Author: ChrisRackauckas

src/vector_of_array.jl

@@ -1404,7 +1404,11 @@ function Base.similar(
 end
 
 @inline function Base.similar(VA::VectorOfArray, ::Type{T} = eltype(VA)) where {T}
-    return VectorOfArray(similar.(VA.u, T))
+    if eltype(VA.u) <: Union{AbstractArray, AbstractVectorOfArray}
+        return VectorOfArray(similar.(VA.u, T))
+    else
+        return VectorOfArray(similar(VA.u, T))
+    end
 end
 
 @inline function Base.similar(VA::VectorOfArray, dims::N) where {N <: Number}
@@ -1420,7 +1424,7 @@ end
 # For DiffEqArray it ignores ts and fills only u
 function Base.fill!(VA::AbstractVectorOfArray, x)
     for i in 1:length(VA.u)
-        if VA[:, i] isa AbstractArray
+        if VA[:, i] isa Union{AbstractArray, AbstractVectorOfArray}
             if ArrayInterface.ismutable(VA.u[i]) || VA.u[i] isa AbstractVectorOfArray
                 fill!(VA[:, i], x)
             else

test/utils_test.jl

@@ -141,6 +141,23 @@ end
     @test u1.u[2] isa SVector
 end
 
+@testset "VectorOfArray similar with nested scalar leaves" begin
+    a = VectorOfArray([ones(2), VectorOfArray([1.0, 1.0])])
+    b = similar(a, Float64)
+    @test b isa typeof(a)
+    @test b.u[1] isa Vector{Float64}
+    @test b.u[2] isa typeof(a.u[2])
+    @test b.u[2].u isa Vector{Float64}
+    @test length(b.u[2].u) == 2
+end
+
+@testset "recursivefill! with nested union partitions" begin
+    a = VectorOfArray([ones(2), VectorOfArray([1.0, 1.0])])
+    recursivefill!(a, true)
+    @test a.u[1] == ones(2)
+    @test a.u[2].u == ones(2)
+end
+
 # Test recursivefill! with immutable StaticArrays (issue #461)
 @testset "recursivefill! with immutable StaticArrays (issue #461)" begin
     # Test with only immutable SVectors