diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index ac60bb49..c30a05b7 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -28,6 +28,9 @@ jobs: - "pre" group: - "Core" + - "RaggedArrays" + - "ArrayPartitionAnyAll" + - "ShorthandConstructors" - "Downstream" - "nopre" exclude: diff --git a/lib/RecursiveArrayToolsRaggedArrays/Project.toml b/lib/RecursiveArrayToolsRaggedArrays/Project.toml index a6fd0aa6..7a900689 100644 --- a/lib/RecursiveArrayToolsRaggedArrays/Project.toml +++ b/lib/RecursiveArrayToolsRaggedArrays/Project.toml @@ -3,17 +3,26 @@ uuid = "c384ba91-639a-44ca-823a-e1d3691ab84a" version = "1.0.0" [deps] +Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" +StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" [compat] +Adapt = "4" +ArrayInterface = "7" +LinearAlgebra = "1.10" RecursiveArrayTools = "4" +StaticArraysCore = "1.4" SymbolicIndexingInterface = "0.3.35" julia = "1.10" [extras] +SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["SymbolicIndexingInterface", "Test"] +test = ["SparseArrays", "SymbolicIndexingInterface", "Test"] diff --git a/lib/RecursiveArrayToolsRaggedArrays/src/RecursiveArrayToolsRaggedArrays.jl b/lib/RecursiveArrayToolsRaggedArrays/src/RecursiveArrayToolsRaggedArrays.jl index 1faff1d9..009cd044 100644 --- a/lib/RecursiveArrayToolsRaggedArrays/src/RecursiveArrayToolsRaggedArrays.jl +++ b/lib/RecursiveArrayToolsRaggedArrays/src/RecursiveArrayToolsRaggedArrays.jl @@ -1,142 +1,234 @@ -""" - RecursiveArrayToolsRaggedArrays +module RecursiveArrayToolsRaggedArrays + +import RecursiveArrayTools: RecursiveArrayTools, AbstractRaggedVectorOfArray, + AbstractRaggedDiffEqArray, VectorOfArray, DiffEqArray, + AbstractVectorOfArray, AbstractDiffEqArray, AllObserved +using SymbolicIndexingInterface +using SymbolicIndexingInterface: ParameterTimeseriesCollection, ParameterIndexingProxy, + ScalarSymbolic, ArraySymbolic, NotSymbolic, Timeseries, SymbolCache +using Adapt +using ArrayInterface +using StaticArraysCore +using LinearAlgebra: Adjoint -Ragged (non-rectangular) vector-of-array types that preserve the true ragged -structure without zero-padding. These types do **not** subtype `AbstractArray`, -so they avoid the invalidation and AD issues that come with forcing a -rectangular interpretation onto ragged data. +export RaggedVectorOfArray, RaggedDiffEqArray -Separated into its own subpackage because the method overloads (especially -`getindex` with `Colon`) would invalidate hot paths in Base if defined -unconditionally. +# Based on code from M. Bauman Stackexchange answer + Gitter discussion +""" ```julia -using RecursiveArrayToolsRaggedArrays +RaggedVectorOfArray(u::AbstractVector) ``` -# Quick start +A `RaggedVectorOfArray` is an array which has the underlying data structure `Vector{AbstractArray{T}}` +(but, hopefully, concretely typed!). This wrapper over such data structures allows one to lazily +act like it's a higher-dimensional vector, and easily convert it to different forms. The indexing +structure is: ```julia -using RecursiveArrayTools, RecursiveArrayToolsRaggedArrays - -# Create a ragged array -r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) - -r[:, 1] # [1.0, 2.0] — no zero-padding -r[:, 2] # [3.0, 4.0, 5.0] — actual inner array -r[2, 2] # 4.0 - -# Convert from/to VectorOfArray -va = VectorOfArray(r) # rectangular, zero-padded -r2 = RaggedVectorOfArray(va) # back to ragged +A.u[i] # Returns the ith array in the vector of arrays +A[j, i] # Returns the jth component in the ith array +A[j1, ..., jN, i] # Returns the (j1,...,jN) component of the ith array ``` -""" -module RecursiveArrayToolsRaggedArrays - -import RecursiveArrayTools: RecursiveArrayTools, AbstractRaggedVectorOfArray, - AbstractRaggedDiffEqArray, VectorOfArray, DiffEqArray, - AbstractVectorOfArray, AbstractDiffEqArray -using SymbolicIndexingInterface: SymbolicIndexingInterface, SymbolCache, Timeseries, - ParameterTimeseriesCollection, ParameterIndexingProxy, - parameter_values, symbolic_container, - symbolic_type, NotSymbolic, ScalarSymbolic, ArraySymbolic, - is_parameter, is_timeseries_parameter, getu, observed, - variable_symbols, all_variable_symbols - -export RaggedVectorOfArray, RaggedDiffEqArray -# ═══════════════════════════════════════════════════════════════════════════════ -# Concrete types -# ═══════════════════════════════════════════════════════════════════════════════ +which presents itself as a column-major matrix with the columns being the arrays from the vector. +The `AbstractArray` interface is implemented, giving access to `copy`, `push`, `append!`, etc. functions, +which act appropriately. Points to note are: -""" - RaggedVectorOfArray{T, N, A} <: AbstractRaggedVectorOfArray{T, N, A} + - The length is the number of vectors, or `length(A.u)` where `u` is the vector of arrays. + - Iteration follows the linear index and goes over the vectors -A ragged vector-of-arrays that preserves the true shape of each inner array. -Unlike `VectorOfArray`, indexing does **not** zero-pad: `A[:, i]` returns the -`i`-th inner array with its original size. +Additionally, the `convert(Array,VA::AbstractRaggedVectorOfArray)` function is provided, which transforms +the `RaggedVectorOfArray` into a matrix/tensor. Also, `vecarr_to_vectors(VA::AbstractRaggedVectorOfArray)` +returns a vector of the series for each component, that is, `A[i,:]` for each `i`. -# Fields -- `u::A` — the underlying container of arrays +There is also support for `RaggedVectorOfArray` constructed from multi-dimensional arrays -# Constructors ```julia -RaggedVectorOfArray(vec::AbstractVector) -RaggedVectorOfArray(va::AbstractVectorOfArray) +RaggedVectorOfArray(u::AbstractArray{AT}) where {T, N, AT <: AbstractArray{T, N}} ``` + +where `IndexStyle(typeof(u)) isa IndexLinear`. """ mutable struct RaggedVectorOfArray{T, N, A} <: AbstractRaggedVectorOfArray{T, N, A} - u::A + u::A # A <: AbstractArray{<: AbstractArray{T, N - 1}} end +# RaggedVectorOfArray with an added series for time """ - RaggedDiffEqArray{T, N, A, B, F, S, D} <: AbstractRaggedDiffEqArray{T, N, A} - -A ragged diff-eq array with time vector, parameters, and symbolic system. -Like `RaggedVectorOfArray`, indexing preserves the true ragged structure. - -# Fields -- `u::A` — the underlying container of arrays -- `t::B` — time vector -- `p::F` — parameters -- `sys::S` — symbolic system -- `discretes::D` — discrete parameter timeseries -- `interp::I` — interpolation object (default `nothing`) -- `dense::Bool` — whether dense interpolation is available (default `false`) +```julia +RaggedDiffEqArray(u::AbstractVector, t::AbstractVector) +``` + +This is a `RaggedVectorOfArray`, which stores `A.t` that matches `A.u`. This will plot +`(A.t[i],A[i,:])`. The function `tuples(diffeq_arr)` returns tuples of `(t,u)`. + +To construct a RaggedDiffEqArray + +```julia +t = 0.0:0.1:10.0 +f(t) = t - 1 +f2(t) = t^2 +vals = [[f(tval) f2(tval)] for tval in t] +A = RaggedDiffEqArray(vals, t) +A[1, :] # all time periods for f(t) +A.t +``` """ mutable struct RaggedDiffEqArray{ - T, N, A, B, F, S, D <: Union{Nothing, ParameterTimeseriesCollection}, I, - } <: AbstractRaggedDiffEqArray{T, N, A} - u::A + T, N, A, B, F, S, D <: Union{Nothing, ParameterTimeseriesCollection}, + I, DN, + } <: + AbstractRaggedDiffEqArray{T, N, A} + u::A # A <: AbstractVector{<: AbstractArray{T, N - 1}} t::B p::F sys::S discretes::D interp::I - dense::Bool + dense::DN +end + +### Abstract Interface + +function Base.Array( + VA::AbstractRaggedVectorOfArray{ + T, + N, + A, + } + ) where { + T, N, + A <: AbstractVector{ + <:AbstractVector, + }, + } + return reduce(hcat, VA.u) +end +function Base.Array( + VA::AbstractRaggedVectorOfArray{ + T, + N, + A, + } + ) where { + T, N, + A <: + AbstractVector{<:Number}, + } + return VA.u +end +function Base.Matrix( + VA::AbstractRaggedVectorOfArray{ + T, + N, + A, + } + ) where { + T, N, + A <: AbstractVector{ + <:AbstractVector, + }, + } + return reduce(hcat, VA.u) +end +function Base.Matrix( + VA::AbstractRaggedVectorOfArray{ + T, + N, + A, + } + ) where { + T, N, + A <: + AbstractVector{<:Number}, + } + return Matrix(VA.u) +end +function Base.Vector( + VA::AbstractRaggedVectorOfArray{ + T, + N, + A, + } + ) where { + T, N, + A <: AbstractVector{ + <:AbstractVector, + }, + } + return vec(reduce(hcat, VA.u)) +end +function Base.Vector( + VA::AbstractRaggedVectorOfArray{ + T, + N, + A, + } + ) where { + T, N, + A <: + AbstractVector{<:Number}, + } + return VA.u +end +function Base.Array(VA::AbstractRaggedVectorOfArray) + vecs = vec.(VA.u) + return Array(reshape(reduce(hcat, vecs), size(VA.u[1])..., length(VA.u))) +end +function Base.Array{U}(VA::AbstractRaggedVectorOfArray) where {U} + vecs = vec.(VA.u) + return Array(reshape(reduce(hcat, vecs), size(VA.u[1])..., length(VA.u))) +end + +Base.convert(::Type{AbstractArray}, VA::AbstractRaggedVectorOfArray) = stack(VA.u) + +function Adapt.adapt_structure(to, VA::AbstractRaggedVectorOfArray) + return RaggedVectorOfArray(Adapt.adapt.((to,), VA.u)) end -# ═══════════════════════════════════════════════════════════════════════════════ -# Constructors — RaggedVectorOfArray -# ═══════════════════════════════════════════════════════════════════════════════ +function Adapt.adapt_structure(to, VA::AbstractRaggedDiffEqArray) + return RaggedDiffEqArray( + Adapt.adapt.((to,), VA.u), Adapt.adapt(to, VA.t); + interp = VA.interp, dense = VA.dense + ) +end +function RaggedVectorOfArray(vec::AbstractVector{T}, ::NTuple{N}) where {T, N} + return RaggedVectorOfArray{eltype(T), N, typeof(vec)}(vec) +end +# Assume that the first element is representative of all other elements function RaggedVectorOfArray(vec::AbstractVector) T = eltype(vec[1]) N = ndims(vec[1]) - if all( - x -> x isa Union{<:AbstractArray, <:AbstractVectorOfArray, <:AbstractRaggedVectorOfArray}, - vec - ) + if all(x isa Union{<:AbstractArray, <:AbstractRaggedVectorOfArray} for x in vec) A = Vector{Union{typeof.(vec)...}} else A = typeof(vec) end return RaggedVectorOfArray{T, N + 1, A}(vec) end - -function RaggedVectorOfArray( - vec::AbstractVector{VT} - ) where {T, N, VT <: AbstractArray{T, N}} +function RaggedVectorOfArray(vec::AbstractVector{VT}) where {T, N, VT <: AbstractArray{T, N}} return RaggedVectorOfArray{T, N + 1, typeof(vec)}(vec) end -function RaggedVectorOfArray(vec::AbstractVector{T}, ::NTuple{N}) where {T, N} - return RaggedVectorOfArray{eltype(T), N, typeof(vec)}(vec) -end +# allow multi-dimensional arrays as long as they're linearly indexed. +# currently restricted to arrays whose elements are all the same type +function RaggedVectorOfArray(array::AbstractArray{AT}) where {T, N, AT <: AbstractArray{T, N}} + @assert IndexStyle(typeof(array)) isa IndexLinear -# Convert from VectorOfArray -function RaggedVectorOfArray(va::AbstractVectorOfArray{T, N, A}) where {T, N, A} - return RaggedVectorOfArray{T, N, A}(va.u) + return RaggedVectorOfArray{T, N + 1, typeof(array)}(array) end -# ═══════════════════════════════════════════════════════════════════════════════ -# Constructors — RaggedDiffEqArray -# ═══════════════════════════════════════════════════════════════════════════════ +Base.parent(vec::RaggedVectorOfArray) = vec.u +#### 2-argument + +# first element representative function RaggedDiffEqArray( - vec::AbstractVector, ts::AbstractVector; - discretes = nothing, variables = nothing, parameters = nothing, - independent_variables = nothing, interp = nothing, dense = false + vec::AbstractVector, ts::AbstractVector; discretes = nothing, + variables = nothing, parameters = nothing, independent_variables = nothing, + interp = nothing, dense = false ) sys = SymbolCache( something(variables, []), @@ -146,11 +238,27 @@ function RaggedDiffEqArray( _size = size(vec[1]) T = eltype(vec[1]) return RaggedDiffEqArray{ - T, length(_size) + 1, typeof(vec), typeof(ts), - Nothing, typeof(sys), typeof(discretes), typeof(interp), - }(vec, ts, nothing, sys, discretes, interp, dense) + T, + length(_size) + 1, + typeof(vec), + typeof(ts), + Nothing, + typeof(sys), + typeof(discretes), + typeof(interp), + typeof(dense), + }( + vec, + ts, + nothing, + sys, + discretes, + interp, + dense + ) end +# T and N from type function RaggedDiffEqArray( vec::AbstractVector{VT}, ts::AbstractVector; discretes = nothing, variables = nothing, parameters = nothing, @@ -162,15 +270,71 @@ function RaggedDiffEqArray( something(independent_variables, []) ) return RaggedDiffEqArray{ - eltype(eltype(vec)), N + 1, typeof(vec), typeof(ts), - Nothing, typeof(sys), typeof(discretes), typeof(interp), - }(vec, ts, nothing, sys, discretes, interp, dense) + eltype(eltype(vec)), + N + 1, + typeof(vec), + typeof(ts), + Nothing, + typeof(sys), + typeof(discretes), + typeof(interp), + typeof(dense), + }( + vec, + ts, + nothing, + sys, + discretes, + interp, + dense + ) end +#### 3-argument + +# NTuple, T from type function RaggedDiffEqArray( - vec::AbstractVector, ts::AbstractVector, p; - discretes = nothing, variables = nothing, parameters = nothing, - independent_variables = nothing, interp = nothing, dense = false + vec::AbstractVector{T}, ts::AbstractVector, + ::NTuple{N, Int}; discretes = nothing, interp = nothing, dense = false + ) where {T, N} + return RaggedDiffEqArray{ + eltype(T), N, typeof(vec), typeof(ts), Nothing, Nothing, typeof(discretes), + typeof(interp), typeof(dense), + }( + vec, + ts, + nothing, + nothing, + discretes, + interp, + dense + ) +end + +# NTuple parameter +function RaggedDiffEqArray( + vec::AbstractVector{VT}, ts::AbstractVector, p::NTuple{N2, Int}; + discretes = nothing, interp = nothing, dense = false + ) where {T, N, VT <: AbstractArray{T, N}, N2} + return RaggedDiffEqArray{ + eltype(T), N + 1, typeof(vec), typeof(ts), typeof(p), Nothing, typeof(discretes), + typeof(interp), typeof(dense), + }( + vec, + ts, + p, + nothing, + discretes, + interp, + dense + ) +end + +# first element representative +function RaggedDiffEqArray( + vec::AbstractVector, ts::AbstractVector, p; discretes = nothing, + variables = nothing, parameters = nothing, independent_variables = nothing, + interp = nothing, dense = false ) sys = SymbolCache( something(variables, []), @@ -180,11 +344,27 @@ function RaggedDiffEqArray( _size = size(vec[1]) T = eltype(vec[1]) return RaggedDiffEqArray{ - T, length(_size) + 1, typeof(vec), typeof(ts), - typeof(p), typeof(sys), typeof(discretes), typeof(interp), - }(vec, ts, p, sys, discretes, interp, dense) + T, + length(_size) + 1, + typeof(vec), + typeof(ts), + typeof(p), + typeof(sys), + typeof(discretes), + typeof(interp), + typeof(dense), + }( + vec, + ts, + p, + sys, + discretes, + interp, + dense + ) end +# T and N from type function RaggedDiffEqArray( vec::AbstractVector{VT}, ts::AbstractVector, p; discretes = nothing, variables = nothing, parameters = nothing, @@ -197,10 +377,61 @@ function RaggedDiffEqArray( ) return RaggedDiffEqArray{ eltype(T), N + 1, typeof(vec), typeof(ts), - typeof(p), typeof(sys), typeof(discretes), typeof(interp), - }(vec, ts, p, sys, discretes, interp, dense) + typeof(p), typeof(sys), typeof(discretes), + typeof(interp), typeof(dense), + }( + vec, + ts, + p, + sys, + discretes, + interp, + dense + ) +end + +#### 4-argument + +# NTuple, T from type +function RaggedDiffEqArray( + vec::AbstractVector{T}, ts::AbstractVector, + ::NTuple{N, Int}, p; discretes = nothing, interp = nothing, dense = false + ) where {T, N} + return RaggedDiffEqArray{ + eltype(T), N, typeof(vec), typeof(ts), typeof(p), Nothing, typeof(discretes), + typeof(interp), typeof(dense), + }( + vec, + ts, + p, + nothing, + discretes, + interp, + dense + ) +end + +# NTuple parameter +function RaggedDiffEqArray( + vec::AbstractVector{VT}, ts::AbstractVector, p::NTuple{N2, Int}, sys; + discretes = nothing, interp = nothing, dense = false + ) where {T, N, VT <: AbstractArray{T, N}, N2} + return RaggedDiffEqArray{ + eltype(T), N + 1, typeof(vec), typeof(ts), + typeof(p), typeof(sys), typeof(discretes), + typeof(interp), typeof(dense), + }( + vec, + ts, + p, + sys, + discretes, + interp, + dense + ) end +# first element representative function RaggedDiffEqArray( vec::AbstractVector, ts::AbstractVector, p, sys; discretes = nothing, interp = nothing, dense = false @@ -208,340 +439,283 @@ function RaggedDiffEqArray( _size = size(vec[1]) T = eltype(vec[1]) return RaggedDiffEqArray{ - T, length(_size) + 1, typeof(vec), typeof(ts), - typeof(p), typeof(sys), typeof(discretes), typeof(interp), - }(vec, ts, p, sys, discretes, interp, dense) + T, + length(_size) + 1, + typeof(vec), + typeof(ts), + typeof(p), + typeof(sys), + typeof(discretes), + typeof(interp), + typeof(dense), + }( + vec, + ts, + p, + sys, + discretes, + interp, + dense + ) end +# T and N from type function RaggedDiffEqArray( vec::AbstractVector{VT}, ts::AbstractVector, p, sys; discretes = nothing, interp = nothing, dense = false ) where {T, N, VT <: AbstractArray{T, N}} return RaggedDiffEqArray{ eltype(T), N + 1, typeof(vec), typeof(ts), - typeof(p), typeof(sys), typeof(discretes), typeof(interp), - }(vec, ts, p, sys, discretes, interp, dense) -end - -# Convert from DiffEqArray -function RaggedDiffEqArray(da::AbstractDiffEqArray{T, N, A}) where {T, N, A} - _interp = hasproperty(da, :interp) ? da.interp : nothing - _dense = hasproperty(da, :dense) ? da.dense : false - return RaggedDiffEqArray{ - T, N, A, typeof(da.t), typeof(da.p), typeof(da.sys), - typeof(RecursiveArrayTools.get_discretes(da)), typeof(_interp), - }(da.u, da.t, da.p, da.sys, RecursiveArrayTools.get_discretes(da), _interp, _dense) + typeof(p), typeof(sys), typeof(discretes), + typeof(interp), typeof(dense), + }( + vec, + ts, + p, + sys, + discretes, + interp, + dense + ) end -# ═══════════════════════════════════════════════════════════════════════════════ -# Conversion to rectangular VectorOfArray / DiffEqArray -# ═══════════════════════════════════════════════════════════════════════════════ - -function VectorOfArray(r::AbstractRaggedVectorOfArray{T, N, A}) where {T, N, A} - return VectorOfArray{T, N, A}(r.u) -end +#### 5-argument -function DiffEqArray(r::AbstractRaggedDiffEqArray) - return DiffEqArray( - r.u, r.t, r.p, r.sys; discretes = r.discretes, - interp = r.interp, dense = r.dense +# NTuple, T from type +function RaggedDiffEqArray( + vec::AbstractVector{T}, ts::AbstractVector, + ::NTuple{N, Int}, p, sys; discretes = nothing, + interp = nothing, dense = false + ) where {T, N} + return RaggedDiffEqArray{ + eltype(T), N, typeof(vec), typeof(ts), typeof(p), typeof(sys), typeof(discretes), + typeof(interp), typeof(dense), + }( + vec, + ts, + p, + sys, + discretes, + interp, + dense ) end -# ═══════════════════════════════════════════════════════════════════════════════ -# Core interface — NOT AbstractArray -# ═══════════════════════════════════════════════════════════════════════════════ - -Base.parent(r::RaggedVectorOfArray) = r.u - -# length = number of inner arrays (not total elements) -Base.length(r::AbstractRaggedVectorOfArray) = length(r.u) - -# ndims -Base.ndims(::AbstractRaggedVectorOfArray{T, N}) where {T, N} = N - -# eltype: the inner array type, not the scalar type -Base.eltype(::Type{<:AbstractRaggedVectorOfArray{T, N, A}}) where {T, N, A} = eltype(A) -Base.eltype(r::AbstractRaggedVectorOfArray) = eltype(typeof(r)) - -# Iteration yields inner arrays -Base.iterate(r::AbstractRaggedVectorOfArray) = iterate(r.u) -Base.iterate(r::AbstractRaggedVectorOfArray, state) = iterate(r.u, state) - -Base.firstindex(r::AbstractRaggedVectorOfArray) = firstindex(r.u) -Base.lastindex(r::AbstractRaggedVectorOfArray) = lastindex(r.u) - -# lastindex with dimension — needed for `end` in multi-index expressions like r[end, 2] -# dim N (last) = number of inner arrays; dim 1..N-1 = ragged, use max across inner arrays -function Base.lastindex(r::AbstractRaggedVectorOfArray{T, N}, d::Int) where {T, N} - if d == N - return length(r.u) - else - return isempty(r.u) ? 0 : maximum(size(u, d) for u in r.u) - end -end +has_discretes(::TT) where {TT <: AbstractRaggedDiffEqArray} = hasfield(TT, :discretes) +get_discretes(x) = getfield(x, :discretes) -# axes with dimension — needed for `end` translation and range indexing -function Base.axes(r::AbstractRaggedVectorOfArray{T, N}, d::Int) where {T, N} - return Base.OneTo(lastindex(r, d)) +SymbolicIndexingInterface.is_timeseries(::Type{<:AbstractRaggedVectorOfArray}) = Timeseries() +function SymbolicIndexingInterface.is_parameter_timeseries( + ::Type{ + RaggedDiffEqArray{ + T, N, A, B, + F, S, D, I, DN, + }, + } + ) where {T, N, A, B, F, S, D <: ParameterIndexingProxy, I, DN} + return Timeseries() end - -Base.keys(r::AbstractRaggedVectorOfArray) = keys(r.u) -Base.eachindex(r::AbstractRaggedVectorOfArray) = eachindex(r.u) - -# first/last return inner arrays (column-wise), not scalar elements -Base.first(r::AbstractRaggedVectorOfArray) = first(r.u) -Base.last(r::AbstractRaggedVectorOfArray) = last(r.u) - -# ═══════════════════════════════════════════════════════════════════════════════ -# Indexing — no zero-padding -# ═══════════════════════════════════════════════════════════════════════════════ - -# Linear indexing: A[i] uses column-major order matching AbstractArray convention -# For N==1 (vector of scalars), linear indexing matches column indexing -Base.@propagate_inbounds function Base.getindex( - r::RaggedVectorOfArray{T, 1}, i::Int - ) where {T} - return r.u[i] +SymbolicIndexingInterface.state_values(A::AbstractRaggedDiffEqArray) = A.u +SymbolicIndexingInterface.current_time(A::AbstractRaggedDiffEqArray) = A.t +SymbolicIndexingInterface.parameter_values(A::AbstractRaggedDiffEqArray) = A.p +SymbolicIndexingInterface.symbolic_container(A::AbstractRaggedDiffEqArray) = A.sys +function SymbolicIndexingInterface.get_parameter_timeseries_collection(A::AbstractRaggedDiffEqArray) + return get_discretes(A) end -Base.@propagate_inbounds function Base.getindex( - r::RaggedDiffEqArray{T, 1}, i::Int - ) where {T} - return r.u[i] -end - -# For N>1, walk through inner arrays in column-major order -@inline function _ragged_linear_getindex(r, i::Int) - offset = 0 - for col in eachindex(r.u) - n = length(r.u[col]) - if i <= offset + n - return r.u[col][i - offset] - end - offset += n - end - throw(BoundsError(r, i)) +# Callable interface for interpolation +function (A::RaggedDiffEqArray)(t, deriv = Val{0}; idxs = nothing, continuity = :left) + A.interp === nothing && + error("No interpolation data is available. Provide an interpolation object via the `interp` keyword.") + return A.interp(t, idxs, deriv, A.p, continuity) end -@inline function _ragged_linear_setindex!(r, v, i::Int) - offset = 0 - for col in eachindex(r.u) - n = length(r.u[col]) - if i <= offset + n - r.u[col][i - offset] = v - return v - end - offset += n - end - throw(BoundsError(r, i)) -end +Base.IndexStyle(A::AbstractRaggedVectorOfArray) = Base.IndexStyle(typeof(A)) +Base.IndexStyle(::Type{<:AbstractRaggedVectorOfArray}) = IndexCartesian() -Base.@propagate_inbounds function Base.getindex( - r::RaggedVectorOfArray{T, N}, i::Int - ) where {T, N} - return _ragged_linear_getindex(r, i) +@inline Base.length(VA::AbstractRaggedVectorOfArray) = length(VA.u) +@inline function Base.eachindex(VA::AbstractRaggedVectorOfArray) + return eachindex(VA.u) end - -Base.@propagate_inbounds function Base.getindex( - r::RaggedDiffEqArray{T, N}, i::Int +@inline function Base.eachindex( + ::IndexLinear, VA::AbstractRaggedVectorOfArray{T, N, <:AbstractVector{T}} ) where {T, N} - return _ragged_linear_getindex(r, i) + return eachindex(IndexLinear(), VA.u) end - -Base.@propagate_inbounds function Base.setindex!( - r::RaggedVectorOfArray{T, 1}, v, i::Int - ) where {T} - r.u[i] = v - return v +@inline Base.IteratorSize(::Type{<:AbstractRaggedVectorOfArray}) = Base.HasLength() +@inline Base.first(VA::AbstractRaggedVectorOfArray) = first(VA.u) +@inline Base.last(VA::AbstractRaggedVectorOfArray) = last(VA.u) +function Base.firstindex(VA::AbstractRaggedVectorOfArray{T, N, A}) where {T, N, A} + return firstindex(VA.u) end -Base.@propagate_inbounds function Base.setindex!( - r::RaggedDiffEqArray{T, 1}, v, i::Int - ) where {T} - r.u[i] = v - return v +function Base.lastindex(VA::AbstractRaggedVectorOfArray{T, N, A}) where {T, N, A} + return lastindex(VA.u) end -Base.@propagate_inbounds function Base.setindex!( - r::RaggedVectorOfArray{T, N}, v, i::Int - ) where {T, N} - return _ragged_linear_setindex!(r, v, i) +# Always return RaggedEnd for type stability. Use dim=0 to indicate a plain index stored in offset. +# _resolve_ragged_index and _column_indices handle the dim=0 case to extract the actual index value. +@inline function Base.lastindex(VA::AbstractRaggedVectorOfArray, d::Integer) + if d == ndims(VA) + return RaggedEnd(0, Int(lastindex(VA.u))) + elseif d < ndims(VA) + isempty(VA.u) && return RaggedEnd(0, 0) + return RaggedEnd(Int(d), 0) + else + return RaggedEnd(0, 1) + end end -Base.@propagate_inbounds function Base.setindex!( - r::RaggedDiffEqArray{T, N}, v, i::Int - ) where {T, N} - return _ragged_linear_setindex!(r, v, i) -end +Base.getindex(A::AbstractRaggedVectorOfArray, I::Int) = A.u[I] +Base.getindex(A::AbstractRaggedVectorOfArray, I::AbstractArray{Int}) = A.u[I] +Base.getindex(A::AbstractRaggedDiffEqArray, I::Int) = A.u[I] +Base.getindex(A::AbstractRaggedDiffEqArray, I::AbstractArray{Int}) = A.u[I] -# A[j, i] returns the j-th component of the i-th inner array (no zero-padding) -function Base.getindex( - r::RaggedVectorOfArray{T, N}, I::Vararg{Int, N} - ) where {T, N} - col = I[N] - inner_I = Base.front(I) - return r.u[col][inner_I...] -end +__parameterless_type(T) = Base.typename(T).wrapper -function Base.getindex( - r::RaggedDiffEqArray{T, N}, I::Vararg{Int, N} - ) where {T, N} - col = I[N] - inner_I = Base.front(I) - return r.u[col][inner_I...] +# `end` support for ragged inner arrays +# Use runtime fields instead of type parameters for type stability +struct RaggedEnd + dim::Int + offset::Int end +RaggedEnd(dim::Int) = RaggedEnd(dim, 0) -function Base.setindex!( - r::AbstractRaggedVectorOfArray{T, N}, v, I::Vararg{Int, N} - ) where {T, N} - col = I[N] - inner_I = Base.front(I) - r.u[col][inner_I...] = v - return v -end +Base.:+(re::RaggedEnd, n::Integer) = RaggedEnd(re.dim, re.offset + Int(n)) +Base.:-(re::RaggedEnd, n::Integer) = RaggedEnd(re.dim, re.offset - Int(n)) +Base.:+(n::Integer, re::RaggedEnd) = re + n -# A[:, i] returns the i-th inner array directly (no zero-padding) -Base.getindex(r::RaggedVectorOfArray, ::Colon, i::Int) = r.u[i] -Base.setindex!(r::RaggedVectorOfArray, v, ::Colon, i::Int) = (r.u[i] = v) -Base.getindex(r::RaggedDiffEqArray, ::Colon, i::Int) = r.u[i] -Base.setindex!(r::RaggedDiffEqArray, v, ::Colon, i::Int) = (r.u[i] = v) +# Make RaggedEnd and RaggedRange broadcast as scalars to avoid +# issues with collect/length in broadcasting contexts (e.g., SymbolicIndexingInterface) +Base.broadcastable(x::RaggedEnd) = Ref(x) -# A[:, :] returns a copy of the ragged array -function Base.getindex(r::RaggedVectorOfArray, ::Colon, ::Colon) - return RaggedVectorOfArray(copy(r.u)) +struct RaggedRange + dim::Int + start::Int + step::Int + offset::Int end -function Base.getindex(r::RaggedDiffEqArray, ::Colon, ::Colon) - return RaggedDiffEqArray( - copy(r.u), copy(r.t), r.p, r.sys; discretes = r.discretes, - interp = r.interp, dense = r.dense - ) +Base.:(:)(stop::RaggedEnd) = RaggedRange(stop.dim, 1, 1, stop.offset) +function Base.:(:)(start::Integer, stop::RaggedEnd) + return RaggedRange(stop.dim, Int(start), 1, stop.offset) end - -# A[:, idx_array] returns a subset -function Base.getindex( - r::RaggedVectorOfArray, ::Colon, - I::Union{AbstractArray{Int}, AbstractArray{Bool}} - ) - return RaggedVectorOfArray(r.u[I]) +function Base.:(:)(start::Integer, step::Integer, stop::RaggedEnd) + return RaggedRange(stop.dim, Int(start), Int(step), stop.offset) end - -function Base.getindex( - r::RaggedDiffEqArray, ::Colon, - I::Union{AbstractArray{Int}, AbstractArray{Bool}} - ) - return RaggedDiffEqArray( - r.u[I], r.t[I], r.p, r.sys; discretes = r.discretes, - interp = r.interp, dense = r.dense - ) +function Base.:(:)(start::RaggedEnd, stop::RaggedEnd) + return RaggedRange(stop.dim, start.offset, 1, stop.offset) end - -# A[j, :] returns a vector of the j-th component from each inner array -function Base.getindex(r::RaggedVectorOfArray, i::Int, ::Colon) - return [u[i] for u in r.u] +function Base.:(:)(start::RaggedEnd, step::Integer, stop::RaggedEnd) + return RaggedRange(stop.dim, start.offset, Int(step), stop.offset) end -function Base.getindex(r::RaggedDiffEqArray, i::Int, ::Colon) - return [u[i] for u in r.u] +function Base.:(:)(start::RaggedEnd, stop::Integer) + return RaggedRange(start.dim, start.offset, 1, Int(stop)) end - -# Range indexing into inner arrays: r[1:3, i] or r[2:end, i] -function Base.getindex( - r::AbstractRaggedVectorOfArray, I::AbstractRange, col::Int - ) - return r.u[col][I] +function Base.:(:)(start::RaggedEnd, step::Integer, stop::Integer) + return RaggedRange(start.dim, start.offset, Int(step), Int(stop)) end +Base.broadcastable(x::RaggedRange) = Ref(x) -# r[1:end, :] — range of components across all columns -function Base.getindex( - r::AbstractRaggedVectorOfArray, I::AbstractRange, ::Colon - ) - return [u[I] for u in r.u] +@inline function _is_ragged_dim(VA::AbstractRaggedVectorOfArray, d::Integer) + length(VA.u) <= 1 && return false + first_size = size(VA.u[1], d) + @inbounds for idx in 2:length(VA.u) + size(VA.u[idx], d) == first_size || return true + end + return false end -# A[:, range] — column subset with UnitRange (needed for A[:, 2:end]) -function Base.getindex( - r::RaggedVectorOfArray, ::Colon, I::AbstractRange - ) - return RaggedVectorOfArray(r.u[I]) -end -function Base.getindex( - r::RaggedDiffEqArray, ::Colon, I::AbstractRange +Base.@propagate_inbounds function _getindex( + A::AbstractRaggedVectorOfArray, ::NotSymbolic, ::Colon, I::Int ) - return RaggedDiffEqArray( - r.u[I], r.t[I], r.p, r.sys; - discretes = r.discretes, interp = r.interp, dense = r.dense - ) -end - -# ═══════════════════════════════════════════════════════════════════════════════ -# Views -# ═══════════════════════════════════════════════════════════════════════════════ - -function Base.view(r::AbstractRaggedVectorOfArray, ::Colon, i::Int) - return view(r.u[i], :) + return A.u[I] end -function Base.view(r::AbstractRaggedVectorOfArray, I, i::Int) - return view(r.u[i], I) +Base.@propagate_inbounds function _getindex( + A::AbstractRaggedDiffEqArray, ::NotSymbolic, ::Colon, I::Int + ) + return A.u[I] end -# ═══════════════════════════════════════════════════════════════════════════════ -# Symbolic Indexing Dispatch -# ═══════════════════════════════════════════════════════════════════════════════ - -# Main getindex entry point for multi-arg indexing — mirrors AbstractVectorOfArray -Base.@propagate_inbounds function Base.getindex( - r::AbstractRaggedDiffEqArray, _arg, args... +Base.@propagate_inbounds function _getindex( + A::AbstractRaggedVectorOfArray, ::NotSymbolic, + I::Union{Int, AbstractArray{Int}, AbstractArray{Bool}, Colon}... ) - symtype = symbolic_type(_arg) - elsymtype = symbolic_type(eltype(_arg)) - - return if symtype == NotSymbolic() && elsymtype == NotSymbolic() - if _arg isa Union{Tuple, AbstractArray} && - any(x -> symbolic_type(x) != NotSymbolic(), _arg) - _ragged_getindex(r, symtype, elsymtype, _arg, args...) - else - _ragged_getindex(r, symtype, _arg, args...) - end + return if last(I) isa Int + A.u[last(I)][Base.front(I)...] else - _ragged_getindex(r, symtype, elsymtype, _arg, args...) + stack(getindex.(A.u[last(I)], tuple.(Base.front(I))...)) end end -# Non-symbolic multi-arg: Colon + Int already handled above, handle remaining cases -Base.@propagate_inbounds function _ragged_getindex( - r::AbstractRaggedDiffEqArray, ::NotSymbolic, +Base.@propagate_inbounds function _getindex( + A::AbstractRaggedDiffEqArray, ::NotSymbolic, I::Union{Int, AbstractArray{Int}, AbstractArray{Bool}, Colon}... ) - if last(I) isa Int - return r.u[last(I)][Base.front(I)...] + return if last(I) isa Int + A.u[last(I)][Base.front(I)...] else - col_idxs = last(I) isa Colon ? eachindex(r.u) : last(I) + col_idxs = last(I) + # Only preserve RaggedDiffEqArray type if all prefix indices are Colons (selecting whole inner arrays) if all(idx -> idx isa Colon, Base.front(I)) - u_slice = [r.u[col][Base.front(I)...] for col in col_idxs] + # For Colon, select all columns + if col_idxs isa Colon + col_idxs = eachindex(A.u) + end + # For RaggedDiffEqArray, we need to preserve the time values and type + # Create a vector of sliced arrays instead of stacking into higher-dim array + u_slice = [A.u[col][Base.front(I)...] for col in col_idxs] + # Return as RaggedDiffEqArray with sliced time values return RaggedDiffEqArray( - u_slice, r.t[col_idxs], r.p, r.sys; - discretes = r.discretes, interp = r.interp, dense = r.dense + u_slice, A.t[col_idxs], parameter_values(A), symbolic_container(A); + interp = A.interp, dense = A.dense ) else - return [r.u[col][Base.front(I)...] for col in col_idxs] + # Prefix indices are not all Colons - do the same as RaggedVectorOfArray + # (stack the results into a higher-dimensional array) + return stack(getindex.(A.u[col_idxs], tuple.(Base.front(I))...)) end end end +Base.@propagate_inbounds function _getindex( + VA::AbstractRaggedVectorOfArray, ::NotSymbolic, ii::CartesianIndex + ) + ti = Tuple(ii) + i = last(ti) + jj = CartesianIndex(Base.front(ti)) + return VA.u[i][jj] +end + +Base.@propagate_inbounds function _getindex( + A::AbstractRaggedVectorOfArray, ::NotSymbolic, ::Colon, + I::Union{AbstractArray{Int}, AbstractArray{Bool}} + ) + return RaggedVectorOfArray(A.u[I]) +end -# ParameterIndexingError — reuse from RecursiveArrayTools -struct RaggedParameterIndexingError +Base.@propagate_inbounds function _getindex( + A::AbstractRaggedDiffEqArray, ::NotSymbolic, ::Colon, + I::Union{AbstractArray{Int}, AbstractArray{Bool}} + ) + return RaggedDiffEqArray( + A.u[I], A.t[I], parameter_values(A), symbolic_container(A); + interp = A.interp, dense = A.dense + ) +end + +struct ParameterIndexingError <: Exception sym::Any end -function Base.showerror(io::IO, pie::RaggedParameterIndexingError) + +function Base.showerror(io::IO, pie::ParameterIndexingError) return print( io, "Indexing with parameters is deprecated. Use `getp(A, $(pie.sym))` for parameter indexing." ) end -# Symbolic indexing methods +# Symbolic Indexing Methods for (symtype, elsymtype, valtype, errcheck) in [ ( ScalarSymbolic, SymbolicIndexingInterface.SymbolicTypeTrait, Any, @@ -557,341 +731,998 @@ for (symtype, elsymtype, valtype, errcheck) in [ :(all(x -> is_parameter(A, x) && !is_timeseries_parameter(A, x), sym)), ), ] - @eval Base.@propagate_inbounds function _ragged_getindex( + @eval Base.@propagate_inbounds function _getindex( A::AbstractRaggedDiffEqArray, ::$symtype, ::$elsymtype, sym::$valtype, arg... ) if $errcheck - throw(RaggedParameterIndexingError(sym)) + throw(ParameterIndexingError(sym)) end return getu(A, sym)(A, arg...) end end -Base.@propagate_inbounds function _ragged_getindex( +Base.@propagate_inbounds function _getindex( A::AbstractRaggedDiffEqArray, ::ScalarSymbolic, ::NotSymbolic, ::SymbolicIndexingInterface.SolvedVariables, args... ) return getindex(A, variable_symbols(A), args...) end -Base.@propagate_inbounds function _ragged_getindex( +Base.@propagate_inbounds function _getindex( A::AbstractRaggedDiffEqArray, ::ScalarSymbolic, ::NotSymbolic, ::SymbolicIndexingInterface.AllVariables, args... ) return getindex(A, all_variable_symbols(A), args...) end -# ═══════════════════════════════════════════════════════════════════════════════ -# Observed functions -# ═══════════════════════════════════════════════════════════════════════════════ +@inline _column_indices(VA::AbstractRaggedVectorOfArray, idx) = idx +@inline _column_indices(VA::AbstractRaggedVectorOfArray, idx::Colon) = eachindex(VA.u) +@inline function _column_indices(VA::AbstractRaggedVectorOfArray, idx::AbstractArray{Bool}) + return findall(idx) +end +@inline function _column_indices(VA::AbstractRaggedVectorOfArray, idx::RaggedEnd) + # RaggedEnd with dim=0 means it's just a plain index stored in offset + return idx.dim == 0 ? idx.offset : idx +end + +@inline function _column_indices(VA::AbstractRaggedVectorOfArray, idx::RaggedRange) + # RaggedRange with dim=0 means it's a column range with pre-resolved indices + if idx.dim == 0 + # Create a range with the offset as the stop value + return Base.range(idx.start; step = idx.step, stop = idx.offset) + else + # dim != 0 means it's an inner-dimension range that needs column expansion + return idx + end +end + +@inline _resolve_ragged_index(idx, ::AbstractRaggedVectorOfArray, ::Any) = idx +@inline function _resolve_ragged_index(idx::RaggedEnd, VA::AbstractRaggedVectorOfArray, col) + if idx.dim == 0 + # Special case: dim=0 means the offset contains the actual index value + return idx.offset + else + return lastindex(VA.u[col], idx.dim) + idx.offset + end +end +@inline function _resolve_ragged_index(idx::RaggedRange, VA::AbstractRaggedVectorOfArray, col) + stop_val = if idx.dim == 0 + # dim == 0 is the sentinel for an already-resolved plain index stored in offset + idx.offset + else + lastindex(VA.u[col], idx.dim) + idx.offset + end + return Base.range(idx.start; step = idx.step, stop = stop_val) +end +@inline function _resolve_ragged_index( + idx::AbstractRange{<:RaggedEnd}, VA::AbstractRaggedVectorOfArray, col + ) + return Base.range( + _resolve_ragged_index(first(idx), VA, col); step = step(idx), + stop = _resolve_ragged_index(last(idx), VA, col) + ) +end +@inline function _resolve_ragged_index(idx::Base.Slice, VA::AbstractRaggedVectorOfArray, col) + return Base.Slice(_resolve_ragged_index(idx.indices, VA, col)) +end +@inline function _resolve_ragged_index(idx::CartesianIndex, VA::AbstractRaggedVectorOfArray, col) + return CartesianIndex(_resolve_ragged_indices(Tuple(idx), VA, col)...) +end +@inline function _resolve_ragged_index( + idx::AbstractArray{<:RaggedEnd}, VA::AbstractRaggedVectorOfArray, col + ) + return map(i -> _resolve_ragged_index(i, VA, col), idx) +end +@inline function _resolve_ragged_index( + idx::AbstractArray{<:RaggedRange}, VA::AbstractRaggedVectorOfArray, col + ) + return map(i -> _resolve_ragged_index(i, VA, col), idx) +end +@inline function _resolve_ragged_index(idx::AbstractArray, VA::AbstractRaggedVectorOfArray, col) + return _has_ragged_end(idx) ? map(i -> _resolve_ragged_index(i, VA, col), idx) : idx +end + +@inline function _resolve_ragged_indices(idxs::Tuple, VA::AbstractRaggedVectorOfArray, col) + return map(i -> _resolve_ragged_index(i, VA, col), idxs) +end + +@inline function _has_ragged_end(x) + x isa RaggedEnd && return true + x isa RaggedRange && return true + x isa Base.Slice && return _has_ragged_end(x.indices) + x isa CartesianIndex && return _has_ragged_end(Tuple(x)) + x isa AbstractRange && return eltype(x) <: Union{RaggedEnd, RaggedRange} + if x isa AbstractArray + el = eltype(x) + return el <: Union{RaggedEnd, RaggedRange} || + (el === Any && any(_has_ragged_end, x)) + end + x isa Tuple && return any(_has_ragged_end, x) + return false +end +@inline _has_ragged_end(x, xs...) = _has_ragged_end(x) || _has_ragged_end(xs) + +# Helper function to resolve RaggedEnd objects in a tuple of arguments +@inline function _resolve_ragged_end_args(A::AbstractRaggedVectorOfArray, args::Tuple) + # Handle empty tuple case + length(args) == 0 && return args + if !_has_ragged_end(args...) + return args + end + # For now, we need to resolve only the last argument if it's RaggedEnd (column selector) + # This handles the common case sol[:x, end] where end gets converted to RaggedEnd(0, lastindex) + if args[end] isa RaggedEnd + resolved_last = _column_indices(A, args[end]) + if length(args) == 1 + return (resolved_last,) + else + return (Base.front(args)..., resolved_last) + end + elseif args[end] isa RaggedRange + # Only pre-resolve if it's an inner-dimension range (dim != 0) + # Column ranges (dim == 0) are handled later by _column_indices + if args[end].dim == 0 + # Column range - let _column_indices handle it + return args + else + resolved_last = _resolve_ragged_index(args[end], A, 1) + if length(args) == 1 + return (resolved_last,) + else + return (Base.front(args)..., resolved_last) + end + end + end + return args +end + +# Helper function to preserve RaggedDiffEqArray type when slicing +@inline function _preserve_array_type(A::AbstractRaggedVectorOfArray, u_slice, col_idxs) + return RaggedVectorOfArray(u_slice) +end + +@inline function _preserve_array_type(A::AbstractRaggedDiffEqArray, u_slice, col_idxs) + return RaggedDiffEqArray( + u_slice, A.t[col_idxs], parameter_values(A), symbolic_container(A); + interp = A.interp, dense = A.dense + ) +end + +@inline function _ragged_getindex(A::AbstractRaggedVectorOfArray, I...) + n = ndims(A) + # Special-case when user provided one fewer index than ndims(A): last index is column selector. + if length(I) == n - 1 + return _ragged_getindex_nm1dims(A, I...) + else + return _ragged_getindex_full(A, I...) + end +end + +@inline function _ragged_getindex_nm1dims(A::AbstractRaggedVectorOfArray, I...) + raw_cols = last(I) + # Determine if we're doing column selection (preserve type) or inner-dimension selection (don't preserve) + is_column_selection = if raw_cols isa RaggedEnd && raw_cols.dim != 0 + false # Inner dimension - don't preserve type + elseif raw_cols isa RaggedRange && raw_cols.dim != 0 + true # Inner dimension range converted to column range - DO preserve type + else + true # Column selection (dim == 0 or not ragged) + end + + # If the raw selector is a RaggedEnd/RaggedRange referring to inner dims, reinterpret as column selector. + cols = if raw_cols isa RaggedEnd && raw_cols.dim != 0 + lastindex(A.u) + raw_cols.offset + elseif raw_cols isa RaggedRange && raw_cols.dim != 0 + # Convert inner-dimension range to column range by resolving bounds + start_val = raw_cols.start < 0 ? lastindex(A.u) + raw_cols.start : raw_cols.start + stop_val = lastindex(A.u) + raw_cols.offset + Base.range(start_val; step = raw_cols.step, stop = stop_val) + else + _column_indices(A, raw_cols) + end + prefix = Base.front(I) + if cols isa Int + resolved_prefix = _resolve_ragged_indices(prefix, A, cols) + inner_nd = ndims(A.u[cols]) + n_missing = inner_nd - length(resolved_prefix) + padded = if n_missing > 0 + if all(idx -> idx === Colon(), resolved_prefix) + (resolved_prefix..., ntuple(_ -> Colon(), n_missing)...) + else + ( + resolved_prefix..., + (lastindex(A.u[cols], length(resolved_prefix) + i) for i in 1:n_missing)..., + ) + end + else + resolved_prefix + end + return A.u[cols][padded...] + else + u_slice = [ + begin + resolved_prefix = _resolve_ragged_indices(prefix, A, col) + inner_nd = ndims(A.u[col]) + n_missing = inner_nd - length(resolved_prefix) + padded = if n_missing > 0 + if all(idx -> idx === Colon(), resolved_prefix) + ( + resolved_prefix..., + ntuple(_ -> Colon(), n_missing)..., + ) + else + ( + resolved_prefix..., + ( + lastindex( + A.u[col], + length(resolved_prefix) + i + ) for i in 1:n_missing + )..., + ) + end + else + resolved_prefix + end + A.u[col][padded...] + end + for col in cols + ] + # Only preserve RaggedDiffEqArray type if we're selecting actual columns, not inner dimensions + if is_column_selection + return _preserve_array_type(A, u_slice, cols) + else + return RaggedVectorOfArray(u_slice) + end + end +end + +@inline function _padded_resolved_indices(prefix, A::AbstractRaggedVectorOfArray, col) + resolved = _resolve_ragged_indices(prefix, A, col) + inner_nd = ndims(A.u[col]) + padded = (resolved..., ntuple(_ -> Colon(), max(inner_nd - length(resolved), 0))...) + return padded +end + +@inline function _ragged_getindex_full(A::AbstractRaggedVectorOfArray, I...) + # Otherwise, use the full-length interpretation (last index is column selector; missing columns default to Colon()). + n = ndims(A) + cols, prefix = if length(I) == n + last(I), Base.front(I) + else + Colon(), I + end + if cols isa Int + if all(idx -> idx === Colon(), prefix) + return A.u[cols] + end + return A.u[cols][_padded_resolved_indices(prefix, A, cols)...] + else + col_idxs = _column_indices(A, cols) + # Resolve sentinel RaggedEnd/RaggedRange (dim==0) for column selection + if col_idxs isa RaggedEnd || col_idxs isa RaggedRange + col_idxs = _resolve_ragged_index(col_idxs, A, 1) + end + # If we're selecting whole inner arrays (all leading indices are Colons), + # keep the result as a RaggedVectorOfArray to match non-ragged behavior. + if all(idx -> idx === Colon(), prefix) + if col_idxs isa Int + return A.u[col_idxs] + else + return _preserve_array_type(A, A.u[col_idxs], col_idxs) + end + end + # If col_idxs resolved to a single Int, handle it directly + vals = map(col_idxs) do col + A.u[col][_padded_resolved_indices(prefix, A, col)...] + end + if col_idxs isa Int + return vals + else + return stack(vals) + end + end +end + +@inline function _checkbounds_ragged(::Type{Bool}, VA::AbstractRaggedVectorOfArray, idxs...) + cols = _column_indices(VA, last(idxs)) + prefix = Base.front(idxs) + if cols isa Int + resolved = _resolve_ragged_indices(prefix, VA, cols) + return checkbounds(Bool, VA.u, cols) && checkbounds(Bool, VA.u[cols], resolved...) + else + for col in cols + resolved = _resolve_ragged_indices(prefix, VA, col) + checkbounds(Bool, VA.u, col) || return false + checkbounds(Bool, VA.u[col], resolved...) || return false + end + return true + end +end + +Base.@propagate_inbounds function Base.getindex(A::AbstractRaggedVectorOfArray, _arg, args...) + symtype = symbolic_type(_arg) + elsymtype = symbolic_type(eltype(_arg)) + + return if symtype == NotSymbolic() && elsymtype == NotSymbolic() + if _has_ragged_end(_arg, args...) + return _ragged_getindex(A, _arg, args...) + end + if _arg isa Union{Tuple, AbstractArray} && + any(x -> symbolic_type(x) != NotSymbolic(), _arg) + _getindex(A, symtype, elsymtype, _arg, args...) + else + _getindex(A, symtype, _arg, args...) + end + else + # Resolve any RaggedEnd objects in args before passing to symbolic indexing + resolved_args = _resolve_ragged_end_args(A, args) + _getindex(A, symtype, elsymtype, _arg, resolved_args...) + end +end + +Base.@propagate_inbounds function Base.getindex( + A::Adjoint{T, <:AbstractRaggedVectorOfArray}, idxs... + ) where {T} + return getindex(A.parent, reverse(to_indices(A, idxs))...) +end -function _ragged_observed(A::AbstractRaggedDiffEqArray{T, N}, sym, i::Int) where {T, N} +function _observed(A::AbstractRaggedDiffEqArray{T, N}, sym, i::Int) where {T, N} return observed(A, sym)(A.u[i], A.p, A.t[i]) end -function _ragged_observed( - A::AbstractRaggedDiffEqArray{T, N}, sym, i::AbstractArray{Int} - ) where {T, N} +function _observed(A::AbstractRaggedDiffEqArray{T, N}, sym, i::AbstractArray{Int}) where {T, N} return observed(A, sym).(A.u[i], (A.p,), A.t[i]) end -function _ragged_observed( - A::AbstractRaggedDiffEqArray{T, N}, sym, ::Colon - ) where {T, N} +function _observed(A::AbstractRaggedDiffEqArray{T, N}, sym, ::Colon) where {T, N} return observed(A, sym).(A.u, (A.p,), A.t) end -# ═══════════════════════════════════════════════════════════════════════════════ -# Sizes of individual inner arrays -# ═══════════════════════════════════════════════════════════════════════════════ +Base.@propagate_inbounds function Base.setindex!( + VA::AbstractRaggedVectorOfArray{T, N}, v, + ::Colon, I::Int + ) where {T, N} + return VA.u[I] = v +end -""" - inner_sizes(r::AbstractRaggedVectorOfArray) +Base.@propagate_inbounds Base.setindex!(VA::AbstractRaggedVectorOfArray, v, I::Int) = Base.setindex!( + VA.u, v, I +) -Return a vector of the sizes of each inner array. -""" -inner_sizes(r::AbstractRaggedVectorOfArray) = size.(r.u) +Base.@propagate_inbounds function Base.setindex!( + VA::AbstractRaggedVectorOfArray{T, N}, v, + ::Colon, I::Colon + ) where {T, N} + return VA.u[I] = v +end -""" - inner_lengths(r::AbstractRaggedVectorOfArray) +Base.@propagate_inbounds Base.setindex!(VA::AbstractRaggedVectorOfArray, v, I::Colon) = Base.setindex!( + VA.u, v, I +) -Return a vector of the lengths of each inner array. -""" -inner_lengths(r::AbstractRaggedVectorOfArray) = length.(r.u) +Base.@propagate_inbounds function Base.setindex!( + VA::AbstractRaggedVectorOfArray{T, N}, v, + ::Colon, I::AbstractArray{Int} + ) where {T, N} + return VA.u[I] = v +end + +Base.@propagate_inbounds Base.setindex!(VA::AbstractRaggedVectorOfArray, v, I::AbstractArray{Int}) = Base.setindex!( + VA.u, v, I +) + +Base.@propagate_inbounds function Base.setindex!( + VA::AbstractRaggedVectorOfArray{T, N}, v, i::Int, + ::Colon + ) where {T, N} + for j in 1:length(VA.u) + VA.u[j][i] = v[j] + end + return v +end +Base.@propagate_inbounds function Base.setindex!( + VA::AbstractRaggedVectorOfArray{T, N}, x, + ii::CartesianIndex + ) where {T, N} + ti = Tuple(ii) + i = last(ti) + jj = CartesianIndex(Base.front(ti)) + return VA.u[i][jj] = x +end + +Base.@propagate_inbounds function Base.setindex!( + VA::AbstractRaggedVectorOfArray{T, N}, + x, + idxs::Union{Int, Colon, CartesianIndex, AbstractArray{Int}, AbstractArray{Bool}}... + ) where { + T, N, + } + v = view(VA, idxs...) + # error message copied from Base by running `ones(3, 3, 3)[:, 2, :] = 2` + if length(v) != length(x) + throw(ArgumentError("indexed assignment with a single value to possibly many locations is not supported; perhaps use broadcasting `.=` instead?")) + end + for (i, j) in zip(eachindex(v), eachindex(x)) + v[i] = x[j] + end + return x +end + +# Interface for the two-dimensional indexing, a more standard AbstractArray interface +@inline Base.size(VA::AbstractRaggedVectorOfArray) = (size(VA.u[1])..., length(VA.u)) +@inline Base.size(VA::AbstractRaggedVectorOfArray, i) = size(VA)[i] +@inline Base.size(A::Adjoint{T, <:AbstractRaggedVectorOfArray}) where {T} = reverse(size(A.parent)) +@inline Base.size(A::Adjoint{T, <:AbstractRaggedVectorOfArray}, i) where {T} = size(A)[i] +Base.axes(VA::AbstractRaggedVectorOfArray) = Base.OneTo.(size(VA)) +Base.axes(VA::AbstractRaggedVectorOfArray, d::Int) = Base.OneTo(size(VA)[d]) + +Base.@propagate_inbounds function Base.setindex!( + VA::AbstractRaggedVectorOfArray{T, N}, v, + I::Int... + ) where {T, N} + return VA.u[I[end]][Base.front(I)...] = v +end -export inner_sizes, inner_lengths +function Base.:(==)(A::AbstractRaggedVectorOfArray, B::AbstractRaggedVectorOfArray) + return A.u == B.u +end +function Base.:(==)(A::AbstractRaggedVectorOfArray, B::AbstractArray) + return A.u == B +end +Base.:(==)(A::AbstractArray, B::AbstractRaggedVectorOfArray) = B == A -# ═══════════════════════════════════════════════════════════════════════════════ -# Copy, zero, similar, fill! -# ═══════════════════════════════════════════════════════════════════════════════ +# The iterator will be over the subarrays of the container, not the individual elements +# unlike an true AbstractArray +function Base.iterate(VA::AbstractRaggedVectorOfArray, state = 1) + return state >= length(VA.u) + 1 ? nothing : (VA[:, state], state + 1) +end +tuples(VA::RaggedDiffEqArray) = tuple.(VA.t, VA.u) -function _ragged_copyfield(r, fname) +# Growing the array simply adds to the container vector +function _copyfield(VA, fname) return if fname == :u - copy(r.u) + copy(VA.u) elseif fname == :t - copy(r.t) + copy(VA.t) else - getfield(r, fname) + getfield(VA, fname) end end - -function Base.copy(r::AbstractRaggedVectorOfArray) - return typeof(r)((_ragged_copyfield(r, fname) for fname in fieldnames(typeof(r)))...) +function Base.copy(VA::AbstractRaggedVectorOfArray) + return typeof(VA)((_copyfield(VA, fname) for fname in fieldnames(typeof(VA)))...) end -function Base.zero(r::AbstractRaggedVectorOfArray) - T = typeof(r) - u_zero = [zero(u) for u in r.u] - fields = [ - fname == :u ? u_zero : _ragged_copyfield(r, fname) - for fname in fieldnames(T) - ] - return T(fields...) +function Base.zero(VA::AbstractRaggedVectorOfArray) + val = copy(VA) + val.u .= zero.(VA.u) + return val end -function Base.similar(r::RaggedVectorOfArray) - return RaggedVectorOfArray(similar.(r.u)) -end +Base.sizehint!(VA::AbstractRaggedVectorOfArray{T, N}, i) where {T, N} = sizehint!(VA.u, i) -function Base.similar(r::RaggedVectorOfArray, ::Type{T}) where {T} - return RaggedVectorOfArray(similar.(r.u, T)) +Base.reverse!(VA::AbstractRaggedVectorOfArray) = reverse!(VA.u) +Base.reverse(VA::AbstractRaggedVectorOfArray) = RaggedVectorOfArray(reverse(VA.u)) +function Base.reverse(VA::AbstractRaggedDiffEqArray) + return RaggedDiffEqArray( + reverse(VA.u), VA.t, parameter_values(VA), symbolic_container(VA); + interp = VA.interp, dense = VA.dense + ) end -function Base.fill!(r::AbstractRaggedVectorOfArray, x) - for u in r.u - fill!(u, x) +function Base.resize!(VA::AbstractRaggedVectorOfArray, i::Integer) + if Base.hasproperty(VA, :sys) && VA.sys !== nothing + error("resize! is not allowed on AbstractRaggedVectorOfArray with a sys") + end + Base.resize!(VA.u, i) + return if Base.hasproperty(VA, :t) && VA.t !== nothing + Base.resize!(VA.t, i) end - return r end -# ═══════════════════════════════════════════════════════════════════════════════ -# Mutation: push!, append!, resize! -# ═══════════════════════════════════════════════════════════════════════════════ +function Base.pointer(VA::AbstractRaggedVectorOfArray) + return Base.pointer(VA.u) +end -function Base.push!(r::AbstractRaggedVectorOfArray, new_item::AbstractArray) - return push!(r.u, new_item) +function Base.push!(VA::AbstractRaggedVectorOfArray{T, N}, new_item::AbstractArray) where {T, N} + return push!(VA.u, new_item) end function Base.append!( - r::AbstractRaggedVectorOfArray, other::AbstractRaggedVectorOfArray - ) - for item in copy(other.u) - push!(r, item) + VA::AbstractRaggedVectorOfArray{T, N}, + new_item::AbstractRaggedVectorOfArray{T, N} + ) where {T, N} + for item in copy(new_item) + push!(VA, item) end - return r + return VA end -Base.sizehint!(r::AbstractRaggedVectorOfArray, i) = sizehint!(r.u, i) +function Base.stack(VA::AbstractRaggedVectorOfArray; dims = :) + return stack(stack.(VA.u); dims) +end -Base.reverse!(r::AbstractRaggedVectorOfArray) = (reverse!(r.u); r) -Base.reverse(r::RaggedVectorOfArray) = RaggedVectorOfArray(reverse(r.u)) -function Base.reverse(r::RaggedDiffEqArray) - return RaggedDiffEqArray( - reverse(r.u), r.t, r.p, r.sys; discretes = r.discretes, - interp = r.interp, dense = r.dense +# AbstractArray methods +function Base.view( + A::AbstractRaggedVectorOfArray{T, N, <:AbstractVector{T}}, + I::Vararg{Any, M} + ) where {T, N, M} + @inline + if length(I) == 1 + J = map(i -> Base.unalias(A, i), to_indices(A, I)) + elseif length(I) == 2 && (I[1] == Colon() || I[1] == 1) + J = map(i -> Base.unalias(A, i), to_indices(A, Base.tail(I))) + else + J = map(i -> Base.unalias(A, i), to_indices(A, I)) + end + @boundscheck checkbounds(A, J...) + return SubArray(A, J) +end +function Base.view(A::AbstractRaggedVectorOfArray, I::Vararg{Any, M}) where {M} + @inline + # 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)) + @boundscheck checkbounds(A, J...) + return SubArray(A, J) +end +function Base.SubArray(parent::AbstractRaggedVectorOfArray, indices::Tuple) + @inline + return SubArray( + IndexStyle(Base.viewindexing(indices), IndexStyle(parent)), parent, + Base.ensure_indexable(indices), Base.index_dimsum(indices...) ) end +Base.isassigned(VA::AbstractRaggedVectorOfArray, idxs...) = checkbounds(Bool, VA, idxs...) +function Base.check_parent_index_match( + ::AbstractRaggedVectorOfArray{T, N}, ::NTuple{N, Bool} + ) where {T, N} + return nothing +end +Base.ndims(::AbstractRaggedVectorOfArray{T, N}) where {T, N} = N +Base.ndims(::Type{<:AbstractRaggedVectorOfArray{T, N}}) where {T, N} = N +function Base.checkbounds( + ::Type{Bool}, VA::AbstractRaggedVectorOfArray{T, N, <:AbstractVector{T}}, + idxs... + ) where {T, N} + if _has_ragged_end(idxs...) + return _checkbounds_ragged(Bool, VA, idxs...) + end + if length(idxs) == 2 && (idxs[1] == Colon() || idxs[1] == 1) + return checkbounds(Bool, VA.u, idxs[2]) + end + return checkbounds(Bool, VA.u, idxs...) +end +function Base.checkbounds(::Type{Bool}, VA::AbstractRaggedVectorOfArray, idx...) + if _has_ragged_end(idx...) + return _checkbounds_ragged(Bool, VA, idx...) + end + checkbounds(Bool, VA.u, last(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 +end +function Base.checkbounds(VA::AbstractRaggedVectorOfArray, idx...) + return checkbounds(Bool, VA, idx...) || throw(BoundsError(VA, idx)) +end function Base.copyto!( - dest::AbstractRaggedVectorOfArray, src::AbstractRaggedVectorOfArray - ) + dest::AbstractRaggedVectorOfArray{T, N}, + src::AbstractRaggedVectorOfArray{T2, N} + ) where {T, T2, N} for (i, j) in zip(eachindex(dest.u), eachindex(src.u)) - copyto!(dest.u[i], src.u[j]) + if ArrayInterface.ismutable(dest.u[i]) || dest.u[i] isa AbstractRaggedVectorOfArray + copyto!(dest.u[i], src.u[j]) + else + dest.u[i] = StaticArraysCore.similar_type(dest.u[i])(src.u[j]) + end + end + return +end +function Base.copyto!( + dest::AbstractRaggedVectorOfArray{T, N}, src::AbstractArray{T2, N} + ) where {T, T2, N} + for (i, slice) in zip(eachindex(dest.u), eachslice(src, dims = ndims(src))) + if ArrayInterface.ismutable(dest.u[i]) || dest.u[i] isa AbstractRaggedVectorOfArray + copyto!(dest.u[i], slice) + else + dest.u[i] = StaticArraysCore.similar_type(dest.u[i])(slice) + end end return dest end +function Base.copyto!( + dest::AbstractRaggedVectorOfArray{T, N, <:AbstractVector{T}}, + src::AbstractVector{T2} + ) where {T, T2, N} + copyto!(dest.u, src) + return dest +end +# Required for broadcasted setindex! when slicing across subarrays +# E.g. if `va = RaggedVectorOfArray([rand(3, 3) for i in 1:5])` +# Need this method for `va[2, :, :] .= 3.0` +Base.@propagate_inbounds function Base.maybeview(A::AbstractRaggedVectorOfArray, I...) + return view(A, I...) +end -# ═══════════════════════════════════════════════════════════════════════════════ -# Equality -# ═══════════════════════════════════════════════════════════════════════════════ +# Operations +function Base.isapprox( + A::AbstractRaggedVectorOfArray, + B::Union{AbstractRaggedVectorOfArray, AbstractArray}; + kwargs... + ) + return all(isapprox.(A, B; kwargs...)) +end -function Base.:(==)(a::AbstractRaggedVectorOfArray, b::AbstractRaggedVectorOfArray) - return a.u == b.u +function Base.isapprox(A::AbstractArray, B::AbstractRaggedVectorOfArray; kwargs...) + return all(isapprox.(A, B; kwargs...)) end -# ═══════════════════════════════════════════════════════════════════════════════ -# Broadcasting — column-wise, preserving ragged structure -# ═══════════════════════════════════════════════════════════════════════════════ +for op in [:(Base.:-), :(Base.:+)] + @eval function ($op)(A::AbstractRaggedVectorOfArray, B::AbstractRaggedVectorOfArray) + return ($op).(A, B) + end + @eval Base.@propagate_inbounds function ($op)( + A::AbstractRaggedVectorOfArray, + B::AbstractArray + ) + @boundscheck length(A) == length(B) + return RaggedVectorOfArray([($op).(a, b) for (a, b) in zip(A, B)]) + end + @eval Base.@propagate_inbounds function ($op)( + A::AbstractArray, B::AbstractRaggedVectorOfArray + ) + @boundscheck length(A) == length(B) + return RaggedVectorOfArray([($op).(a, b) for (a, b) in zip(A, B)]) + end +end -struct RaggedVectorOfArrayStyle{N} <: Broadcast.BroadcastStyle end -RaggedVectorOfArrayStyle{N}(::Val{N}) where {N} = RaggedVectorOfArrayStyle{N}() -RaggedVectorOfArrayStyle(::Val{N}) where {N} = RaggedVectorOfArrayStyle{N}() +for op in [:(Base.:/), :(Base.:\), :(Base.:*)] + if op !== :(Base.:/) + @eval ($op)(A::Number, B::AbstractRaggedVectorOfArray) = ($op).(A, B) + end + if op !== :(Base.:\) + @eval ($op)(A::AbstractRaggedVectorOfArray, B::Number) = ($op).(A, B) + end +end -function Broadcast.BroadcastStyle(::Type{<:AbstractRaggedVectorOfArray{T, N}}) where {T, N} - return RaggedVectorOfArrayStyle{N}() +function Base.CartesianIndices(VA::AbstractRaggedVectorOfArray) + if !allequal(size.(VA.u)) + error("CartesianIndices only valid for non-ragged arrays") + end + return CartesianIndices((size(VA.u[1])..., length(VA.u))) end -# Make ragged arrays broadcastable without being collected -Broadcast.broadcastable(r::AbstractRaggedVectorOfArray) = r +# Tools for creating similar objects +Base.eltype(::Type{<:AbstractRaggedVectorOfArray{T}}) where {T} = T -# RaggedVectorOfArrayStyle wins over DefaultArrayStyle at all dims -function Broadcast.BroadcastStyle( - a::RaggedVectorOfArrayStyle, ::Base.Broadcast.DefaultArrayStyle{0} - ) - return a -end -function Broadcast.BroadcastStyle( - a::RaggedVectorOfArrayStyle{N}, ::Base.Broadcast.DefaultArrayStyle{M} - ) where {M, N} - return RaggedVectorOfArrayStyle(Val(max(M, N))) +@inline function Base.similar(VA::AbstractRaggedVectorOfArray, args...) + if args[end] isa Type + return Base.similar(eltype(VA)[], args..., size(VA)) + else + return Base.similar(eltype(VA)[], args...) + end end -# RaggedVectorOfArrayStyle wins over itself at different dims -function Broadcast.BroadcastStyle( - ::RaggedVectorOfArrayStyle{M}, ::RaggedVectorOfArrayStyle{N} - ) where {M, N} - return RaggedVectorOfArrayStyle(Val(max(M, N))) + +function Base.similar( + vec::RaggedVectorOfArray{ + T, N, AT, + } + ) where {T, N, AT <: AbstractArray{<:AbstractArray{T}}} + return RaggedVectorOfArray(similar.(Base.parent(vec))) end -# Number of inner arrays in a broadcast -_ragged_narrays(::Any) = 0 -_ragged_narrays(r::AbstractRaggedVectorOfArray) = length(r.u) -_ragged_narrays(bc::Broadcast.Broadcasted) = __ragged_narrays(bc.args) -function __ragged_narrays(args::Tuple) - a = _ragged_narrays(args[1]) - b = __ragged_narrays(Base.tail(args)) - return a == 0 ? b : ( - b == 0 ? a : ( - a == b ? a : - throw(DimensionMismatch("number of arrays must be equal")) - ) - ) +function Base.similar( + vec::RaggedVectorOfArray{ + T, N, AT, + } + ) where {T, N, AT <: AbstractArray{<:StaticArraysCore.StaticVecOrMat{T}}} + # this avoids behavior such as similar(SVector) returning an MVector + return RaggedVectorOfArray(similar(Base.parent(vec))) end -__ragged_narrays(args::Tuple{Any}) = _ragged_narrays(args[1]) -__ragged_narrays(::Tuple{}) = 0 -# Unpack the i-th inner array from broadcast arguments -_ragged_unpack(x, ::Any) = x -_ragged_unpack(r::AbstractRaggedVectorOfArray, i) = r.u[i] -function _ragged_unpack(bc::Broadcast.Broadcasted{Style}, i) where {Style} - return Broadcast.Broadcasted{Style}(bc.f, _ragged_unpack_args(i, bc.args)) +@inline function Base.similar(VA::RaggedVectorOfArray, ::Type{T} = eltype(VA)) where {T} + if eltype(VA.u) <: Union{AbstractArray, AbstractRaggedVectorOfArray} + return RaggedVectorOfArray(similar.(VA.u, T)) + else + return RaggedVectorOfArray(similar(VA.u, T)) + end end -function _ragged_unpack(bc::Broadcast.Broadcasted{<:RaggedVectorOfArrayStyle}, i) - return Broadcast.Broadcasted(bc.f, _ragged_unpack_args(i, bc.args)) + +@inline function Base.similar(VA::RaggedVectorOfArray, dims::N) where {N <: Number} + l = length(VA) + return if dims <= l + RaggedVectorOfArray(similar.(VA.u[1:dims])) + else + RaggedVectorOfArray([similar.(VA.u); [similar(VA.u[end]) for _ in (l + 1):dims]]) + end end -function _ragged_unpack_args(i, args::Tuple) - return (_ragged_unpack(args[1], i), _ragged_unpack_args(i, Base.tail(args))...) + +# fill! +# For RaggedDiffEqArray it ignores ts and fills only u +function Base.fill!(VA::AbstractRaggedVectorOfArray, x) + for i in 1:length(VA.u) + if VA[:, i] isa Union{AbstractArray, AbstractRaggedVectorOfArray} + if ArrayInterface.ismutable(VA.u[i]) || VA.u[i] isa AbstractRaggedVectorOfArray + fill!(VA[:, i], x) + else + # For immutable arrays like SVector, create a new filled array + VA.u[i] = fill(x, StaticArraysCore.similar_type(VA.u[i])) + end + else + VA[:, i] = x + end + end + return VA end -_ragged_unpack_args(i, args::Tuple{Any}) = (_ragged_unpack(args[1], i),) -_ragged_unpack_args(::Any, ::Tuple{}) = () -# copy: create new RaggedVectorOfArray from broadcast -@inline function Base.copy(bc::Broadcast.Broadcasted{<:RaggedVectorOfArrayStyle}) - bc = Broadcast.flatten(bc) - N = _ragged_narrays(bc) - u = map(1:N) do i - copy(_ragged_unpack(bc, i)) +Base.reshape(A::AbstractRaggedVectorOfArray, dims...) = Base.reshape(Array(A), dims...) + +# Need this for ODE_DEFAULT_UNSTABLE_CHECK from DiffEqBase to work properly +@inline Base.any(f, VA::AbstractRaggedVectorOfArray) = any(any(f, u) for u in VA.u) +@inline Base.all(f, VA::AbstractRaggedVectorOfArray) = all(all(f, u) for u in VA.u) + +# conversion tools +vecarr_to_vectors(VA::AbstractRaggedVectorOfArray) = [VA[i, :] for i in eachindex(VA.u[1])] +Base.vec(VA::AbstractRaggedVectorOfArray) = vec(convert(Array, VA)) # Allocates +# stack non-ragged arrays to convert them +function Base.convert(::Type{Array}, VA::AbstractRaggedVectorOfArray) + if !allequal(size.(VA.u)) + error("Can only convert non-ragged RaggedVectorOfArray to Array") end - return RaggedVectorOfArray(u) + return Array(VA) end -# Override materialize to skip instantiate (which needs axes/size on non-AbstractArray types) -# This causes 1 invalidation tree from Base.Broadcast.materialize(::Broadcasted), but it's -# unavoidable for non-AbstractArray types that need custom broadcasting, and this sublibrary -# is opt-in to isolate exactly this kind of invalidation from the main path. -@inline function Broadcast.materialize(bc::Broadcast.Broadcasted{<:RaggedVectorOfArrayStyle}) - return copy(bc) +# statistics +@inline Base.sum(VA::AbstractRaggedVectorOfArray; kwargs...) = sum(identity, VA; kwargs...) +@inline function Base.sum(f, VA::AbstractRaggedVectorOfArray; kwargs...) + return mapreduce(f, Base.add_sum, VA; kwargs...) +end +@inline Base.prod(VA::AbstractRaggedVectorOfArray; kwargs...) = prod(identity, VA; kwargs...) +@inline function Base.prod(f, VA::AbstractRaggedVectorOfArray; kwargs...) + return mapreduce(f, Base.mul_prod, VA; kwargs...) end -# materialize! bypass: skip shape checking since ragged arrays have no rectangular axes -@inline function Broadcast.materialize!( - dest::AbstractRaggedVectorOfArray, - bc::Broadcast.Broadcasted{<:RaggedVectorOfArrayStyle} - ) - return copyto!(dest, bc) +@inline Base.adjoint(VA::AbstractRaggedVectorOfArray) = Adjoint(VA) + +# linear algebra +ArrayInterface.issingular(va::AbstractRaggedVectorOfArray) = ArrayInterface.issingular(Matrix(va)) + +# make it show just like its data +function Base.show(io::IO, m::MIME"text/plain", x::AbstractRaggedVectorOfArray) + (println(io, summary(x), ':'); show(io, m, x.u)) +end +function Base.summary(A::AbstractRaggedVectorOfArray{T, N}) where {T, N} + return string("RaggedVectorOfArray{", T, ",", N, "}") end -@inline function Broadcast.materialize!( - dest::AbstractRaggedVectorOfArray, - bc::Broadcast.Broadcasted - ) - return copyto!(dest, bc) +function Base.show(io::IO, m::MIME"text/plain", x::AbstractRaggedDiffEqArray) + (print(io, "t: "); show(io, m, x.t); println(io); print(io, "u: "); show(io, m, x.u)) end -@inline function Broadcast.materialize!( - dest::AbstractRaggedVectorOfArray, - x::Any - ) - return Broadcast.materialize!(dest, Broadcast.broadcasted(identity, x)) +Base.map(f, A::AbstractRaggedVectorOfArray) = map(f, A.u) + +function Base.mapreduce(f, op, A::AbstractRaggedVectorOfArray; kwargs...) + return mapreduce(f, op, view(A, ntuple(_ -> :, ndims(A))...); kwargs...) +end +function Base.mapreduce( + f, op, A::AbstractRaggedVectorOfArray{T, 1, <:AbstractVector{T}}; kwargs... + ) where {T} + return mapreduce(f, op, A.u; kwargs...) end -# copyto!: in-place broadcast -@inline function Base.copyto!( - dest::AbstractRaggedVectorOfArray, - bc::Broadcast.Broadcasted{<:RaggedVectorOfArrayStyle} - ) +## broadcasting + +struct RaggedVectorOfArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end # N is only used when voa sees other abstract arrays +RaggedVectorOfArrayStyle{N}(::Val{N}) where {N} = RaggedVectorOfArrayStyle{N}() +RaggedVectorOfArrayStyle(::Val{N}) where {N} = RaggedVectorOfArrayStyle{N}() + +# The order is important here. We want to override Base.Broadcast.DefaultArrayStyle to return another Base.Broadcast.DefaultArrayStyle. +Broadcast.BroadcastStyle(a::RaggedVectorOfArrayStyle, ::Base.Broadcast.DefaultArrayStyle{0}) = a +function Broadcast.BroadcastStyle( + ::RaggedVectorOfArrayStyle{N}, + a::Base.Broadcast.DefaultArrayStyle{M} + ) where {M, N} + return Base.Broadcast.DefaultArrayStyle(Val(max(M, N))) +end +function Broadcast.BroadcastStyle( + ::RaggedVectorOfArrayStyle{N}, + a::Base.Broadcast.AbstractArrayStyle{M} + ) where {M, N} + return typeof(a)(Val(max(M, N))) +end +function Broadcast.BroadcastStyle( + ::RaggedVectorOfArrayStyle{M}, + ::RaggedVectorOfArrayStyle{N} + ) where {M, N} + return RaggedVectorOfArrayStyle(Val(max(M, N))) +end +function Broadcast.BroadcastStyle(::Type{<:AbstractRaggedVectorOfArray{T, N}}) where {T, N} + return RaggedVectorOfArrayStyle{N}() +end +# make vectorofarrays broadcastable so they aren't collected +Broadcast.broadcastable(x::AbstractRaggedVectorOfArray) = x + +# recurse through broadcast arguments and return a parent array for +# the first RaggedVoA or RaggedDiffEqArray in the bc arguments +function find_RaggedVoA_parent(args) + arg = Base.first(args) + if arg isa AbstractRaggedDiffEqArray + # if first(args) is a RaggedDiffEqArray, use the underlying + # field `u` of RaggedDiffEqArray as a parent array. + return arg.u + elseif arg isa AbstractRaggedVectorOfArray + return parent(arg) + else + return find_RaggedVoA_parent(Base.tail(args)) + end +end + +@inline function Base.copy(bc::Broadcast.Broadcasted{<:RaggedVectorOfArrayStyle}) bc = Broadcast.flatten(bc) - N = _ragged_narrays(bc) - @inbounds for i in 1:N - copyto!(dest.u[i], _ragged_unpack(bc, i)) + p = find_RaggedVoA_parent(bc.args) + + u = if p isa AbstractVector + # this is the default behavior in v3.15.0 + N = narrays(bc) + map(1:N) do i + copy(unpack_voa(bc, i)) + end + else # if p isa AbstractArray + map(enumerate(Iterators.product(axes(p)...))) do (i, _) + copy(unpack_voa(bc, i)) + end end - return dest + return RaggedVectorOfArray(rewrap(p, u)) end -# ═══════════════════════════════════════════════════════════════════════════════ -# Show -# ═══════════════════════════════════════════════════════════════════════════════ +rewrap(::Array, u) = u +rewrap(p, u) = convert(typeof(p), u) -function Base.show(io::IO, m::MIME"text/plain", r::AbstractRaggedVectorOfArray) - println(io, summary(r), ':') - return show(io, m, r.u) +for (type, N_expr) in [ + (Broadcast.Broadcasted{<:RaggedVectorOfArrayStyle}, :(narrays(bc))), + (Broadcast.Broadcasted{<:Broadcast.DefaultArrayStyle}, :(length(dest.u))), + ] + @eval @inline function Base.copyto!( + dest::AbstractRaggedVectorOfArray, + bc::$type + ) + bc = Broadcast.flatten(bc) + N = $N_expr + @inbounds for i in 1:N + if dest[:, i] isa AbstractArray + if ArrayInterface.ismutable(dest[:, i]) + copyto!(dest[:, i], unpack_voa(bc, i)) + else + unpacked = unpack_voa(bc, i) + arr_type = StaticArraysCore.similar_type(dest[:, i]) + dest[:, i] = if length(unpacked) == 1 && length(dest[:, i]) == 1 + arr_type(unpacked[1]) + elseif length(unpacked) == 1 + fill(copy(unpacked), arr_type) + else + arr_type(unpacked[j] for j in eachindex(unpacked)) + end + end + else + dest[:, i] = copy(unpack_voa(bc, i)) + end + end + return dest + end end -function Base.summary(r::AbstractRaggedVectorOfArray{T, N}) where {T, N} - return string("RaggedVectorOfArray{", T, ",", N, "}") +## broadcasting utils + +""" + narrays(A...) + +Retrieve number of arrays in the AbstractRaggedVectorOfArrays of a broadcast. +""" +narrays(A) = 0 +narrays(A::AbstractRaggedVectorOfArray) = length(A.u) +narrays(bc::Broadcast.Broadcasted) = _narrays(bc.args) +narrays(A, Bs...) = common_length(narrays(A), _narrays(Bs)) + +function common_length(a, b) + return a == 0 ? b : + ( + b == 0 ? a : + ( + a == b ? a : + throw(DimensionMismatch("number of arrays must be equal")) + ) + ) end -function Base.show(io::IO, m::MIME"text/plain", r::AbstractRaggedDiffEqArray) - print(io, "t: ") - show(io, m, r.t) - println(io) - print(io, "u: ") - return show(io, m, r.u) +_narrays(args::AbstractRaggedVectorOfArray) = length(args.u) +@inline _narrays(args::Tuple) = common_length(narrays(args[1]), _narrays(Base.tail(args))) +_narrays(args::Tuple{Any}) = _narrays(args[1]) +_narrays(::Any) = 0 + +# drop axes because it is easier to recompute +@inline function unpack_voa(bc::Broadcast.Broadcasted{Style}, i) where {Style} + return Broadcast.Broadcasted{Style}(bc.f, unpack_args_voa(i, bc.args)) +end +@inline function unpack_voa(bc::Broadcast.Broadcasted{<:RaggedVectorOfArrayStyle}, i) + return Broadcast.Broadcasted(bc.f, unpack_args_voa(i, bc.args)) +end +unpack_voa(x, ::Any) = x +unpack_voa(x::AbstractRaggedVectorOfArray, i) = x.u[i] +function unpack_voa(x::AbstractArray{T, N}, i) where {T, N} + return @view x[ntuple(x -> Colon(), N - 1)..., i] end -function Base.summary(r::AbstractRaggedDiffEqArray{T, N}) where {T, N} - return string("RaggedDiffEqArray{", T, ",", N, "}") +@inline function unpack_args_voa(i, args::Tuple) + return (unpack_voa(args[1], i), unpack_args_voa(i, Base.tail(args))...) end +unpack_args_voa(i, args::Tuple{Any}) = (unpack_voa(args[1], i),) +unpack_args_voa(::Any, args::Tuple{}) = () -# ═══════════════════════════════════════════════════════════════════════════════ -# Callable interface — dense interpolation -# ═══════════════════════════════════════════════════════════════════════════════ +# Conversion methods between Ragged and non-Ragged types -function (r::AbstractRaggedDiffEqArray)( - t, ::Type{deriv} = Val{0}; - idxs = nothing, continuity = :left - ) where {deriv} - r.interp === nothing && - error("No interpolation data is available. Provide an interpolation object via the `interp` keyword.") - return r.interp(t, idxs, deriv, r.p, continuity) -end +""" + VectorOfArray(r::AbstractRaggedVectorOfArray) -# ═══════════════════════════════════════════════════════════════════════════════ -# SymbolicIndexingInterface -# ═══════════════════════════════════════════════════════════════════════════════ +Convert a `RaggedVectorOfArray` to a regular `VectorOfArray`. +""" +function VectorOfArray(r::AbstractRaggedVectorOfArray) + return VectorOfArray(r.u) +end -SymbolicIndexingInterface.is_timeseries(::Type{<:AbstractRaggedVectorOfArray}) = Timeseries() +""" + RaggedVectorOfArray(va::AbstractVectorOfArray) -function SymbolicIndexingInterface.is_parameter_timeseries( - ::Type{ - RaggedDiffEqArray{ - T, N, A, B, - F, S, D, I, - }, - } - ) where {T, N, A, B, F, S, D <: ParameterIndexingProxy, I} - return Timeseries() +Convert a regular `VectorOfArray` to a `RaggedVectorOfArray`. +""" +function RaggedVectorOfArray(va::AbstractVectorOfArray) + return RaggedVectorOfArray(va.u) end -SymbolicIndexingInterface.state_values(r::AbstractRaggedDiffEqArray) = r.u -SymbolicIndexingInterface.current_time(r::AbstractRaggedDiffEqArray) = r.t -SymbolicIndexingInterface.parameter_values(r::AbstractRaggedDiffEqArray) = r.p -SymbolicIndexingInterface.parameter_values(r::AbstractRaggedDiffEqArray, i) = parameter_values(r.p, i) -SymbolicIndexingInterface.symbolic_container(r::AbstractRaggedDiffEqArray) = r.sys +""" + DiffEqArray(r::AbstractRaggedDiffEqArray) -RecursiveArrayTools.has_discretes(::T) where {T <: AbstractRaggedDiffEqArray} = hasfield(T, :discretes) -RecursiveArrayTools.get_discretes(r::AbstractRaggedDiffEqArray) = r.discretes +Convert a `RaggedDiffEqArray` to a regular `DiffEqArray`. +""" +function DiffEqArray(r::AbstractRaggedDiffEqArray) + return DiffEqArray( + r.u, r.t, r.p, r.sys; + discretes = r.discretes, interp = r.interp, dense = r.dense + ) +end + +""" + RaggedDiffEqArray(va::AbstractDiffEqArray) -function SymbolicIndexingInterface.get_parameter_timeseries_collection(r::AbstractRaggedDiffEqArray) - return r.discretes +Convert a regular `DiffEqArray` to a `RaggedDiffEqArray`. +""" +function RaggedDiffEqArray(va::AbstractDiffEqArray) + return RaggedDiffEqArray( + va.u, va.t, va.p, va.sys; + discretes = hasfield(typeof(va), :discretes) ? getfield(va, :discretes) : nothing, + interp = hasfield(typeof(va), :interp) ? getfield(va, :interp) : nothing, + dense = hasfield(typeof(va), :dense) ? getfield(va, :dense) : false + ) end -end # module +# Re-export has_discretes and get_discretes for the non-ragged types +has_discretes(::TT) where {TT <: AbstractDiffEqArray} = hasfield(TT, :discretes) + +end # module RecursiveArrayToolsRaggedArrays diff --git a/lib/RecursiveArrayToolsRaggedArrays/test/runtests.jl b/lib/RecursiveArrayToolsRaggedArrays/test/runtests.jl index c3626499..b3489c41 100644 --- a/lib/RecursiveArrayToolsRaggedArrays/test/runtests.jl +++ b/lib/RecursiveArrayToolsRaggedArrays/test/runtests.jl @@ -1,167 +1,775 @@ -using RecursiveArrayTools -using RecursiveArrayToolsRaggedArrays +using RecursiveArrayTools, RecursiveArrayToolsRaggedArrays +using RecursiveArrayToolsRaggedArrays: RaggedEnd, RaggedRange using SymbolicIndexingInterface +using SymbolicIndexingInterface: SymbolCache using Test @testset "RecursiveArrayToolsRaggedArrays" begin - @testset "RaggedVectorOfArray construction" begin - r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) - @test r.u == [[1.0, 2.0], [3.0, 4.0, 5.0]] - @test length(r) == 2 - @test ndims(r) == 2 + # =================================================================== + # Tests ported from v3 basic_indexing.jl + # =================================================================== - # From typed vector - r2 = RaggedVectorOfArray(Vector{Float64}[[1.0], [2.0, 3.0]]) - @test length(r2) == 2 + @testset "v3 basic indexing (ported)" begin + # Example Problem + recs = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] + testa = cat(recs..., dims = 2) + testva = RaggedVectorOfArray(recs) + @test maximum(testva) == maximum(maximum.(recs)) + + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + @test maximum(testva) == maximum(maximum.(recs)) + + # broadcast with array + X = rand(3, 3) + mulX = sqrt.(abs.(testva .* X)) + ref = mapreduce((x, y) -> sqrt.(abs.(x .* y)), hcat, testva, eachcol(X)) + @test mulX == ref + fill!(mulX, 0) + mulX .= sqrt.(abs.(testva .* X)) + @test mulX == ref + + @test Array(testva) == [ + 1 4 7 + 2 5 8 + 3 6 9 + ] + + @test testa[1:2, 1:2] == [1 4; 2 5] + @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 = RaggedDiffEqArray(recs, t) + @test Array(diffeq) == [ + 1 4 7 + 2 5 8 + 3 6 9 + ] + @test diffeq[1:2, 1:2] == [1 4; 2 5] end - @testset "RaggedVectorOfArray linear indexing (column-major)" begin - r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) + @testset "v3 ndims==2 indexing (ported)" begin + t = 1:10 + recs = [rand(8) for i in 1:10] + testa = cat(recs..., dims = 2) + testva = RaggedVectorOfArray(recs) + + # Array functions + @test size(testva) == (8, 10) + @test ndims(testva) == 2 + @test eltype(testva) == eltype(eltype(recs)) + testvasim = similar(testva) + @test size(testvasim) == size(testva) + @test eltype(testvasim) == eltype(testva) + testvasim = similar(testva, Float32) + @test size(testvasim) == size(testva) + @test eltype(testvasim) == Float32 + testvb = deepcopy(testva) + @test testva == testvb == recs + + # Math operations + @test testva + testvb == testva + recs == 2testva == 2 .* recs + @test testva - testvb == testva - recs == 0 .* recs + @test testva / 2 == recs ./ 2 + @test 2 .\ testva == 2 .\ recs + + # Column indexing + @test testva[:, begin] == first(testva) + @test testva[:, end] == last(testva) + @test testa[:, 1] == recs[1] + @test testva.u == recs + @test testva[:, 2:end] == RaggedVectorOfArray([recs[i] for i in 2:length(recs)]) + + diffeq = RaggedDiffEqArray(recs, t) + @test diffeq[:, 1] == testa[:, 1] + @test diffeq.u == recs + @test diffeq[:, end] == testa[:, end] + @test diffeq[:, 2:end] == RaggedDiffEqArray([recs[i] for i in 2:length(recs)], t[2:end]) + @test diffeq[:, 2:end].t == t[2:end] + @test diffeq[:, (end - 1):end] == RaggedDiffEqArray([recs[i] for i in (length(recs) - 1):length(recs)], t[(length(t) - 1):length(t)]) + @test diffeq[:, (end - 1):end].t == t[(length(t) - 1):length(t)] + @test diffeq[:, (end - 5):8] == RaggedDiffEqArray([recs[i] for i in (length(t) - 5):8], t[(length(t) - 5):8]) + @test diffeq[:, (end - 5):8].t == t[(length(t) - 5):8] + + # (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] + end - # Linear indexing goes through elements in column-major order: - # col 1: [1.0, 2.0], col 2: [3.0, 4.0, 5.0] - @test r[1] == 1.0 - @test r[2] == 2.0 - @test r[3] == 3.0 - @test r[4] == 4.0 - @test r[5] == 5.0 - @test_throws BoundsError r[6] + @testset "v3 ndims==3 indexing (ported)" begin + t = 1:15 + recs = [rand(10, 8) for i in 1:15] + testa = cat(recs..., dims = 3) + testva = RaggedVectorOfArray(recs) + diffeq = RaggedDiffEqArray(recs, t) - # Linear setindex! - r2 = copy(r) - r2[1] = 10.0 - @test r2.u[1][1] == 10.0 - r2[4] = 40.0 - @test r2.u[2][2] == 40.0 + # (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] end - @testset "RaggedVectorOfArray N-ary indexing — no zero-padding" begin - r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0]]) + @testset "v3 ragged arrays (ported)" begin + t = 1:3 + recs = [[1, 2, 3], [3, 5, 6, 7], [8, 9, 10, 11]] + testva = RaggedVectorOfArray(recs) + diffeq = RaggedDiffEqArray(recs, t) + + @test testva[:, 1] == recs[1] + @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] + @test diffeq[:, 1:2] == RaggedDiffEqArray([recs[i] for i in 1:2], t[1:2]) + @test diffeq[:, 1:2].t == t[1:2] + @test diffeq[:, 2:end] == RaggedDiffEqArray([recs[i] for i in 2:3], t[2:end]) + @test diffeq[:, 2:end].t == t[2:end] + @test diffeq[:, (end - 1):end] == RaggedDiffEqArray([recs[i] for i in (length(recs) - 1):length(recs)], t[(length(t) - 1):length(t)]) + @test diffeq[:, (end - 1):end].t == t[(length(t) - 1):length(t)] + end - # A[:, i] returns i-th inner array directly (NO zero-padding) - @test r[:, 1] == [1.0, 2.0] - @test r[:, 2] == [3.0, 4.0, 5.0] - @test r[:, 3] == [6.0] - @test length(r[:, 1]) == 2 - @test length(r[:, 2]) == 3 - @test length(r[:, 3]) == 1 + @testset "v3 heterogeneous views (ported, issue #453)" begin + f = RaggedVectorOfArray([[1.0], [2.0, 3.0]]) + @test length(view(f, :, 1)) == 1 + @test length(view(f, :, 2)) == 2 + @test view(f, :, 1) == [1.0] + @test view(f, :, 2) == [2.0, 3.0] + @test collect(view(f, :, 1)) == f[:, 1] + @test collect(view(f, :, 2)) == f[:, 2] - # A[j, i] returns component (no zero-padding, throws on out-of-bounds) - @test r[1, 1] == 1.0 - @test r[2, 1] == 2.0 - @test r[3, 2] == 5.0 - @test r[1, 3] == 6.0 - @test_throws BoundsError r[3, 1] # only 2 elements in first array - @test_throws BoundsError r[2, 3] # only 1 element in third array + f2 = RaggedVectorOfArray([[1.0, 2.0], [3.0]]) + @test length(view(f2, :, 1)) == 2 + @test length(view(f2, :, 2)) == 1 + @test view(f2, :, 1) == [1.0, 2.0] + @test view(f2, :, 2) == [3.0] + @test collect(view(f2, :, 1)) == f2[:, 1] + @test collect(view(f2, :, 2)) == f2[:, 2] + end - # A[j, :] returns time series of j-th component - @test r[1, :] == [1.0, 3.0, 6.0] + @testset "v3 end indexing with ragged arrays (ported)" begin + ragged = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0]]) + @test ragged[end, 1] == 2.0 + @test ragged[end, 2] == 5.0 + @test ragged[end, 3] == 9.0 + @test ragged[end - 1, 1] == 1.0 + @test ragged[end - 1, 2] == 4.0 + @test ragged[end - 1, 3] == 8.0 + @test ragged[1:end, 1] == [1.0, 2.0] + @test ragged[1:end, 2] == [3.0, 4.0, 5.0] + @test ragged[1:end, 3] == [6.0, 7.0, 8.0, 9.0] + @test ragged[:, end] == [6.0, 7.0, 8.0, 9.0] + @test ragged[:, 2:end] == RaggedVectorOfArray(ragged.u[2:end]) + @test ragged[:, (end - 1):end] == RaggedVectorOfArray(ragged.u[(end - 1):end]) - # A[:, idx_array] returns subset - r_sub = r[:, [1, 3]] - @test r_sub isa RaggedVectorOfArray - @test length(r_sub) == 2 - @test r_sub[:, 1] == [1.0, 2.0] - @test r_sub[:, 2] == [6.0] + ragged2 = RaggedVectorOfArray([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]]) + @test ragged2[end, 1] == 4.0 + @test ragged2[end, 2] == 6.0 + @test ragged2[end, 3] == 9.0 + @test ragged2[end - 1, 1] == 3.0 + @test ragged2[end - 1, 2] == 5.0 + @test ragged2[end - 1, 3] == 8.0 + @test ragged2[end - 2, 1] == 2.0 + @test ragged2[1:end, 1] == [1.0, 2.0, 3.0, 4.0] + @test ragged2[1:end, 2] == [5.0, 6.0] + @test ragged2[1:end, 3] == [7.0, 8.0, 9.0] + @test ragged2[2:end, 1] == [2.0, 3.0, 4.0] + @test ragged2[2:end, 2] == [6.0] + @test ragged2[2:end, 3] == [8.0, 9.0] + @test ragged2[:, end] == [7.0, 8.0, 9.0] + @test ragged2[:, 2:end] == RaggedVectorOfArray(ragged2.u[2:end]) + @test ragged2[1:(end - 1), 1] == [1.0, 2.0, 3.0] + @test ragged2[1:(end - 1), 2] == [5.0] + @test ragged2[1:(end - 1), 3] == [7.0, 8.0] + @test ragged2[:, (end - 1):end] == RaggedVectorOfArray(ragged2.u[(end - 1):end]) + end - # A[:, bool_array] - r_bool = r[:, [true, false, true]] - @test r_bool isa RaggedVectorOfArray - @test length(r_bool) == 2 + @testset "v3 RaggedEnd/RaggedRange broadcast as scalars (ported)" begin + ragged = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0]]) + ragged_idx = lastindex(ragged, 1) + @test ragged_idx isa RaggedEnd + @test Base.broadcastable(ragged_idx) isa Ref + # Broadcasting with RaggedEnd should not error + @test identity.(ragged_idx) === ragged_idx + + ragged_range_idx = 1:lastindex(ragged, 1) + @test ragged_range_idx isa RaggedRange + @test Base.broadcastable(ragged_range_idx) isa Ref + # Broadcasting with RaggedRange should not error + @test identity.(ragged_range_idx) === ragged_range_idx + end - # A[:, :] - r_all = r[:, :] - @test r_all isa RaggedVectorOfArray - @test r_all == r - @test r_all.u !== r.u # copy + @testset "v3 broadcast assignment heterogeneous (ported, issue #454)" begin + u = RaggedVectorOfArray([[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] end - @testset "RaggedVectorOfArray setindex!" begin - r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) - r[1, 1] = 10.0 - @test r[1, 1] == 10.0 + @testset "v3 DiffEqArray with 2D inner arrays (ported)" begin + t = 1:2 + recs_2d = [rand(2, 3), rand(2, 4)] + diffeq_2d = RaggedDiffEqArray(recs_2d, t) + @test diffeq_2d[:, 1] == recs_2d[1] + @test diffeq_2d[:, 2] == recs_2d[2] + @test diffeq_2d[:, 1:2] == RaggedDiffEqArray(recs_2d[1:2], t[1:2]) + @test diffeq_2d[:, 1:2].t == t[1:2] + @test diffeq_2d[:, 2:end] == RaggedDiffEqArray(recs_2d[2:end], t[2:end]) + @test diffeq_2d[:, 2:end].t == t[2:end] + @test diffeq_2d[:, (end - 1):end] == RaggedDiffEqArray(recs_2d[(end - 1):end], t[(end - 1):end]) + @test diffeq_2d[:, (end - 1):end].t == t[(end - 1):end] + end - r[:, 2] = [30.0, 40.0, 50.0] - @test r[:, 2] == [30.0, 40.0, 50.0] + @testset "v3 DiffEqArray with 3D inner arrays (ported)" begin + t = 1:2 + recs_3d = [rand(2, 3, 4), rand(2, 3, 5)] + diffeq_3d = RaggedDiffEqArray(recs_3d, t) + @test diffeq_3d[:, :, :, 1] == recs_3d[1] + @test diffeq_3d[:, :, :, 2] == recs_3d[2] + @test diffeq_3d[:, :, :, 1:2] == RaggedDiffEqArray(recs_3d[1:2], t[1:2]) + @test diffeq_3d[:, :, :, 1:2].t == t[1:2] + @test diffeq_3d[:, :, :, 2:end] == RaggedDiffEqArray(recs_3d[2:end], t[2:end]) + @test diffeq_3d[:, :, :, 2:end].t == t[2:end] + @test diffeq_3d[:, :, :, (end - 1):end] == RaggedDiffEqArray(recs_3d[(end - 1):end], t[(end - 1):end]) + @test diffeq_3d[:, :, :, (end - 1):end].t == t[(end - 1):end] end - @testset "RaggedVectorOfArray iteration" begin - r = RaggedVectorOfArray([[1, 2], [3, 4, 5]]) - collected = collect(r) - @test collected == [[1, 2], [3, 4, 5]] - @test first(r) == [1, 2] - @test last(r) == [3, 4, 5] + @testset "v3 2D inner arrays ragged second dimension (ported)" begin + u = RaggedVectorOfArray([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] + # test scalar indexing with end + @test u[1, 1, end] == u.u[end][1, 1] + @test u[1, end, end] == u.u[end][1, end] + @test u[1, 2:end, end] == vec(u.u[end][1, 2:end]) end - @testset "RaggedVectorOfArray copy, zero, similar, fill!" begin - r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) + @testset "v3 3D inner arrays ragged third dimension (ported)" begin + u = RaggedVectorOfArray([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 u[:, :, end] == u.u[end] + @test u[:, :, 2:end] == RaggedVectorOfArray(u.u[2:end]) + @test u[1, 1, end] == u.u[end][1, 1, end] + @test u[end, 1, end] == u.u[end][end, 1, end] + end - r2 = copy(r) - @test r2 == r - @test r2.u !== r.u + @testset "v3 views can be modified (ported)" begin + f3 = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) + v = view(f3, :, 2) + @test length(v) == 3 + v[1] = 10.0 + @test f3[1, 2] == 10.0 + @test f3.u[2][1] == 10.0 + end - r0 = zero(r) - @test r0[:, 1] == [0.0, 0.0] - @test r0[:, 2] == [0.0, 0.0, 0.0] + @testset "v3 2D inner array conversions (ported)" begin + t = 1:5 + recs = [rand(2, 2) for i in 1:5] + testva = RaggedVectorOfArray(recs) + diffeq = RaggedDiffEqArray(recs, t) + + @test Array(testva) isa Array{Float64, 3} + @test Array(diffeq) isa Array{Float64, 3} + end - rs = similar(r) - @test length(rs[:, 1]) == 2 - @test length(rs[:, 2]) == 3 + @testset "v3 CartesianIndex (ported)" begin + v = RaggedVectorOfArray([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 - fill!(r0, 7.0) - @test r0[:, 1] == [7.0, 7.0] - @test r0[:, 2] == [7.0, 7.0, 7.0] + v = RaggedDiffEqArray([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 end - @testset "RaggedVectorOfArray push!, append!" begin - r = RaggedVectorOfArray([[1.0], [2.0, 3.0]]) - push!(r, [4.0, 5.0, 6.0]) - @test length(r) == 3 - @test r[:, 3] == [4.0, 5.0, 6.0] + @testset "v3 heterogeneous broadcasting (ported)" begin + v = RaggedVectorOfArray([rand(20), rand(10, 10), rand(3, 3, 3)]) + w = v .* v + @test w isa RaggedVectorOfArray + @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 w isa RaggedVectorOfArray + @test w.u == map(x -> x .+ 1, v.u) + + v = RaggedDiffEqArray([rand(20), rand(10, 10), rand(3, 3, 3)], 1:3) + w = v .* v + @test_broken w isa RaggedDiffEqArray # 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 RaggedDiffEqArray # FIXME + @test w.u == map(x -> x .+ 1, v.u) + end - r2 = RaggedVectorOfArray([[7.0]]) - append!(r, r2) - @test length(r) == 4 - @test r[:, 4] == [7.0] + @testset "v3 setindex! (ported)" begin + testva = RaggedVectorOfArray([i * ones(3, 3) for i in 1:5]) + testva[:, 2] = 7ones(3, 3) + @test testva[:, 2] == 7ones(3, 3) + testva[:, :] = [2i * ones(3, 3) for i in 1:5] + for i in 1:5 + @test testva[:, i] == 2i * ones(3, 3) + end + testva[:, 1:2:5] = [5i * ones(3, 3) for i in 1:2:5] + for i in 1:2:5 + @test testva[:, i] == 5i * ones(3, 3) + end + testva[CartesianIndex(3, 3, 5)] = 64.0 + @test testva[:, 5][3, 3] == 64.0 + @test_throws ArgumentError testva[2, 1:2, :] = 108.0 + testva[2, 1:2, :] .= 108.0 + for i in 1:5 + @test all(testva[:, i][2, 1:2] .== 108.0) + end + testva[:, 3, :] = [3i / 7j for i in 1:3, j in 1:5] + for j in 1:5 + for i in 1:3 + @test testva[i, 3, j] == 3i / 7j + end + end end - @testset "RaggedVectorOfArray broadcasting" begin - r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) + @testset "v3 edge cases DiffEqArray scalar broadcast (ported)" begin + x = [1, 2, 3, 4, 5, 6, 7, 8, 9] + testva = RaggedDiffEqArray(x, x) + testvb = RaggedDiffEqArray(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 + end - # Scalar broadcast - r2 = r .* 2.0 - @test r2 isa RaggedVectorOfArray - @test r2[:, 1] == [2.0, 4.0] - @test r2[:, 2] == [6.0, 8.0, 10.0] + @testset "v3 multidimensional parent VectorOfArray (ported)" begin + u_matrix = RaggedVectorOfArray([[1, 2] for i in 1:2, j in 1:3]) + u_vector = RaggedVectorOfArray([[1, 2] for i in 1:6]) + + # test broadcasting + function foo!(u) + @. u += 2 * u * abs(u) + return u + end + foo!(u_matrix) + foo!(u_vector) + @test all(u_matrix .== [3, 10]) + @test all(vec(u_matrix) .≈ vec(u_vector)) + + # test that, for VectorOfArray with multi-dimensional parent arrays, + # broadcast and `similar` preserve the structure of the parent array + @test typeof(parent(similar(u_matrix))) == typeof(parent(u_matrix)) + @test typeof(parent((x -> x).(u_matrix))) == typeof(parent(u_matrix)) + + # test efficiency (allow 1 alloc on Julia 1.10 LTS) + num_allocs = @allocations foo!(u_matrix) + @test num_allocs <= 1 + end + + @testset "v3 issue 354 (ported)" begin + @test RaggedVectorOfArray(ones(1))[:] == ones(1) + end - # Element-wise broadcast between ragged arrays of same structure - r3 = r .+ r - @test r3 isa RaggedVectorOfArray - @test r3[:, 1] == [2.0, 4.0] - @test r3[:, 2] == [6.0, 8.0, 10.0] + # =================================================================== + # Tests ported from v3 interface_tests.jl + # =================================================================== - # In-place broadcast - r4 = copy(r) - r4 .= r .+ 1.0 - @test r4[:, 1] == [2.0, 3.0] - @test r4[:, 2] == [4.0, 5.0, 6.0] + @testset "v3 iteration (ported)" begin + t = 1:3 + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + testda = RaggedDiffEqArray([[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 + end + + @testset "v3 push! and copy (ported)" begin + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + testda = RaggedDiffEqArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], 1:3) + + 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] + end + + @testset "v3 append! (ported)" begin + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]) + testda = RaggedDiffEqArray([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], 1:4) + + 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] + end + + @testset "v3 push! making array ragged (ported)" begin + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]]) + testda = RaggedDiffEqArray([[1, 2, 3], [4, 5, 6], [7, 8, 9], [10, 11, 12]], 1:4) + + # After appending, add ragged element + append!(testva, testva) + append!(testda, testda) + + push!(testva, [-1, -2, -3, -4]) + push!(testda, [-1, -2, -3, -4]) + @test testva[1:2, 5:6] == [1 4; 2 5] + @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] + end + + @testset "v3 ndims (ported)" begin + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + @test ndims(testva) == 2 + @test ndims(typeof(testva)) == 2 + end + + @testset "v3 push! shape enforcement (ported)" begin + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + testda = RaggedDiffEqArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], 1:3) + + @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]) + end + + @testset "v3 type inference (ported)" begin + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + @inferred sum(testva) + @inferred sum(RaggedVectorOfArray([RaggedVectorOfArray([zeros(4, 4)])])) + @inferred mapreduce(string, *, testva) + # Type stability for `end` indexing (issue #525) + testva_end = RaggedVectorOfArray(fill(fill(2.0, 2), 10)) + # Use lastindex directly since `end` doesn't work in SafeTestsets + last_col = lastindex(testva_end, 2) + @inferred testva_end[1, last_col] + @inferred testva_end[1, 1:last_col] + @test testva_end[1, last_col] == 2.0 + last_row = lastindex(testva_end, 1) + @inferred testva_end[last_row, 1] + @inferred testva_end[1:last_row, 1] + @test testva_end[last_row, 1] == 2.0 + end + + @testset "v3 mapreduce (ported)" begin + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + @test mapreduce(x -> string(x) * "q", *, testva) == "1q2q3q4q5q6q7q8q9q" + + testvb = RaggedVectorOfArray([rand(1:10, 3, 3, 3) for _ in 1:4]) + arrvb = Array(testvb) + for i in 1:ndims(arrvb) + @test sum(arrvb; dims = i) == sum(testvb; dims = i) + @test prod(arrvb; dims = i) == prod(testvb; dims = i) + @test mapreduce(string, *, arrvb; dims = i) == mapreduce(string, *, testvb; dims = i) + end + end - # .= zero(r) should work (the issue from JoshuaLampert) - r5 = copy(r) - r5 .= zero(r5) - @test r5[:, 1] == [0.0, 0.0] - @test r5[:, 2] == [0.0, 0.0, 0.0] + @testset "v3 mapreduce ndims==1 (ported)" begin + testvb = RaggedVectorOfArray(collect(1.0:0.1:2.0)) + arrvb = Array(testvb) + @test sum(arrvb) == sum(testvb) + @test prod(arrvb) == prod(testvb) + @test mapreduce(string, *, arrvb) == mapreduce(string, *, testvb) end - @testset "VectorOfArray ragged .= zero also works" begin - # This is the exact bug from JoshuaLampert's report - f = VectorOfArray([[1.0], [2.0, 3.0]]) - f .= zero(f) - @test f.u[1] == [0.0] - @test f.u[2] == [0.0, 0.0] + @testset "v3 view (ported)" begin + testvc = RaggedVectorOfArray([rand(1:10, 3, 3) for _ in 1:3]) + arrvc = Array(testvc) + for idxs in [ + (2, 2, :), (2, :, 2), (:, 2, 2), (:, :, 2), (:, 2, :), (2, :, :), (:, :, :), + (1:2, 1:2, Bool[1, 0, 1]), (1:2, Bool[1, 0, 1], 1:2), (Bool[1, 0, 1], 1:2, 1:2), + ] + arr_view = view(arrvc, idxs...) + voa_view = view(testvc, idxs...) + @test size(arr_view) == size(voa_view) + @test all(arr_view .== voa_view) + end + + testvc = RaggedVectorOfArray(collect(1:10)) + arrvc = Array(testvc) + bool_idx = rand(Bool, 10) + for (voaidx, arridx) in [ + ((:,), (:,)), + ((3:5,), (3:5,)), + ((bool_idx,), (bool_idx,)), + ] + arr_view = view(arrvc, arridx...) + voa_view = view(testvc.u, voaidx...) + @test size(arr_view) == size(voa_view) + @test all(arr_view .== voa_view) + end end - @testset "RaggedVectorOfArray ↔ VectorOfArray conversion" begin + @testset "v3 stack (ported)" begin + testva = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) + @test stack(testva) == [1 4 7; 2 5 8; 3 6 9] + @test stack(testva; dims = 1) == [1 2 3; 4 5 6; 7 8 9] + + testva = RaggedVectorOfArray([RaggedVectorOfArray([ones(2, 2), 2ones(2, 2)]), 3ones(2, 2, 2)]) + @test stack(testva) == + [1.0 1.0; 1.0 1.0;;; 2.0 2.0; 2.0 2.0;;;; 3.0 3.0; 3.0 3.0;;; 3.0 3.0; 3.0 3.0] + end + + @testset "v3 convert Array (ported)" begin + t = 1:8 + recs = [rand(10, 7) for i in 1:8] + testva = RaggedVectorOfArray(recs) + testda = RaggedDiffEqArray(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 = RaggedVectorOfArray(recs) + testda = RaggedDiffEqArray(recs, t) + + @test size(convert(Array, testva)) == (2, 2, 3) + @test size(convert(Array, testda)) == (2, 2, 3) + end + + @testset "v3 similar (ported)" begin + recs = [rand(6) for i in 1:4] + testva = RaggedVectorOfArray(recs) + + testva2 = similar(testva) + @test typeof(testva2) == typeof(testva) + @test size(testva2) == size(testva) + + testva3 = similar(testva, 10) + @test typeof(testva3) == typeof(testva) + @test length(testva3) == 10 + end + + @testset "v3 fill! and all (ported)" begin + recs = [rand(6) for i in 1:4] + testva = RaggedVectorOfArray(recs) + + testval = 3.0 + testva2 = similar(testva) + fill!(testva2, testval) + @test all(x -> (x == testval), testva2) + testts = rand(Float64, size(testva.u)) + testda = RaggedDiffEqArray(recursivecopy(testva.u), testts) + fill!(testda, testval) + @test all(x -> (x == testval), testda) + end + + @testset "v3 copyto! (ported)" begin + testva = RaggedVectorOfArray(collect(0.1:0.1:1.0)) + arr = 0.2:0.2:2.0 + copyto!(testva, arr) + @test Array(testva) == arr + testva = RaggedVectorOfArray([i * ones(3, 2) for i in 1:4]) + arr = rand(3, 2, 4) + copyto!(testva, arr) + @test Array(testva) == arr + testva = RaggedVectorOfArray( + [ + ones(3, 2, 2), + RaggedVectorOfArray( + [ + 2ones(3, 2), + RaggedVectorOfArray([3ones(3), 4ones(3)]), + ] + ), + RaggedDiffEqArray( + [ + 5ones(3, 2), + RaggedVectorOfArray([6ones(3), 7ones(3)]), + ], + [0.1, 0.2], + [100.0, 200.0], + SymbolCache([:x, :y], [:a, :b], :t) + ), + ] + ) + arr = rand(3, 2, 2, 3) + copyto!(testva, arr) + @test Array(testva) == arr + # ensure structure and fields are maintained + @test testva.u[1] isa Array + @test testva.u[2] isa RaggedVectorOfArray + @test testva.u[2].u[2] isa RaggedVectorOfArray + @test testva.u[3] isa RaggedDiffEqArray + @test testva.u[3].u[2] isa RaggedVectorOfArray + @test testva.u[3].t == [0.1, 0.2] + @test testva.u[3].p == [100.0, 200.0] + @test testva.u[3].sys isa SymbolCache + end + + @testset "v3 any (ported)" begin + recs = [collect(1:5), collect(6:10), collect(11:15)] + testts = rand(5) + testva = RaggedVectorOfArray(recs) + testda = RaggedDiffEqArray(recs, testts) + testval1 = 4 + testval2 = 17 + @test any(x -> (x == testval1), testva) + @test !any(x -> (x == testval2), testda) + end + + @testset "v3 empty creation (ported)" begin + emptyva = RaggedVectorOfArray(Array{Vector{Float64}}([])) + @test isempty(emptyva) + emptyda = RaggedDiffEqArray(Array{Vector{Float64}}([]), Vector{Float64}()) + @test isempty(emptyda) + end + + @testset "v3 map (ported)" begin + A = RaggedVectorOfArray(map(i -> rand(2, 4), 1:7)) + @test map(x -> maximum(x), A) isa Vector + + DA = RaggedDiffEqArray(map(i -> rand(2, 4), 1:7), 1:7) + @test map(x -> maximum(x), DA) isa Vector + end + + @testset "v3 zero and resize! (ported)" begin + u = RaggedVectorOfArray([fill(2.0, 2), ones(2)]) + @test typeof(zero(u)) <: typeof(u) + resize!(u, 3) + @test pointer(u) === pointer(u.u) + end + + @testset "v3 DiffEqArray constructor ambiguity (ported, issue SciMLBase#889)" begin + darr = RaggedDiffEqArray([[1.0, 1.0]], [1.0], ()) + @test darr.p == () + @test darr.sys === nothing + @test ndims(darr) == 2 + darr = RaggedDiffEqArray([[1.0, 1.0]], [1.0], (), "A") + @test darr.p == () + @test darr.sys == "A" + @test ndims(darr) == 2 + darr = RaggedDiffEqArray([ones(2, 2)], [1.0], (1, 1, 1)) + @test darr.p == (1, 1, 1) + @test darr.sys === nothing + @test ndims(darr) == 3 + darr = RaggedDiffEqArray([ones(2, 2)], [1.0], (1, 1, 1), "A") + @test darr.p == (1, 1, 1) + @test darr.sys == "A" + @test ndims(darr) == 3 + end + + @testset "v3 system retained in 4-argument constructor (ported)" begin + darr = RaggedDiffEqArray([ones(2)], [1.0], :params, :sys) + @test darr.sys == :sys + end + + # issparse test skipped: ragged types are not AbstractArray, so + # SparseArrays.issparse is not defined for them (issue #486 only + # applies to AbstractVectorOfArray which is now AbstractArray). + + # =================================================================== + # Additional tests for new v4 functionality (interp, conversion, SII) + # =================================================================== + + @testset "RaggedVectorOfArray construction (typed)" begin + r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) + @test r.u == [[1.0, 2.0], [3.0, 4.0, 5.0]] + @test length(r) == 2 + @test ndims(r) == 2 + + # From typed vector + r2 = RaggedVectorOfArray(Vector{Float64}[[1.0], [2.0, 3.0]]) + @test length(r2) == 2 + end + + # Conversion, SII, interp/dense, and type hierarchy tests below + + @testset "RaggedVectorOfArray <-> VectorOfArray conversion" begin r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) # To VectorOfArray (rectangular, zero-padded) @@ -195,12 +803,6 @@ using Test @test rr[:, 3] == [1.0] end - @testset "RaggedVectorOfArray inner_sizes/inner_lengths" begin - r = RaggedVectorOfArray([[1.0, 2.0], [3.0], [4.0, 5.0, 6.0]]) - @test inner_lengths(r) == [2, 1, 3] - @test inner_sizes(r) == [(2,), (1,), (3,)] - end - @testset "RaggedDiffEqArray construction" begin r = RaggedDiffEqArray([[1.0, 2.0], [3.0, 4.0, 5.0]], [0.0, 1.0]) @test r.u == [[1.0, 2.0], [3.0, 4.0, 5.0]] @@ -221,10 +823,9 @@ using Test @test r[2, 2] == 4.0 @test_throws BoundsError r[3, 1] - # Linear indexing (column-major) - @test r[1] == 1.0 - @test r[2] == 2.0 - @test r[3] == 3.0 + # Linear indexing returns inner arrays (v3 behavior) + @test r[1] == [1.0, 2.0] + @test r[2] == [3.0, 4.0, 5.0] # Colon-colon preserves DiffEqArray type r_all = r[:, :] @@ -238,7 +839,7 @@ using Test @test r_sub[:, 1] == [3.0, 4.0, 5.0] end - @testset "RaggedDiffEqArray ↔ DiffEqArray conversion" begin + @testset "RaggedDiffEqArray <-> DiffEqArray conversion" begin r = RaggedDiffEqArray([[1.0, 2.0], [3.0, 4.0, 5.0]], [0.0, 1.0]) da = DiffEqArray(r) @@ -299,7 +900,7 @@ using Test @test r[:b] == [2.0, 4.0] # Parameter indexing should throw - @test_throws RecursiveArrayToolsRaggedArrays.RaggedParameterIndexingError r[:p] + @test_throws RecursiveArrayToolsRaggedArrays.ParameterIndexingError r[:p] end @testset "2D inner arrays" begin @@ -345,175 +946,6 @@ using Test @test r4.dense == true end - # =================================================================== - # Tests ported from v3 VectorOfArray ragged behavior - # These were the standard ragged tests before v4's AbstractArray - # subtyping changed VectorOfArray to zero-pad. The ragged sublibrary - # preserves the original non-zero-padded behavior. - # =================================================================== - - @testset "v3 ragged indexing (ported)" begin - recs = [[1, 2, 3], [3, 5, 6, 7], [8, 9, 10, 11]] - r = RaggedVectorOfArray(recs) - rd = RaggedDiffEqArray(recs, 1:3) - - # Colon indexing returns actual inner array (no zero-padding) - @test r[:, 1] == recs[1] - @test rd[:, 1] == recs[1] - - # Subarray indexing into compatible region - @test r[1, 1] == 1 - @test r[2, 1] == 2 - @test r[1, 2] == 3 - @test r[2, 2] == 5 - - # DiffEqArray column slicing preserves time - rd_sub = rd[:, 1:2] - @test rd_sub isa RaggedDiffEqArray - @test rd_sub.t == [1, 2] - rd_sub2 = rd[:, 2:3] - @test rd_sub2 isa RaggedDiffEqArray - @test rd_sub2.t == [2, 3] - end - - @testset "v3 heterogeneous views (ported, issue #453)" begin - f = RaggedVectorOfArray([[1.0], [2.0, 3.0]]) - # Column access respects actual inner array size - @test length(f[:, 1]) == 1 - @test length(f[:, 2]) == 2 - @test f[:, 1] == [1.0] - @test f[:, 2] == [2.0, 3.0] - - # view also respects actual inner array size - @test length(view(f, :, 1)) == 1 - @test length(view(f, :, 2)) == 2 - @test view(f, :, 1) == [1.0] - @test view(f, :, 2) == [2.0, 3.0] - @test collect(view(f, :, 1)) == f[:, 1] - @test collect(view(f, :, 2)) == f[:, 2] - - f2 = RaggedVectorOfArray([[1.0, 2.0], [3.0]]) - @test length(f2[:, 1]) == 2 - @test length(f2[:, 2]) == 1 - @test f2[:, 1] == [1.0, 2.0] - @test f2[:, 2] == [3.0] - @test length(view(f2, :, 1)) == 2 - @test length(view(f2, :, 2)) == 1 - @test view(f2, :, 1) == [1.0, 2.0] - @test view(f2, :, 2) == [3.0] - end - - @testset "v3 end indexing with ragged arrays (ported)" begin - ragged = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0]]) - # A[j, i] accesses the j-th element of i-th inner array - @test ragged[2, 1] == 2.0 - @test ragged[3, 2] == 5.0 - @test ragged[4, 3] == 9.0 - @test ragged[1, 1] == 1.0 - @test ragged[2, 2] == 4.0 - @test ragged[3, 3] == 8.0 - - # `end` in the last (column) dimension - @test ragged[:, end] == [6.0, 7.0, 8.0, 9.0] - @test ragged[:, (end - 1):end] isa RaggedVectorOfArray - @test length(ragged[:, (end - 1):end]) == 2 - - # `end` in first dim = max inner array length (4 here) - # So ragged[end, 1] = ragged[4, 1] which is out of bounds for inner array 1 - @test_throws BoundsError ragged[end, 1] # inner[1] has only 2 elements - @test_throws BoundsError ragged[end, 2] # inner[2] has only 3 elements - @test ragged[end, 3] == 9.0 # inner[3] has 4 elements, end=4 works - - # Range indexing into inner arrays - @test ragged[1:2, 1] == [1.0, 2.0] - @test ragged[1:3, 2] == [3.0, 4.0, 5.0] - @test ragged[1:4, 3] == [6.0, 7.0, 8.0, 9.0] - # 1:end uses max inner length (4), so only works for longest column - @test_throws BoundsError ragged[1:end, 1] # 1:4 but inner[1] has 2 - @test_throws BoundsError ragged[1:end, 2] # 1:4 but inner[2] has 3 - @test ragged[1:end, 3] == [6.0, 7.0, 8.0, 9.0] # 1:4 matches inner[3] - - # Colon returns actual arrays - @test ragged[:, 1] == [1.0, 2.0] - @test ragged[:, 2] == [3.0, 4.0, 5.0] - @test ragged[:, 3] == [6.0, 7.0, 8.0, 9.0] - - # Subset selection via range - @test ragged[:, 2:end] isa RaggedVectorOfArray - @test ragged[:, 2:end][:, 1] == [3.0, 4.0, 5.0] - - ragged2 = RaggedVectorOfArray([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]]) - # end in first dim = max length = 4 - @test ragged2[end, 1] == 4.0 # inner[1] has 4 elements - @test_throws BoundsError ragged2[end, 2] # inner[2] has only 2 - @test_throws BoundsError ragged2[end, 3] # inner[3] has only 3 - @test ragged2[end - 1, 1] == 3.0 # end-1 = 3 - @test_throws BoundsError ragged2[end - 1, 2] # 3 > length([5,6]) - @test ragged2[end - 1, 3] == 9.0 # 3 <= length([7,8,9]) - @test ragged2[:, 1] == [1.0, 2.0, 3.0, 4.0] - @test ragged2[:, 2] == [5.0, 6.0] - @test ragged2[:, 3] == [7.0, 8.0, 9.0] - # Explicit range indexing (not using end) - @test ragged2[1:4, 1] == [1.0, 2.0, 3.0, 4.0] - @test ragged2[1:2, 2] == [5.0, 6.0] - @test ragged2[1:3, 3] == [7.0, 8.0, 9.0] - @test ragged2[2:4, 1] == [2.0, 3.0, 4.0] - @test ragged2[2:2, 2] == [6.0] - @test ragged2[2:3, 3] == [8.0, 9.0] - # end in last dim works fine - @test ragged2[:, end] == [7.0, 8.0, 9.0] - @test ragged2[:, 2:end] isa RaggedVectorOfArray - # end in first dim = 4 (max), 1:(end-1) = 1:3 - @test ragged2[1:(end - 1), 1] == [1.0, 2.0, 3.0] - @test_throws BoundsError ragged2[1:(end - 1), 2] # 1:3 but inner[2] has 2 - @test ragged2[1:(end - 1), 3] == [7.0, 8.0, 9.0] # 1:3 matches inner[3] - end - - @testset "v3 push! making array ragged (ported)" begin - r = RaggedVectorOfArray([[1, 2, 3], [4, 5, 6]]) - push!(r, [-1, -2, -3, -4]) - - # Can still index compatible region - @test r[1, 1] == 1 - @test r[2, 1] == 2 - @test r[1, 2] == 4 - @test r[2, 2] == 5 - - # Out-of-bounds on shorter arrays throws - @test_throws BoundsError r[4, 1] - @test_throws BoundsError r[4, 2] - - # Full column access of the new ragged element - @test r[:, 3] == [-1, -2, -3, -4] - @test length(r) == 3 - end - - @testset "v3 broadcast assignment (ported, issue #454)" begin - u = RaggedVectorOfArray([[1.0], [2.0, 3.0]]) - @test length(u[:, 1]) == 1 - @test length(u[:, 2]) == 2 - - # Broadcast assignment into a column - u[:, 2] = [10.0, 11.0] - @test u.u[2] == [10.0, 11.0] - end - - @testset "v3 DiffEqArray 2D ragged inner arrays (ported)" begin - recs_2d = [rand(2, 3), rand(2, 4)] - rd = RaggedDiffEqArray(recs_2d, 1:2) - @test rd[:, 1] == recs_2d[1] - @test rd[:, 2] == recs_2d[2] - @test size(rd[:, 1]) == (2, 3) - @test size(rd[:, 2]) == (2, 4) - - # Subset preserves time - rd_sub = rd[:, [1, 2]] - @test rd_sub isa RaggedDiffEqArray - @test rd_sub.t == [1, 2] - @test rd_sub[:, 1] == recs_2d[1] - @test rd_sub[:, 2] == recs_2d[2] - end - @testset "v3 zero and fill on ragged (ported)" begin r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) diff --git a/test/runtests.jl b/test/runtests.jl index b486f64a..7bdb76c8 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -53,26 +53,25 @@ end @time @safetestset "SymbolicIndexingInterface API test" include("symbolic_indexing_interface_test.jl") end - if GROUP == "Subpackages" || GROUP == "All" - # Test that loading RecursiveArrayToolsArrayPartitionAnyAll overrides any/all - Pkg.develop( - PackageSpec( - path = joinpath(dirname(@__DIR__), "lib", "RecursiveArrayToolsArrayPartitionAnyAll") - ) - ) - @time @safetestset "ArrayPartition AnyAll Subpackage" begin - using RecursiveArrayTools, RecursiveArrayToolsArrayPartitionAnyAll, Test - # Verify optimized methods are active - m_any = which(any, Tuple{Function, ArrayPartition}) - m_all = which(all, Tuple{Function, ArrayPartition}) - @test occursin("ArrayPartitionAnyAll", string(m_any.module)) - @test occursin("ArrayPartitionAnyAll", string(m_all.module)) - # Verify correctness - @test any(isnan, ArrayPartition([NaN], [1.0])) - @test !any(isnan, ArrayPartition([1.0], [2.0])) - @test all(isnan, ArrayPartition([NaN], [NaN])) - @test !all(isnan, ArrayPartition([NaN], [1.0])) - end + if GROUP == "RaggedArrays" + Pkg.activate(joinpath(dirname(@__DIR__), "lib", "RecursiveArrayToolsRaggedArrays")) + Pkg.develop(PackageSpec(path = dirname(@__DIR__))) + Pkg.instantiate() + Pkg.test() + end + + if GROUP == "ArrayPartitionAnyAll" + Pkg.activate(joinpath(dirname(@__DIR__), "lib", "RecursiveArrayToolsArrayPartitionAnyAll")) + Pkg.develop(PackageSpec(path = dirname(@__DIR__))) + Pkg.instantiate() + Pkg.test() + end + + if GROUP == "ShorthandConstructors" + Pkg.activate(joinpath(dirname(@__DIR__), "lib", "RecursiveArrayToolsShorthandConstructors")) + Pkg.develop(PackageSpec(path = dirname(@__DIR__))) + Pkg.instantiate() + Pkg.test() end if GROUP == "Downstream"