Merge pull request #148 from sharanry/sy/add_DiffEqArray_special_case

Add special cases for DiffEqArray constructor

Author: ChrisRackauckas

src/vector_of_array.jl

@@ -39,11 +39,12 @@ 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}) where {T, N} = DiffEqArray{eltype(T), N, typeof(vec), typeof(ts), Nothing, Nothing, Nothing, Nothing}(vec, ts, nothing, nothing, nothing, nothing)
+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) = 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), Nothing, Nothing, Nothing, Nothing}(vec, ts, nothing, nothing, nothing, nothing)
-DiffEqArray(vec::AbstractVector{VT},ts::AbstractVector, syms::Vector{Symbol}, indepsym::Symbol, observed::Function, p) where {T, N, VT<:AbstractArray{T, N}} = DiffEqArray{T, N+1, typeof(vec), typeof(ts), typeof(syms), typeof(indepsym), typeof(observed), typeof(p)}(vec, ts, syms, indepsym, observed, p)
+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}}
+  DiffEqArray{T, N+1, typeof(vec), typeof(ts), typeof(syms), typeof(indepsym), typeof(observed), typeof(p)}(vec, ts, syms, indepsym, observed, p)
+end
 
 # Interface for the linear indexing. This is just a view of the underlying nested structure
 @inline Base.firstindex(VA::AbstractVectorOfArray) = firstindex(VA.u)
@@ -102,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
@@ -149,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)

test/interface_tests.jl

@@ -1,51 +1,83 @@
 using RecursiveArrayTools, Test
 
-recs = [[1, 2, 3], [4, 5, 6], [7, 8, 9]]
-testva = VectorOfArray(recs)
+t = 1:3
+testva = VectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+testda = DiffEqArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]],t)
 
 for (i, elem) in enumerate(testva)
     @test elem == testva[i]
 end
 
+for (i, elem) in enumerate(testda)
+    @test elem == testda[i]
+end
+
 push!(testva, [10, 11, 12])
 @test testva[:, end] == [10, 11, 12]
+push!(testda, [10, 11, 12])
+@test testda[:, end] == [10, 11, 12]
+
 testva2 = copy(testva)
 push!(testva2, [13, 14, 15])
+testda2 = copy(testva)
+push!(testda2, [13, 14, 15])
+
 # make sure we copy when we pass containers
 @test size(testva) == (3, 4)
 @test testva2[:, end] == [13, 14, 15]
+@test size(testda) == (3, 4)
+@test testda2[:, end] == [13, 14, 15]
 
 append!(testva, testva)
 @test testva[1:2, 5:6] == [1 4; 2 5]
+append!(testda, testda)
+@test testda[1:2, 5:6] == [1 4; 2 5]
 
 # Test that adding a array of different dimension makes the array ragged
 push!(testva, [-1, -2, -3, -4])
+push!(testda, [-1, -2, -3, -4])
 #testva #TODO: this screws up printing, try to make a fallback
 @test testva[1:2, 5:6] == [1 4; 2 5] # we just let the indexing happen if it works
-testva[4, 9] # == testva.data[9][4]
+@test testda[1:2, 5:6] == [1 4; 2 5]
+
 @test_throws BoundsError testva[4:5, 5:6]
+@test_throws BoundsError testda[4:5, 5:6]
+
 @test testva[9] == [-1, -2, -3, -4]
 @test testva[end] == [-1, -2, -3, -4]
+@test testda[9] == [-1, -2, -3, -4]
+@test testda[end] == [-1, -2, -3, -4]
 
 # Currently we enforce the general shape, they can just be different lengths, ie we
 # can't do
 # Decide if this is desired, or remove this restriction
 @test_throws MethodError push!(testva, [-1 -2 -3 -4])
 @test_throws MethodError push!(testva, [-1 -2; -3 -4])
+@test_throws MethodError push!(testda, [-1 -2 -3 -4])
+@test_throws MethodError push!(testda, [-1 -2; -3 -4])
 
-# convert array from VectorOfArray
+# convert array from VectorOfArray/DiffEqArray
+t = 1:8
 recs = [rand(10, 7) for i = 1:8]
 testva = VectorOfArray(recs)
+testda = DiffEqArray(recs,t)
 testa = cat(recs...,dims=3)
+
 @test convert(Array,testva) == testa
+@test convert(Array,testda) == testa
 
+t = 1:3
 recs = [[1 2; 3 4], [3 5; 6 7], [8 9; 10 11]]
 testva = VectorOfArray(recs)
+testda = DiffEqArray(recs,t)
+
 @test size(convert(Array,testva)) == (2,2,3)
+@test size(convert(Array,testda)) == (2,2,3)
 
 # create similar VectorOfArray
 recs = [rand(6) for i = 1:4]
 testva = VectorOfArray(recs)
+
 testva2 = similar(testva)
 @test typeof(testva2) == typeof(testva)
 @test size(testva2) == size(testva)
@@ -77,3 +109,6 @@ emptyda = DiffEqArray(Array{Vector{Float64}}([]), Vector{Float64}())
 
 A = VectorOfArray(map(i->rand(2,4),1:7))
 @test map(x->maximum(x),A) isa Vector
+
+DA = DiffEqArray(map(i->rand(2,4),1:7), 1:7)
+@test map(x->maximum(x),DA) isa Vector