Merge pull request #504 from JoshuaLampert/fix-broadcasting-ragged-multi-d

Fix broadcasting and `view` for ragged `VectorOfArray`

Author: ChrisRackauckas

src/vector_of_array.jl

@@ -168,7 +168,7 @@ 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. 
+# 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
@@ -675,13 +675,19 @@ function Base.view(A::AbstractVectorOfArray{T, N, <:AbstractVector{T}},
 end
 function Base.view(A::AbstractVectorOfArray, I::Vararg{Any, M}) where {M}
     @inline
-    # Special handling for heterogeneous arrays when viewing a single column
-    # The issue is that to_indices uses axes, which is based on the first element's size
-    # For heterogeneous arrays, we need to use the actual size of the specific column
-    if length(I) == 2 && I[1] == Colon() && I[2] isa Int
-        @boundscheck checkbounds(A.u, I[2])
-        # Use the actual size of the specific column instead of relying on axes/to_indices
-        J = (Base.OneTo(length(A.u[I[2]])), I[2])
+    # Generalized handling for heterogeneous arrays when the last index selects a column (Int)
+    # The issue is that `to_indices` uses `axes(A)` which is based on the first element's size.
+    # For heterogeneous arrays, use the actual axes of the specific selected inner array.
+    if length(I) >= 1 && I[end] isa Int
+        i = I[end]
+        @boundscheck checkbounds(A.u, i)
+        frontI = Base.front(I)
+        # Normalize indices against the selected inner array's axes
+        frontJ = to_indices(A.u[i], frontI)
+        # Unalias indices and construct the full index tuple
+        J = (map(j -> Base.unalias(A, j), frontJ)..., i)
+        # Bounds check against the selected inner array to avoid relying on A's axes
+        @boundscheck checkbounds(Bool, A.u[i], frontJ...) || throw(BoundsError(A, I))
         return SubArray(A, J)
     end
     J = map(i -> Base.unalias(A, i), to_indices(A, I))
@@ -711,10 +717,14 @@ function Base.checkbounds(
 end
 function Base.checkbounds(::Type{Bool}, VA::AbstractVectorOfArray, idx...)
     checkbounds(Bool, VA.u, last(idx)) || return false
-    for i in last(idx)
-        checkbounds(Bool, VA.u[i], Base.front(idx)...) || return false
+    if last(idx) isa Int
+        return checkbounds(Bool, VA.u[last(idx)], Base.front(idx)...)
+    else
+        for i in last(idx)
+            checkbounds(Bool, VA.u[i], Base.front(idx)...) || return false
+        end
+        return true
     end
-    return true
 end
 function Base.checkbounds(VA::AbstractVectorOfArray, idx...)
     checkbounds(Bool, VA, idx...) || throw(BoundsError(VA, idx))
@@ -950,13 +960,13 @@ end
 # make vectorofarrays broadcastable so they aren't collected
 Broadcast.broadcastable(x::AbstractVectorOfArray) = x
 
-# recurse through broadcast arguments and return a parent array for 
+# recurse through broadcast arguments and return a parent array for
 # the first VoA or DiffEqArray in the bc arguments
 function find_VoA_parent(args)
     arg = Base.first(args)
     if arg isa AbstractDiffEqArray
-        # if first(args) is a DiffEqArray, use the underlying 
-        # field `u` of DiffEqArray as a parent array. 
+        # if first(args) is a DiffEqArray, use the underlying
+        # field `u` of DiffEqArray as a parent array.
         return arg.u
     elseif arg isa AbstractVectorOfArray
         return parent(arg)
@@ -975,7 +985,7 @@ end
         map(1:N) do i
             copy(unpack_voa(bc, i))
         end
-    else # if parent isa AbstractArray            
+    else # if parent isa AbstractArray
         map(enumerate(Iterators.product(axes(parent)...))) do (i, _)
             copy(unpack_voa(bc, i))
         end

test/basic_indexing.jl

@@ -162,6 +162,38 @@ f2 = VectorOfArray([[1.0, 2.0], [3.0]])
 @test collect(view(f2, :, 1)) == f2[:, 1]
 @test collect(view(f2, :, 2)) == f2[:, 2]
 
+# Broadcasting of heterogeneous arrays (issue #454)
+u = VectorOfArray([[1.0], [2.0, 3.0]])
+@test length(view(u, :, 1)) == 1
+@test length(view(u, :, 2)) == 2
+# broadcast assignment into selected column (last index Int)
+u[:, 2] .= [10.0, 11.0]
+@test u.u[2] == [10.0, 11.0]
+
+# 2D inner arrays (matrices) with ragged second dimension
+u = VectorOfArray([zeros(1, n) for n in (2, 3)])
+@test length(view(u, 1, :, 1)) == 2
+@test length(view(u, 1, :, 2)) == 3
+u[1, :, 2] .= [1.0, 2.0, 3.0]
+@test u.u[2] == [1.0 2.0 3.0]
+# partial column selection by indices
+u[1, [1, 3], 2] .= [7.0, 9.0]
+@test u.u[2] == [7.0 2.0 9.0]
+
+# 3D inner arrays (tensors) with ragged third dimension
+u = VectorOfArray([zeros(2, 1, n) for n in (2, 3)])
+@test size(view(u, :, :, :, 1)) == (2, 1, 2)
+@test size(view(u, :, :, :, 2)) == (2, 1, 3)
+# assign into a slice of the second inner array using last index Int
+u[2, 1, :, 2] .= [7.0, 8.0, 9.0]
+@test vec(u.u[2][2, 1, :]) == [7.0, 8.0, 9.0]
+# check mixed slicing with range on front dims
+u[1:2, 1, [1, 3], 2] .= [1.0 3.0; 2.0 4.0]
+@test u.u[2][1, 1, 1] == 1.0
+@test u.u[2][2, 1, 1] == 2.0
+@test u.u[2][1, 1, 3] == 3.0
+@test u.u[2][2, 1, 3] == 4.0
+
 # Test that views can be modified
 f3 = VectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]])
 v = view(f3, :, 2)
@@ -259,14 +291,14 @@ a[1:8]
 a[[1, 3, 8]]
 
 ####################################################################
-# test when VectorOfArray is constructed from a linearly indexed 
+# test when VectorOfArray is constructed from a linearly indexed
 # multidimensional array of arrays
 ####################################################################
 
 u_matrix = VectorOfArray([[1, 2] for i in 1:2, j in 1:3])
 u_vector = VectorOfArray([[1, 2] for i in 1:6])
 
-# test broadcasting 
+# test broadcasting
 function foo!(u)
     @. u += 2 * u * abs(u)
     return u
@@ -281,7 +313,7 @@ foo!(u_vector)
 @test typeof(parent(similar(u_matrix))) == typeof(parent(u_matrix))
 @test typeof(parent((x -> x).(u_matrix))) == typeof(parent(u_matrix))
 
-# test efficiency 
+# test efficiency
 num_allocs = @allocations foo!(u_matrix)
 @test num_allocs == 0