Flatten the broadcast before calculating narray and handle edge cases

Project.toml

@@ -1,7 +1,7 @@
 name = "RecursiveArrayTools"
 uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "2.4.1"
+version = "2.4.2"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/vector_of_array.jl

@@ -15,12 +15,12 @@ function Base.Array(VA::AbstractVectorOfArray)
   Array(reshape(reduce(hcat,vecs),size(VA.u[1])...,length(VA.u)))
 end
 
-VectorOfArray(vec::AbstractVector{T}, dims::NTuple{N}) where {T, N} = VectorOfArray{eltype(T), N, typeof(vec)}(vec)
+VectorOfArray(vec::AbstractVector{T}, ::NTuple{N}) where {T, N} = VectorOfArray{eltype(T), N, typeof(vec)}(vec)
 # Assume that the first element is representative of all other elements
 VectorOfArray(vec::AbstractVector) = VectorOfArray(vec, (size(vec[1])..., length(vec)))
 VectorOfArray(vec::AbstractVector{VT}) where {T, N, VT<:AbstractArray{T, N}} = VectorOfArray{T, N+1, typeof(vec)}(vec)
 
-DiffEqArray(vec::AbstractVector{T}, ts, dims::NTuple{N}) where {T, N} = DiffEqArray{eltype(T), N, typeof(vec), typeof(ts)}(vec, ts)
+DiffEqArray(vec::AbstractVector{T}, ts, ::NTuple{N}) where {T, N} = DiffEqArray{eltype(T), N, typeof(vec), typeof(ts)}(vec, ts)
 # Assume that the first element is representative of all other elements
 DiffEqArray(vec::AbstractVector,ts::AbstractVector) = DiffEqArray(vec, ts, (size(vec[1])..., length(vec)))
 DiffEqArray(vec::AbstractVector{VT},ts::AbstractVector) where {T, N, VT<:AbstractArray{T, N}} = DiffEqArray{T, N+1, typeof(vec), typeof(ts)}(vec, ts)
@@ -94,7 +94,11 @@ recursivecopy(VA::VectorOfArray) = VectorOfArray(copy.(VA.u))
 # For DiffEqArray it ignores ts and fills only u
 function Base.fill!(VA::AbstractVectorOfArray, x)
     for i in eachindex(VA)
-        fill!(VA[i], x)
+        if VA[i] isa AbstractArray
+            fill!(VA[i], x)
+        else
+            VA[i] = x
+        end
     end
     return VA
 end
@@ -160,17 +164,22 @@ Broadcast.BroadcastStyle(::VectorOfArrayStyle{M}, ::VectorOfArrayStyle{N}) where
 Broadcast.BroadcastStyle(::Type{<:AbstractVectorOfArray{T,N}}) where {T,N} = VectorOfArrayStyle{N}()
 
 @inline function Base.copy(bc::Broadcast.Broadcasted{<:VectorOfArrayStyle})
+    bc = Broadcast.flatten(bc)
     N = narrays(bc)
-    x = unpack_voa(bc, 1)
     VectorOfArray(map(1:N) do i
         copy(unpack_voa(bc, i))
     end)
 end
 
 @inline function Base.copyto!(dest::AbstractVectorOfArray, bc::Broadcast.Broadcasted{<:VectorOfArrayStyle})
+    bc = Broadcast.flatten(bc)
     N = narrays(bc)
     @inbounds for i in 1:N
-        copyto!(dest[i], unpack_voa(bc, i))
+        if dest[i] isa AbstractArray
+            copyto!(dest[i], unpack_voa(bc, i))
+        else
+            dest[i] = copy(unpack_voa(bc, i))
+        end
     end
     dest
 end

test/basic_indexing.jl

@@ -7,11 +7,11 @@ testva = VectorOfArray(recs)
 
 # broadcast with array
 X = rand(3, 3)
-mulX = testva .* X
-ref = mapreduce((x,y)->x.*y, hcat, testva, eachcol(X))
+mulX = sqrt.(abs.(testva .* X))
+ref = mapreduce((x,y)->sqrt.(abs.(x.*y)), hcat, testva, eachcol(X))
 @test mulX == ref
 fill!(mulX, 0)
-mulX .= testva .* X
+mulX .= sqrt.(abs.(testva .* X))
 @test mulX == ref
 
 t = [1,2,3]
@@ -107,3 +107,14 @@ x .= v .* v
 w = v .+ 1
 @test w isa VectorOfArray
 @test w.u == map(x -> x .+ 1, v.u)
+
+# edges cases
+x = [1, 2, 3, 4, 5, 6, 7, 8, 9]
+testva = DiffEqArray(x, x)
+testvb = DiffEqArray(x, x)
+mulX = sqrt.(abs.(testva .* testvb))
+ref = sqrt.(abs.(x .* x))
+@test mulX == ref
+fill!(mulX, 0)
+mulX .= sqrt.(abs.(testva .* testvb))
+@test mulX == ref