Add basic indexing tests for DiffEqArray

Author: sharanry

src/vector_of_array.jl

@@ -39,7 +39,7 @@ VectorOfArray(vec::AbstractVector{T}, ::NTuple{N}) where {T, N} = VectorOfArray{
 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, ::NTuple{N}, syms, indepsym, observed, p) where {T, N} = DiffEqArray{eltype(T), N, typeof(vec), typeof(ts), typeof(syms), typeof(indepsym), typeof(observed), typeof(p)}(vec, ts, syms, indepsym, observed, p)
+DiffEqArray(vec::AbstractVector{T}, ts, ::NTuple{N}, syms=nothing, indepsym=nothing, observed=nothing, p=nothing) where {T, N} = DiffEqArray{eltype(T), N, typeof(vec), typeof(ts), typeof(syms), typeof(indepsym), typeof(observed), typeof(p)}(vec, ts, syms, indepsym, observed, p)
 # Assume that the first element is representative of all other elements
 DiffEqArray(vec::AbstractVector,ts::AbstractVector, syms=nothing, indepsym=nothing, observed=nothing, p=nothing) = DiffEqArray(vec, ts, (size(vec[1])..., length(vec)), syms, indepsym, observed, p)
 function DiffEqArray(vec::AbstractVector{VT},ts::AbstractVector, syms=nothing, indepsym=nothing, observed=nothing, p=nothing) where {T, N, VT<:AbstractArray{T, N}}
@@ -103,6 +103,13 @@ Base.@propagate_inbounds function Base.getindex(A::AbstractDiffEqArray{T, N},sym
 end
 Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, I::Int...) where {T, N} = A.u[I[end]][Base.front(I)...]
 Base.@propagate_inbounds Base.getindex(A::AbstractDiffEqArray{T, N}, i::Int) where {T, N} = A.u[i]
+Base.@propagate_inbounds function Base.getindex(VA::AbstractDiffEqArray{T,N}, ii::CartesianIndex) where {T, N}
+  ti = Tuple(ii)
+  i = last(ti)
+  jj = CartesianIndex(Base.front(ti))
+  return VA.u[i][jj]
+end
+
 function observed(A::AbstractDiffEqArray{T, N},sym,i::Int) where {T, N}
     A.observed(sym,A.u[i],A.p,A.t[i])
 end
@@ -150,6 +157,14 @@ end
 tuples(VA::DiffEqArray) = tuple.(VA.t,VA.u)
 
 # Growing the array simply adds to the container vector
+Base.copy(VA::AbstractDiffEqArray) = typeof(VA)(
+    copy(VA.u),
+    copy(VA.t),
+    (VA.syms===nothing) ? nothing : copy(VA.syms),
+    (VA.indepsym===nothing) ? nothing : copy(VA.indepsym),
+    (VA.observed===nothing) ? nothing : copy(VA.observed),
+    (VA.p===nothing) ? nothing : copy(VA.p)
+  )
 Base.copy(VA::AbstractVectorOfArray) = typeof(VA)(copy(VA.u))
 Base.sizehint!(VA::AbstractVectorOfArray{T, N}, i) where {T, N} = sizehint!(VA.u, i)
 Base.push!(VA::AbstractVectorOfArray{T, N}, new_item::AbstractVector) where {T, N} = push!(VA.u, new_item)

test/basic_indexing.jl

@@ -15,8 +15,6 @@ fill!(mulX, 0)
 mulX .= sqrt.(abs.(testva .* X))
 @test mulX == ref
 
-t = [1,2,3]
-diffeq = DiffEqArray(recs,t)
 @test Array(testva) == [1 4 7
                         2 5 8
                         3 6 9]
@@ -25,7 +23,15 @@ diffeq = DiffEqArray(recs,t)
 @test testva[1:2, 1:2] == [1 4; 2 5]
 @test testa[1:2, 1:2] == [1 4; 2 5]
 
+t = [1,2,3]
+diffeq = DiffEqArray(recs,t)
+@test Array(diffeq) == [1 4 7
+                        2 5 8
+                        3 6 9]
+@test diffeq[1:2, 1:2] == [1 4; 2 5]
+
 # # ndims == 2
+t = 1:10
 recs = [rand(8) for i in 1:10]
 testa = cat(recs...,dims=2)
 testva = VectorOfArray(recs)
@@ -36,63 +42,100 @@ testva = VectorOfArray(recs)
 @test testva[end] == testa[:, end]
 @test testva[2:end] == VectorOfArray([recs[i] for i = 2:length(recs)])
 
+diffeq = DiffEqArray(recs,t)
+@test diffeq[1] == testa[:, 1]
+@test diffeq[:] == recs
+@test diffeq[end] == testa[:, end]
+@test diffeq[2:end] == DiffEqArray([recs[i] for i = 2:length(recs)], t)
+
 # ## (Int, Int)
 @test testa[5, 4] == testva[5, 4]
+@test testa[5, 4] == diffeq[5, 4]
 
 # ## (Int, Range) or (Range, Int)
 @test testa[1, 2:3] == testva[1, 2:3]
 @test testa[5:end, 1] == testva[5:end, 1]
 @test testa[:, 1] == testva[:, 1]
 @test testa[3, :] == testva[3, :]
 
