Merge pull request #99 from SciML/myb/cart

Support CartesianIndices for VectorOfArray

Author: ChrisRackauckas

Project.toml

@@ -1,7 +1,7 @@
 name = "RecursiveArrayTools"
 uuid = "731186ca-8d62-57ce-b412-fbd966d074cd"
 authors = ["Chris Rackauckas <[email protected]>"]
-version = "2.3.3"
+version = "2.3.4"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"

src/array_partition.jl

@@ -142,7 +142,7 @@ end
 @inline Base.firstindex(A::ArrayPartition) = 1
 @inline Base.lastindex(A::ArrayPartition) = length(A)
 
-@inline function Base.getindex(A::ArrayPartition, i::Int)
+Base.@propagate_inbounds function Base.getindex(A::ArrayPartition, i::Int)
   @boundscheck checkbounds(A, i)
   @inbounds for j in 1:length(A.x)
     i -= length(A.x[j])
@@ -157,7 +157,7 @@ end
 
 Return the entry at index `j...` of the `i`th partition of `A`.
 """
-@inline function Base.getindex(A::ArrayPartition, i::Int, j...)
+Base.@propagate_inbounds function Base.getindex(A::ArrayPartition, i::Int, j...)
   @boundscheck 0 < i <= length(A.x) || throw(BoundsError(A.x, i))
   @inbounds b = A.x[i]
   @boundscheck checkbounds(b, j...)
@@ -171,7 +171,7 @@ Return vector with all elements of array partition `A`.
 """
 Base.getindex(A::ArrayPartition{T,S}, ::Colon) where {T,S} = T[a for a in Chain(A.x)]
 
-@inline function Base.setindex!(A::ArrayPartition, v, i::Int)
+Base.@propagate_inbounds function Base.setindex!(A::ArrayPartition, v, i::Int)
   @boundscheck checkbounds(A, i)
   @inbounds for j in 1:length(A.x)
     i -= length(A.x[j])
@@ -187,7 +187,7 @@ end
 
 Set the entry at index `j...` of the `i`th partition of `A` to `v`.
 """
-@inline function Base.setindex!(A::ArrayPartition, v, i::Int, j...)
+Base.@propagate_inbounds function Base.setindex!(A::ArrayPartition, v, i::Int, j...)
   @boundscheck 0 < i <= length(A.x) || throw(BoundsError(A.x, i))
   @inbounds b = A.x[i]
   @boundscheck checkbounds(b, j...)

src/vector_of_array.jl

@@ -34,25 +34,38 @@ DiffEqArray(vec::AbstractVector{VT},ts::AbstractVector) where {T, N, VT<:Abstrac
 @inline Base.IteratorSize(VA::AbstractVectorOfArray) = Base.HasLength()
 # Linear indexing will be over the container elements, not the individual elements
 # unlike an true AbstractArray
-@inline Base.getindex(VA::AbstractVectorOfArray{T, N}, I::Int) where {T, N} = VA.u[I]
-@inline Base.getindex(VA::AbstractVectorOfArray{T, N}, I::Colon) where {T, N} = VA.u[I]
-@inline Base.getindex(VA::AbstractVectorOfArray{T, N}, I::AbstractArray{Int}) where {T, N} = VectorOfArray(VA.u[I])
-@inline Base.getindex(VA::AbstractDiffEqArray{T, N}, I::AbstractArray{Int}) where {T, N} = DiffEqArray(VA.u[I],VA.t[I])
-@inline Base.getindex(VA::AbstractVectorOfArray{T, N}, i::Int,::Colon) where {T, N} = [VA.u[j][i] for j in 1:length(VA)]
-@inline Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, I::Int) where {T, N} = VA.u[I] = v
-@inline Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, I::Colon) where {T, N} = VA.u[I] = v
-@inline Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, I::AbstractArray{Int}) where {T, N} = VA.u[I] = v
-@inline function Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, i::Int,::Colon) where {T, N}
+Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, I::Int) where {T, N} = VA.u[I]
+Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, I::Colon) where {T, N} = VA.u[I]
+Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, I::AbstractArray{Int}) where {T, N} = VectorOfArray(VA.u[I])
+Base.@propagate_inbounds Base.getindex(VA::AbstractDiffEqArray{T, N}, I::AbstractArray{Int}) where {T, N} = DiffEqArray(VA.u[I],VA.t[I])
+Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, i::Int,::Colon) where {T, N} = [VA.u[j][i] for j in 1:length(VA)]
+Base.@propagate_inbounds function Base.getindex(VA::AbstractVectorOfArray{T,N}, ii::CartesianIndex) where {T, N}
+    ti = Tuple(ii)
+    i = first(ti)
+    jj = CartesianIndex(Base.tail(ti))
+    return VA.u[i][jj]
+end
+Base.@propagate_inbounds Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, I::Int) where {T, N} = VA.u[I] = v
+Base.@propagate_inbounds Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, I::Colon) where {T, N} = VA.u[I] = v
+Base.@propagate_inbounds Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, I::AbstractArray{Int}) where {T, N} = VA.u[I] = v
+Base.@propagate_inbounds function Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, i::Int,::Colon) where {T, N}
   for j in 1:length(VA)
     VA.u[j][i] = v[j]
   end
+  return v
+end
+Base.@propagate_inbounds function Base.setindex!(VA::AbstractVectorOfArray{T,N}, x, ii::CartesianIndex) where {T, N}
+    ti = Tuple(ii)
+    i = first(ti)
+    jj = CartesianIndex(Base.tail(ti))
+    return VA.u[i][jj] = x
 end
 
 # Interface for the two dimensional indexing, a more standard AbstractArray interface
 @inline Base.size(VA::AbstractVectorOfArray) = (size(VA.u[1])..., length(VA.u))
-@inline Base.getindex(VA::AbstractVectorOfArray{T, N}, I::Int...) where {T, N} = VA.u[I[end]][Base.front(I)...]
-@inline Base.getindex(VA::AbstractVectorOfArray{T, N}, ::Colon, I::Int) where {T, N} = VA.u[I]
-@inline Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, I::Int...) where {T, N} = VA.u[I[end]][Base.front(I)...] = v
+Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, I::Int...) where {T, N} = VA.u[I[end]][Base.front(I)...]
+Base.@propagate_inbounds Base.getindex(VA::AbstractVectorOfArray{T, N}, ::Colon, I::Int) where {T, N} = VA.u[I]
+Base.@propagate_inbounds Base.setindex!(VA::AbstractVectorOfArray{T, N}, v, I::Int...) where {T, N} = VA.u[I[end]][Base.front(I)...] = v
 
 # The iterator will be over the subarrays of the container, not the individual elements
 # unlike an true AbstractArray

test/basic_indexing.jl

@@ -80,3 +80,8 @@ a = testva .+ rand(3,3)
 recs = [rand(2,2) for i in 1:5]
 testva = VectorOfArray(recs)
 @test Array(testva) isa Array{Float64,3}
+
+v = VectorOfArray([zeros(20), zeros(10,10), zeros(3,3,3)])
+v[CartesianIndex((3, 2, 3, 2))] = 1
+@test v[CartesianIndex((3, 2, 3, 2))] == 1
+@test v.u[3][2, 3, 2] == 1