+@test testa[1, 2:3] == diffeq[1, 2:3]
+@test testa[5:end, 1] == diffeq[5:end, 1]
+@test testa[:, 1] == diffeq[:, 1]
+@test testa[3, :] == diffeq[3, :]
+
 # ## (Range, Range)
 @test testa[5:end, 1:2] == testva[5:end, 1:2]
+@test testa[5:end, 1:2] == diffeq[5:end, 1:2]
 
 # # ndims == 3
+t = 1:15
 recs = recs = [rand(10, 8) for i in 1:15]
 testa = cat(recs...,dims=3)
 testva = VectorOfArray(recs)
+diffeq = DiffEqArray(recs,t)
 
 # ## (Int, Int, Int)
 @test testa[1, 7, 14] == testva[1, 7, 14]
+@test testa[1, 7, 14] == diffeq[1, 7, 14]
 
 # ## (Int, Int, Range)
 @test testa[2, 3, 1:2] == testva[2, 3, 1:2]
+@test testa[2, 3, 1:2] == diffeq[2, 3, 1:2]
 
 # ## (Int, Range, Int)
 @test testa[2, 3:4, 1] == testva[2, 3:4, 1]
+@test testa[2, 3:4, 1] == diffeq[2, 3:4, 1]
 
 # ## (Int, Range, Range)
 @test testa[2, 3:4, 1:2] == testva[2, 3:4, 1:2]
+@test testa[2, 3:4, 1:2] == diffeq[2, 3:4, 1:2]
 
 # ## (Range, Int, Range)
 @test testa[2:3, 1, 1:2] == testva[2:3, 1, 1:2]
+@test testa[2:3, 1, 1:2] == diffeq[2:3, 1, 1:2]
 
 # ## (Range, Range, Int)
 @test testa[1:2, 2:3, 1] == testva[1:2, 2:3, 1]
+@test testa[1:2, 2:3, 1] == diffeq[1:2, 2:3, 1]
 
 # ## (Range, Range, Range)
 @test testa[2:3, 2:3, 1:2] == testva[2:3, 2:3, 1:2]
+@test testa[2:3, 2:3, 1:2] == diffeq[2:3, 2:3, 1:2]
 
 # ## Make sure that 1:1 like ranges are not collapsed
 @test testa[1:1, 2:3, 1:2] == testva[1:1, 2:3, 1:2]
+@test testa[1:1, 2:3, 1:2] == diffeq[1:1, 2:3, 1:2]
 
 # ## Test ragged arrays work, or give errors as needed
 #TODO: I am not really sure what the behavior of this is, what does Mathematica do?
+t = 1:3
 recs = [[1, 2, 3], [3, 5, 6, 7], [8, 9, 10, 11]]
 testva = VectorOfArray(recs) #TODO: clearly this printed form is nonsense
+diffeq = DiffEqArray(recs,t)
+
 @test testva[:, 1] == recs[1]
-testva[1:2, 1:2]
+@test testva[1:2, 1:2] == [1 3; 2 5]
+@test diffeq[:, 1] == recs[1]
+@test diffeq[1:2, 1:2] == [1 3; 2 5]
 
+t = 1:5
 recs = [rand(2,2) for i in 1:5]
 testva = VectorOfArray(recs)
+diffeq = DiffEqArray(recs,t)
+
 @test Array(testva) isa Array{Float64,3}
+@test Array(diffeq) isa Array{Float64,3}
 
 v = VectorOfArray([zeros(20), zeros(10,10), zeros(3,3,3)])
 v[CartesianIndex((2, 3, 2, 3))] = 1
 @test v[CartesianIndex((2, 3, 2, 3))] == 1
 @test v.u[3][2, 3, 2] == 1
 
+v = DiffEqArray([zeros(20), zeros(10,10), zeros(3,3,3)], 1:3)
+v[CartesianIndex((2, 3, 2, 3))] = 1
+@test v[CartesianIndex((2, 3, 2, 3))] == 1
+@test v.u[3][2, 3, 2] == 1
+
 v = VectorOfArray([rand(20), rand(10,10), rand(3,3,3)])
 w = v .* v
 @test w isa VectorOfArray
@@ -103,12 +146,25 @@ w = v .* v
 x = copy(v)
 x .= v .* v
 @test x.u == w.u
-
-# broadcast with number
 w = v .+ 1
 @test w isa VectorOfArray
 @test w.u == map(x -> x .+ 1, v.u)
 
+
+v = DiffEqArray([rand(20), rand(10,10), rand(3,3,3)], 1:3)
+w = v .* v
+@test_broken w isa DiffEqArray # FIXME
+@test w[1] isa Vector
+@test w[1] == v[1] .* v[1]
+@test w[2] == v[2] .* v[2]
+@test w[3] == v[3] .* v[3]
+x = copy(v)
+x .= v .* v
+@test x.u == w.u
+w = v .+ 1
+@test_broken w isa DiffEqArray # FIXME
+@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)