diff --git a/Project.toml b/Project.toml index 18c4c5d8..f4f16060 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "RecursiveArrayTools" uuid = "731186ca-8d62-57ce-b412-fbd966d074cd" +version = "4.0.0" authors = ["Chris Rackauckas "] -version = "3.51.0" [deps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" @@ -15,6 +15,7 @@ StaticArraysCore = "1e83bf80-4336-4d27-bf5d-d5a4f845583c" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" [weakdeps] +CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -29,6 +30,7 @@ Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [extensions] +RecursiveArrayToolsCUDAExt = "CUDA" RecursiveArrayToolsFastBroadcastExt = "FastBroadcast" RecursiveArrayToolsForwardDiffExt = "ForwardDiff" RecursiveArrayToolsKernelAbstractionsExt = "KernelAbstractions" @@ -46,6 +48,7 @@ RecursiveArrayToolsZygoteExt = "Zygote" Adapt = "4" Aqua = "0.8" ArrayInterface = "7.16" +CUDA = "5" DocStringExtensions = "0.9.3" FastBroadcast = "1.1" ForwardDiff = "0.10.38, 1" @@ -61,7 +64,6 @@ Random = "1" RecipesBase = "1.3.4" ReverseDiff = "1.15" SafeTestsets = "0.1" -SciMLBase = "2.103" SparseArrays = "1.10" StaticArrays = "1.6" StaticArraysCore = "1.4.2" @@ -86,7 +88,6 @@ NLsolve = "2774e3e8-f4cf-5e23-947b-6d7e65073b56" Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" -SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" @@ -97,4 +98,4 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" [targets] -test = ["Aqua", "FastBroadcast", "ForwardDiff", "KernelAbstractions", "Measurements", "NLsolve", "Pkg", "Random", "SafeTestsets", "SciMLBase", "SparseArrays", "StaticArrays", "Statistics", "StructArrays", "Tables", "Test", "Unitful", "Zygote"] +test = ["Aqua", "FastBroadcast", "ForwardDiff", "KernelAbstractions", "Measurements", "NLsolve", "Pkg", "Random", "SafeTestsets", "SparseArrays", "StaticArrays", "Statistics", "StructArrays", "Tables", "Test", "Unitful", "Zygote"] diff --git a/docs/Project.toml b/docs/Project.toml index 56beb125..d7b8ccfd 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -4,4 +4,4 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" [compat] Documenter = "1.3" -RecursiveArrayTools = "3" +RecursiveArrayTools = "4" diff --git a/docs/pages.jl b/docs/pages.jl index b049a133..f4d44797 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -2,7 +2,9 @@ pages = [ "Home" => "index.md", + "breaking_changes_v4.md", "AbstractVectorOfArrayInterface.md", "array_types.md", + "plotting.md", "recursive_array_functions.md", ] diff --git a/docs/src/assets/Project.toml b/docs/src/assets/Project.toml new file mode 100644 index 00000000..56beb125 --- /dev/null +++ b/docs/src/assets/Project.toml @@ -0,0 +1,7 @@ +[deps] +Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" + +[compat] +Documenter = "1.3" +RecursiveArrayTools = "3" diff --git a/docs/src/breaking_changes_v4.md b/docs/src/breaking_changes_v4.md new file mode 100644 index 00000000..a9cbb83f --- /dev/null +++ b/docs/src/breaking_changes_v4.md @@ -0,0 +1,123 @@ +# Breaking Changes in v4.0: AbstractArray Interface + +## Summary + +`AbstractVectorOfArray{T, N, A}` now subtypes `AbstractArray{T, N}`. This means +all `VectorOfArray` and `DiffEqArray` objects are proper Julia `AbstractArray`s, +and all standard `AbstractArray` operations work out of the box, including linear +algebra, broadcasting with plain arrays, and generic algorithms. + +## Key Changes + +### Linear Indexing + +Previously, `A[i]` returned the `i`th inner array (`A.u[i]`). Now, `A[i]` returns +the `i`th element in column-major linear order, matching standard Julia `AbstractArray` +behavior. + +```julia +A = VectorOfArray([[1, 2], [3, 4]]) +# Old: A[1] == [1, 2] (first inner array) +# New: A[1] == 1 (first element, column-major) +# To access inner arrays: A.u[1] or A[:, 1] +``` + +### Size and Ragged Arrays + +For ragged arrays (inner arrays of different sizes), `size(A)` now reports the +**maximum** size in each dimension. Out-of-bounds elements are treated as zero +(sparse representation): + +```julia +A = VectorOfArray([[1, 2], [3, 4, 5]]) +size(A) # (3, 2) — max inner length is 3 +A[3, 1] # 0 — implicit zero (inner array 1 has only 2 elements) +A[3, 2] # 5 — actual stored value +Array(A) # [1 3; 2 4; 0 5] — zero-padded dense array +``` + +This means ragged `VectorOfArray`s can be used directly with linear algebra +operations, treating the data as a rectangular matrix with zero padding. + +### Iteration + +Iteration now goes over scalar elements in column-major order, matching +`AbstractArray` behavior: + +```julia +A = VectorOfArray([[1, 2], [3, 4]]) +collect(A) # [1 3; 2 4] — 2x2 matrix +# To iterate over inner arrays: for u in A.u ... end +``` + +### `length` + +`length(A)` now returns `prod(size(A))` (total number of elements including +ragged zeros), not the number of inner arrays. Use `length(A.u)` for the number +of inner arrays. + +### `map` + +`map(f, A)` now maps over individual elements, not inner arrays. Use +`map(f, A.u)` to map over inner arrays. + +### `first` / `last` + +`first(A)` and `last(A)` return the first/last scalar element, not the first/last +inner array. Use `first(A.u)` / `last(A.u)` for inner arrays. + +### `eachindex` + +`eachindex(A)` returns `CartesianIndices(size(A))` for the full rectangular shape, +not indices into `A.u`. + +## Migration Guide + +| Old Code | New Code | +|----------|----------| +| `A[i]` (get inner array) | `A.u[i]` or `A[:, i]` | +| `length(A)` (number of arrays) | `length(A.u)` | +| `for elem in A` (iterate columns) | `for elem in A.u` | +| `first(A)` (first inner array) | `first(A.u)` | +| `map(f, A)` (map over columns) | `map(f, A.u)` | +| `A == vec_of_vecs` | `A.u == vec_of_vecs` | + +## Interpolation Interface on DiffEqArray + +`DiffEqArray` now has `interp` and `dense` fields for interpolation support: + +```julia +# Create with interpolation +da = DiffEqArray(u, t, p, sys; interp = my_interp, dense = true) + +# Callable syntax +da(0.5) # interpolate at t=0.5 +da(0.5, Val{1}) # first derivative at t=0.5 +da([0.1, 0.5, 0.9]) # interpolate at multiple times +da(0.5; idxs = 1) # interpolate single component +da(0.5; idxs = [1, 2]) # interpolate subset of components +``` + +The interpolation object must be callable as `interp(t, idxs, deriv, p, continuity)`, +matching the protocol used by SciMLBase's `LinearInterpolation`, `HermiteInterpolation`, +and `ConstantInterpolation`. + +When `dense = true` and `interp` is provided, `plot(da)` automatically generates +dense interpolated output instead of plotting only the saved time points. + +## Ragged Arrays Sublibrary + +`RaggedVectorOfArray` and `RaggedDiffEqArray` are available via: + +```julia +using RecursiveArrayToolsRaggedArrays +``` + +These types preserve the true ragged structure without zero-padding, and do **not** +subtype `AbstractArray`. See the `RecursiveArrayToolsRaggedArrays` subpackage for details. + +## Zygote Compatibility + +Some Zygote adjoint rules need updating for the new `AbstractArray` subtyping. +ForwardDiff continues to work correctly. Zygote support will be updated in a +follow-up release. diff --git a/docs/src/plotting.md b/docs/src/plotting.md new file mode 100644 index 00000000..16e4f600 --- /dev/null +++ b/docs/src/plotting.md @@ -0,0 +1,154 @@ +# Plotting + +`AbstractVectorOfArray` and `AbstractDiffEqArray` types include +[Plots.jl](https://github.com/JuliaPlots/Plots.jl) recipes so they can be +visualised directly with `plot(A)`. + +## VectorOfArray + +A `VectorOfArray` plots as a matrix where each inner array is a column: + +```julia +using RecursiveArrayTools, Plots +A = VectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) +plot(A) # 3 series (components) vs column index +``` + +## DiffEqArray + +A `DiffEqArray` plots as component time series against `A.t`: + +```julia +u = [[sin(t), cos(t)] for t in 0:0.1:2pi] +t = collect(0:0.1:2pi) +A = DiffEqArray(u, t) +plot(A) # plots sin and cos vs t +``` + +If the `DiffEqArray` carries a symbolic system (via `variables` and +`independent_variables` keyword arguments), the axis labels and series names +are set automatically from the symbol names. + +## Dense (Interpolated) Plotting + +When a `DiffEqArray` has an interpolation object in its `interp` field and +`dense = true`, calling `plot(A)` generates a smooth curve by evaluating the +interpolation at many points rather than connecting only the saved time steps. + +```julia +using SciMLBase: LinearInterpolation + +u = [[1.0, 0.0], [0.0, 1.0], [-1.0, 0.0], [0.0, -1.0]] +t = [0.0, 1.0, 2.0, 3.0] +interp = LinearInterpolation(t, u) +A = DiffEqArray(u, t; interp = interp, dense = true) +plot(A) # smooth interpolated curve with 1000+ points +``` + +### Plot Recipe Keyword Arguments + +The `AbstractDiffEqArray` recipe accepts the following keyword arguments, +which can be passed directly to `plot`: + +| Keyword | Type | Default | Description | +|---------|------|---------|-------------| +| `denseplot` | `Bool` | `A.dense && A.interp !== nothing` | Use dense interpolation for smooth curves. Set `false` to show only saved points. | +| `plotdensity` | `Int` | `max(1000, 10 * length(A.u))` | Number of evenly-spaced points to evaluate when `denseplot = true`. | +| `tspan` | `Tuple` or `nothing` | `nothing` | Restrict the time window. E.g. `tspan = (0.0, 5.0)`. | +| `idxs` | varies | `nothing` | Select which components to plot (see below). | + +Example: + +```julia +plot(A; denseplot = true, plotdensity = 5000, tspan = (0.0, 2.0)) +``` + +## Selecting Variables with `idxs` + +The `idxs` keyword controls which variables appear in the plot. It supports +several formats: + +### Single component + +```julia +plot(A; idxs = 1) # plot component 1 vs time +plot(A; idxs = 2) # plot component 2 vs time +``` + +### Multiple components + +```julia +plot(A; idxs = [1, 3, 5]) # plot components 1, 3, 5 vs time +``` + +### Phase-space plots (component vs component) + +Use a tuple where index `0` represents the independent variable (time): + +```julia +plot(A; idxs = (1, 2)) # component 1 vs component 2 +plot(A; idxs = (0, 1)) # time vs component 1 (same as default) +plot(A; idxs = (1, 2, 3)) # 3D plot of components 1, 2, 3 +``` + +### Symbolic indexing + +When the `DiffEqArray` carries a symbolic system, variables can be referenced +by symbol: + +```julia +A = DiffEqArray(u, t; variables = [:x, :y], independent_variables = [:t]) +plot(A; idxs = :x) # plot x vs time +plot(A; idxs = [:x, :y]) # plot both +plot(A; idxs = (:x, :y)) # phase plot of x vs y +``` + +### Custom transformations + +A function can be applied to the plotted values: + +```julia +plot(A; idxs = (norm, 0, 1, 2)) # plot norm(u1, u2) vs time +``` + +The tuple format is `(f, xvar, yvar)` or `(f, xvar, yvar, zvar)` where `f` +is applied element-wise. + +## Callable Interface + +Any `AbstractDiffEqArray` with an `interp` field supports callable syntax for +interpolation, independent of plotting: + +```julia +A(0.5) # interpolate all components at t=0.5 +A(0.5; idxs = 1) # interpolate component 1 at t=0.5 +A([0.1, 0.5, 0.9]) # interpolate at multiple times (returns DiffEqArray) +A(0.5, Val{1}) # first derivative at t=0.5 +A(0.5; continuity = :right) # right-continuity at discontinuities +``` + +The interpolation object must be callable as +`interp(t, idxs, deriv, p, continuity)`, matching the protocol used by +SciMLBase's `LinearInterpolation`, `HermiteInterpolation`, and +`ConstantInterpolation`. + +When no interpolation is available (`interp === nothing`), calling `A(t)` +throws an error. + +## ODE Solution Plotting + +ODE solutions from DifferentialEquations.jl are subtypes of +`AbstractDiffEqArray` and inherit all of the above functionality, plus +additional features: + +- **Automatic dense plotting**: `denseplot` defaults to `true` when the solver + provides dense output. +- **Analytic solution overlay**: `plot(sol; plot_analytic = true)` overlays the + exact solution if the problem defines one. +- **Discrete variables**: Time-varying parameters are plotted as step functions + with dashed lines and markers. +- **Symbolic indexing**: Full symbolic indexing through ModelingToolkit is + supported, including observed (derived) variables. + +These advanced features are defined in SciMLBase and activate automatically +when plotting solution objects. diff --git a/ext/RecursiveArrayToolsCUDAExt.jl b/ext/RecursiveArrayToolsCUDAExt.jl new file mode 100644 index 00000000..950bf709 --- /dev/null +++ b/ext/RecursiveArrayToolsCUDAExt.jl @@ -0,0 +1,13 @@ +module RecursiveArrayToolsCUDAExt + +using RecursiveArrayTools: AbstractVectorOfArray +import CUDA: CuArray + +# Disambiguate CuArray(::AbstractVectorOfArray) vs CuArray(::AbstractArray{T,N}) from CUDA.jl. +# This is the exact signature Julia's ambiguity error requests. +# Uses stack to stay on GPU (avoids GPU→CPU→GPU round-trip). +function CuArray(VA::AbstractVectorOfArray{T, N}) where {T, N} + return CuArray{T, N}(stack(VA.u)) +end + +end diff --git a/ext/RecursiveArrayToolsZygoteExt.jl b/ext/RecursiveArrayToolsZygoteExt.jl index 94aa30d2..e4ef4c1f 100644 --- a/ext/RecursiveArrayToolsZygoteExt.jl +++ b/ext/RecursiveArrayToolsZygoteExt.jl @@ -5,67 +5,15 @@ using RecursiveArrayTools using Zygote using Zygote: FillArrays, ChainRulesCore, literal_getproperty, @adjoint -# Define a new species of projection operator for this type: -# ChainRulesCore.ProjectTo(x::VectorOfArray) = ChainRulesCore.ProjectTo{VectorOfArray}() - function ChainRulesCore.rrule( T::Type{<:RecursiveArrayTools.GPUArraysCore.AbstractGPUArray}, xs::AbstractVectorOfArray ) - return T(xs), ȳ -> (ChainRulesCore.NoTangent(), ȳ) -end - -@adjoint function getindex( - VA::AbstractVectorOfArray, - i::Union{BitArray, AbstractArray{Bool}} - ) - function AbstractVectorOfArray_getindex_adjoint(Δ) - Δ′ = [ - (i[j] ? Δ[j] : FillArrays.Fill(zero(eltype(x)), size(x))) - for (x, j) in zip(VA.u, 1:length(VA)) - ] - (VectorOfArray(Δ′), nothing) - end - VA[:, i], AbstractVectorOfArray_getindex_adjoint -end - -@adjoint function getindex(VA::AbstractVectorOfArray, i::AbstractArray{Int}) - function AbstractVectorOfArray_getindex_adjoint(Δ) - iter = 0 - Δ′ = [ - (j ∈ i ? Δ[iter += 1] : FillArrays.Fill(zero(eltype(x)), size(x))) - for (x, j) in zip(VA.u, 1:length(VA)) - ] - (VectorOfArray(Δ′), nothing) - end - VA[:, i], AbstractVectorOfArray_getindex_adjoint -end - -@adjoint function getindex(VA::AbstractVectorOfArray, i::Colon) - function AbstractVectorOfArray_getindex_adjoint(Δ) - (VectorOfArray(Δ), nothing) - end - VA.u[i], AbstractVectorOfArray_getindex_adjoint + return T(xs), ȳ -> (ChainRulesCore.NoTangent(), ȳ) end -@adjoint function getindex( - VA::AbstractVectorOfArray, i::Int, - j::Union{ - Int, AbstractArray{Int}, CartesianIndex, - Colon, BitArray, AbstractArray{Bool}, - }... - ) - function AbstractVectorOfArray_getindex_adjoint(Δ) - Δ′ = VectorOfArray([zero(x) for (x, j) in zip(VA.u, 1:length(VA))]) - if isempty(j) - Δ′.u[i] = Δ - else - Δ′[i, j...] = Δ - end - (Δ′, nothing, map(_ -> nothing, j)...) - end - VA[i, j...], AbstractVectorOfArray_getindex_adjoint -end +# getindex adjoints are inherited from Zygote's AbstractArray rules +# since AbstractVectorOfArray <: AbstractArray @adjoint function ArrayPartition( x::S, @@ -88,15 +36,22 @@ end @adjoint function VectorOfArray(u) VectorOfArray(u), y -> begin - y isa Ref && (y = VectorOfArray(y[].u)) - ( - VectorOfArray( + if y isa Ref + y = VectorOfArray(y[].u) + end + # Return a plain Vector of arrays as gradient for `u`, not wrapped in VectorOfArray. + # This avoids issues with downstream pullbacks that index into the gradient + # using linear indexing (which now returns scalar elements for VectorOfArray). + if y isa AbstractVectorOfArray + (y.u,) + else + ( [ y[ntuple(x -> Colon(), ndims(y) - 1)..., i] for i in 1:size(y)[end] - ] - ), - ) + ], + ) + end end end @@ -108,17 +63,20 @@ end @adjoint function DiffEqArray(u, t) DiffEqArray(u, t), y -> begin - y isa Ref && (y = VectorOfArray(y[].u)) - ( - DiffEqArray( + if y isa Ref + y = VectorOfArray(y[].u) + end + if y isa AbstractVectorOfArray + (y.u, nothing) + else + ( [ y[ntuple(x -> Colon(), ndims(y) - 1)..., i] for i in 1:size(y)[end] ], - t - ), - nothing, - ) + nothing, + ) + end end end @@ -156,6 +114,7 @@ end @adjoint function Base.Array(VA::AbstractVectorOfArray) adj = let VA = VA function Array_adjoint(y) + # Return a VectorOfArray so it flows correctly back through VectorOfArray constructor VA = recursivecopy(VA) copyto!(VA, y) return (VA,) @@ -164,243 +123,4 @@ end Array(VA), adj end -@adjoint function Base.view(A::AbstractVectorOfArray, I::Colon...) - view_adjoint = let A = A, I = I - function (y) - A = recursivecopy(A) - copyto!(A, y) - return (A, map(_ -> nothing, I)...) - end - end - return view(A, I...), view_adjoint -end - -@adjoint function Base.view(A::AbstractVectorOfArray, I...) - view_adjoint = let A = A, I = I - function (y) - A = recursivecopy(A) - recursivefill!(A, zero(eltype(A))) - v = view(A, I...) - copyto!(v, y) - return (A, map(_ -> nothing, I)...) - end - end - view(A, I...), view_adjoint -end - -@adjoint function Broadcast.broadcasted( - ::typeof(+), x::AbstractVectorOfArray, - y::Union{Zygote.Numeric, AbstractVectorOfArray} - ) - broadcast(+, x, y), ȳ -> (nothing, map(x -> Zygote.unbroadcast(x, ȳ), (x, y))...) -end -@adjoint function Broadcast.broadcasted( - ::typeof(+), x::Zygote.Numeric, y::AbstractVectorOfArray - ) - broadcast(+, x, y), ȳ -> (nothing, map(x -> Zygote.unbroadcast(x, ȳ), (x, y))...) -end - -_minus(Δ) = .-Δ -_minus(::Nothing) = nothing - -@adjoint function Broadcast.broadcasted( - ::typeof(-), x::AbstractVectorOfArray, - y::Union{AbstractVectorOfArray, Zygote.Numeric} - ) - x .- y, Δ -> (nothing, Zygote.unbroadcast(x, Δ), _minus(Zygote.unbroadcast(y, Δ))) -end -@adjoint function Broadcast.broadcasted( - ::typeof(*), x::AbstractVectorOfArray, - y::Union{AbstractVectorOfArray, Zygote.Numeric} - ) - ( - x .* y, - Δ -> ( - nothing, Zygote.unbroadcast(x, Δ .* conj.(y)), - Zygote.unbroadcast(y, Δ .* conj.(x)), - ), - ) -end -@adjoint function Broadcast.broadcasted( - ::typeof(/), x::AbstractVectorOfArray, - y::Union{AbstractVectorOfArray, Zygote.Numeric} - ) - res = x ./ y - res, - Δ -> ( - nothing, Zygote.unbroadcast(x, Δ ./ conj.(y)), - Zygote.unbroadcast(y, .-Δ .* conj.(res ./ y)), - ) -end -@adjoint function Broadcast.broadcasted( - ::typeof(-), x::Zygote.Numeric, y::AbstractVectorOfArray - ) - x .- y, Δ -> (nothing, Zygote.unbroadcast(x, Δ), _minus(Zygote.unbroadcast(y, Δ))) -end -@adjoint function Broadcast.broadcasted( - ::typeof(*), x::Zygote.Numeric, y::AbstractVectorOfArray - ) - ( - x .* y, - Δ -> ( - nothing, Zygote.unbroadcast(x, Δ .* conj.(y)), - Zygote.unbroadcast(y, Δ .* conj.(x)), - ), - ) -end -@adjoint function Broadcast.broadcasted( - ::typeof(/), x::Zygote.Numeric, y::AbstractVectorOfArray - ) - res = x ./ y - res, - Δ -> ( - nothing, Zygote.unbroadcast(x, Δ ./ conj.(y)), - Zygote.unbroadcast(y, .-Δ .* conj.(res ./ y)), - ) -end -@adjoint function Broadcast.broadcasted(::typeof(-), x::AbstractVectorOfArray) - .-x, Δ -> (nothing, _minus(Δ)) -end - -@adjoint function Broadcast.broadcasted( - ::typeof(Base.literal_pow), ::typeof(^), - x::AbstractVectorOfArray, exp::Val{p} - ) where {p} - y = Base.literal_pow.(^, x, exp) - y, ȳ -> (nothing, nothing, ȳ .* p .* conj.(x .^ (p - 1)), nothing) -end - -@adjoint Broadcast.broadcasted(::typeof(identity), x::AbstractVectorOfArray) = x, - Δ -> (nothing, Δ) - -@adjoint function Broadcast.broadcasted(::typeof(tanh), x::AbstractVectorOfArray) - y = tanh.(x) - y, ȳ -> (nothing, ȳ .* conj.(1 .- y .^ 2)) -end - -@adjoint Broadcast.broadcasted(::typeof(conj), x::AbstractVectorOfArray) = conj.(x), - z̄ -> (nothing, conj.(z̄)) - -@adjoint Broadcast.broadcasted(::typeof(real), x::AbstractVectorOfArray) = real.(x), - z̄ -> (nothing, real.(z̄)) - -@adjoint Broadcast.broadcasted( - ::typeof(imag), x::AbstractVectorOfArray -) = imag.(x), - z̄ -> (nothing, im .* real.(z̄)) - -@adjoint Broadcast.broadcasted( - ::typeof(abs2), - x::AbstractVectorOfArray -) = abs2.(x), - z̄ -> (nothing, 2 .* real.(z̄) .* x) - -@adjoint function Broadcast.broadcasted( - ::typeof(+), a::AbstractVectorOfArray{<:Number}, b::Bool - ) - y = b === false ? a : a .+ b - y, Δ -> (nothing, Δ, nothing) -end -@adjoint function Broadcast.broadcasted( - ::typeof(+), b::Bool, a::AbstractVectorOfArray{<:Number} - ) - y = b === false ? a : b .+ a - y, Δ -> (nothing, nothing, Δ) -end - -@adjoint function Broadcast.broadcasted( - ::typeof(-), a::AbstractVectorOfArray{<:Number}, b::Bool - ) - y = b === false ? a : a .- b - y, Δ -> (nothing, Δ, nothing) -end -@adjoint function Broadcast.broadcasted( - ::typeof(-), b::Bool, a::AbstractVectorOfArray{<:Number} - ) - b .- a, Δ -> (nothing, nothing, .-Δ) -end - -@adjoint function Broadcast.broadcasted( - ::typeof(*), a::AbstractVectorOfArray{<:Number}, b::Bool - ) - if b === false - zero(a), Δ -> (nothing, zero(Δ), nothing) - else - a, Δ -> (nothing, Δ, nothing) - end -end -@adjoint function Broadcast.broadcasted( - ::typeof(*), b::Bool, a::AbstractVectorOfArray{<:Number} - ) - if b === false - zero(a), Δ -> (nothing, nothing, zero(Δ)) - else - a, Δ -> (nothing, nothing, Δ) - end -end - -@adjoint Broadcast.broadcasted( - ::Type{T}, - x::AbstractVectorOfArray -) where { - T <: - Number, -} = T.(x), - ȳ -> (nothing, Zygote._project(x, ȳ)) - -function Zygote.unbroadcast(x::AbstractVectorOfArray, x̄) - N = ndims(x̄) - return if length(x) == length(x̄) - Zygote._project(x, x̄) # ProjectTo handles reshape, offsets, structured matrices, row vectors - else - dims = ntuple(d -> size(x, d) == 1 ? d : ndims(x̄) + 1, ndims(x̄)) - Zygote._project(x, Zygote.accum_sum(x̄; dims = dims)) - end -end - -@adjoint Broadcast.broadcasted( - ::Broadcast.AbstractArrayStyle, f::F, a::AbstractVectorOfArray, - b -) where {F} = _broadcast_generic( - __context__, f, a, b -) -@adjoint Broadcast.broadcasted( - ::Broadcast.AbstractArrayStyle, f::F, a, - b::AbstractVectorOfArray -) where {F} = _broadcast_generic( - __context__, f, a, b -) -@adjoint Broadcast.broadcasted( - ::Broadcast.AbstractArrayStyle, f::F, a::AbstractVectorOfArray, - b::AbstractVectorOfArray -) where {F} = _broadcast_generic( - __context__, f, a, b -) - -@inline function _broadcast_generic(__context__, f::F, args...) where {F} - T = Broadcast.combine_eltypes(f, args) - # Avoid generic broadcasting in two easy cases: - if T == Bool - return (f.(args...), _ -> nothing) - elseif T <: Union{Real, Complex} && isconcretetype(T) && Zygote._dual_purefun(F) && - all(Zygote._dual_safearg, args) && !Zygote.isderiving() - return Zygote.broadcast_forward(f, args...) - end - len = Zygote.inclen(args) - y∂b = Zygote._broadcast((x...) -> Zygote._pullback(__context__, f, x...), args...) - y = broadcast(first, y∂b) - function ∇broadcasted(ȳ) - y∂b = y∂b isa AbstractVectorOfArray ? Iterators.flatten(y∂b.u) : y∂b - ȳ = ȳ isa AbstractVectorOfArray ? Iterators.flatten(ȳ.u) : ȳ - dxs_zip = map(((_, pb), ȳ₁) -> pb(ȳ₁), y∂b, ȳ) - getters = ntuple(i -> Zygote.StaticGetter{i}(), len) - dxs = map(g -> Zygote.collapse_nothings(map(g, dxs_zip)), getters) - return ( - nothing, Zygote.accum_sum(dxs[1]), - map(Zygote.unbroadcast, args, Base.tail(dxs))..., - ) - end - return y, ∇broadcasted -end - end # module diff --git a/lib/RecursiveArrayToolsArrayPartitionAnyAll/Project.toml b/lib/RecursiveArrayToolsArrayPartitionAnyAll/Project.toml new file mode 100644 index 00000000..0be72b75 --- /dev/null +++ b/lib/RecursiveArrayToolsArrayPartitionAnyAll/Project.toml @@ -0,0 +1,16 @@ +name = "RecursiveArrayToolsArrayPartitionAnyAll" +uuid = "172d604e-c495-4f00-97bf-d70957099afa" +version = "1.0.0" + +[deps] +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" + +[compat] +RecursiveArrayTools = "4" +julia = "1.10" + +[extras] +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[targets] +test = ["Test"] diff --git a/lib/RecursiveArrayToolsArrayPartitionAnyAll/README.md b/lib/RecursiveArrayToolsArrayPartitionAnyAll/README.md new file mode 100644 index 00000000..28abf03b --- /dev/null +++ b/lib/RecursiveArrayToolsArrayPartitionAnyAll/README.md @@ -0,0 +1,37 @@ +# RecursiveArrayToolsArrayPartitionAnyAll.jl + +Optimized `any` and `all` for +[RecursiveArrayTools.jl](https://github.com/SciML/RecursiveArrayTools.jl)'s +`ArrayPartition` type. + +`ArrayPartition` stores data as a tuple of arrays. The default `AbstractArray` +`any`/`all` iterates element-by-element through the partition, which incurs +per-element partition lookup overhead. This subpackage provides methods that +iterate partition-by-partition instead, giving ~1.5-1.8x speedup on full scans. + +It is separated from the main package because `any(f::Function, ::ArrayPartition)` +invalidates ~780 compiled specializations of `any(f::Function, ::AbstractArray)`. + +## Installation + +```julia +using Pkg +Pkg.add(url = "https://github.com/SciML/RecursiveArrayTools.jl", + subdir = "lib/RecursiveArrayToolsArrayPartitionAnyAll") +``` + +## Usage + +```julia +using RecursiveArrayTools +using RecursiveArrayToolsArrayPartitionAnyAll + +ap = ArrayPartition(rand(1000), rand(1000), rand(1000)) + +# These now use the optimized partition-by-partition iteration +any(isnan, ap) +all(x -> x > 0, ap) +``` + +Without this package, `any`/`all` use the default `AbstractArray` implementation +which is correct but ~1.5x slower due to per-element partition indexing overhead. diff --git a/lib/RecursiveArrayToolsArrayPartitionAnyAll/src/RecursiveArrayToolsArrayPartitionAnyAll.jl b/lib/RecursiveArrayToolsArrayPartitionAnyAll/src/RecursiveArrayToolsArrayPartitionAnyAll.jl new file mode 100644 index 00000000..c7a47c22 --- /dev/null +++ b/lib/RecursiveArrayToolsArrayPartitionAnyAll/src/RecursiveArrayToolsArrayPartitionAnyAll.jl @@ -0,0 +1,25 @@ +""" + RecursiveArrayToolsArrayPartitionAnyAll + +Optimized `any`/`all` for `ArrayPartition` that iterates partition-by-partition +instead of using the generic `AbstractArray` element iteration. This gives +~1.5-1.8x speedup on full scans because it avoids per-element partition lookup +overhead. Separated into its own subpackage because `any(f::Function, ::ArrayPartition)` +invalidates ~780 specializations of `any(f::Function, ::AbstractArray)`. + +```julia +using RecursiveArrayToolsArrayPartitionAnyAll +``` +""" +module RecursiveArrayToolsArrayPartitionAnyAll + +using RecursiveArrayTools: ArrayPartition + +Base.any(f, A::ArrayPartition) = any((any(f, x) for x in A.x)) +Base.any(f::Function, A::ArrayPartition) = any((any(f, x) for x in A.x)) +Base.any(A::ArrayPartition) = any(identity, A) +Base.all(f, A::ArrayPartition) = all((all(f, x) for x in A.x)) +Base.all(f::Function, A::ArrayPartition) = all((all(f, x) for x in A.x)) +Base.all(A::ArrayPartition) = all(identity, A) + +end diff --git a/lib/RecursiveArrayToolsArrayPartitionAnyAll/test/runtests.jl b/lib/RecursiveArrayToolsArrayPartitionAnyAll/test/runtests.jl new file mode 100644 index 00000000..c765e279 --- /dev/null +++ b/lib/RecursiveArrayToolsArrayPartitionAnyAll/test/runtests.jl @@ -0,0 +1,34 @@ +using RecursiveArrayTools, RecursiveArrayToolsArrayPartitionAnyAll, Test + +@testset "Optimized any" begin + ap = ArrayPartition(collect(1:5), collect(6:10), collect(11:15)) + @test any(x -> x == 4, ap) + @test any(x -> x == 15, ap) + @test !any(x -> x == 17, ap) + @test any(ap .> 10) + @test !any(ap .> 20) +end + +@testset "Optimized all" begin + ap = ArrayPartition(ones(5), ones(5), ones(5)) + @test all(x -> x == 1.0, ap) + @test !all(x -> x == 2.0, ap) + @test all(ap .> 0) + + ap2 = ArrayPartition(ones(5), [1.0, 1.0, 0.0, 1.0, 1.0], ones(5)) + @test !all(x -> x == 1.0, ap2) +end + +@testset "Matches AbstractArray default results" begin + ap = ArrayPartition(rand(100), rand(100), rand(100)) + f = x -> x > 0.5 + + # Results must match + @test any(f, ap) == any(f, collect(ap)) + @test all(f, ap) == all(f, collect(ap)) + + # Edge case: empty + ap_empty = ArrayPartition(Float64[], Float64[]) + @test !any(x -> true, ap_empty) + @test all(x -> true, ap_empty) +end diff --git a/lib/RecursiveArrayToolsRaggedArrays/Project.toml b/lib/RecursiveArrayToolsRaggedArrays/Project.toml new file mode 100644 index 00000000..a6fd0aa6 --- /dev/null +++ b/lib/RecursiveArrayToolsRaggedArrays/Project.toml @@ -0,0 +1,19 @@ +name = "RecursiveArrayToolsRaggedArrays" +uuid = "c384ba91-639a-44ca-823a-e1d3691ab84a" +version = "1.0.0" + +[deps] +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" +SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" + +[compat] +RecursiveArrayTools = "4" +SymbolicIndexingInterface = "0.3.35" +julia = "1.10" + +[extras] +SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[targets] +test = ["SymbolicIndexingInterface", "Test"] diff --git a/lib/RecursiveArrayToolsRaggedArrays/src/RecursiveArrayToolsRaggedArrays.jl b/lib/RecursiveArrayToolsRaggedArrays/src/RecursiveArrayToolsRaggedArrays.jl new file mode 100644 index 00000000..1faff1d9 --- /dev/null +++ b/lib/RecursiveArrayToolsRaggedArrays/src/RecursiveArrayToolsRaggedArrays.jl @@ -0,0 +1,897 @@ +""" + RecursiveArrayToolsRaggedArrays + +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. + +Separated into its own subpackage because the method overloads (especially +`getindex` with `Colon`) would invalidate hot paths in Base if defined +unconditionally. + +```julia +using RecursiveArrayToolsRaggedArrays +``` + +# Quick start + +```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 +``` +""" +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 +# ═══════════════════════════════════════════════════════════════════════════════ + +""" + RaggedVectorOfArray{T, N, A} <: AbstractRaggedVectorOfArray{T, N, A} + +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. + +# Fields +- `u::A` — the underlying container of arrays + +# Constructors +```julia +RaggedVectorOfArray(vec::AbstractVector) +RaggedVectorOfArray(va::AbstractVectorOfArray) +``` +""" +mutable struct RaggedVectorOfArray{T, N, A} <: AbstractRaggedVectorOfArray{T, N, A} + u::A +end + +""" + 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`) +""" +mutable struct RaggedDiffEqArray{ + T, N, A, B, F, S, D <: Union{Nothing, ParameterTimeseriesCollection}, I, + } <: AbstractRaggedDiffEqArray{T, N, A} + u::A + t::B + p::F + sys::S + discretes::D + interp::I + dense::Bool +end + +# ═══════════════════════════════════════════════════════════════════════════════ +# Constructors — RaggedVectorOfArray +# ═══════════════════════════════════════════════════════════════════════════════ + +function RaggedVectorOfArray(vec::AbstractVector) + T = eltype(vec[1]) + N = ndims(vec[1]) + if all( + x -> x isa Union{<:AbstractArray, <:AbstractVectorOfArray, <:AbstractRaggedVectorOfArray}, + 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}} + 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 + +# Convert from VectorOfArray +function RaggedVectorOfArray(va::AbstractVectorOfArray{T, N, A}) where {T, N, A} + return RaggedVectorOfArray{T, N, A}(va.u) +end + +# ═══════════════════════════════════════════════════════════════════════════════ +# Constructors — RaggedDiffEqArray +# ═══════════════════════════════════════════════════════════════════════════════ + +function RaggedDiffEqArray( + vec::AbstractVector, ts::AbstractVector; + discretes = nothing, variables = nothing, parameters = nothing, + independent_variables = nothing, interp = nothing, dense = false + ) + sys = SymbolCache( + something(variables, []), + something(parameters, []), + something(independent_variables, []) + ) + _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) +end + +function RaggedDiffEqArray( + vec::AbstractVector{VT}, ts::AbstractVector; + discretes = nothing, variables = nothing, parameters = nothing, + independent_variables = nothing, interp = nothing, dense = false + ) where {T, N, VT <: AbstractArray{T, N}} + sys = SymbolCache( + something(variables, []), + something(parameters, []), + 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) +end + +function RaggedDiffEqArray( + vec::AbstractVector, ts::AbstractVector, p; + discretes = nothing, variables = nothing, parameters = nothing, + independent_variables = nothing, interp = nothing, dense = false + ) + sys = SymbolCache( + something(variables, []), + something(parameters, []), + something(independent_variables, []) + ) + _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) +end + +function RaggedDiffEqArray( + vec::AbstractVector{VT}, ts::AbstractVector, p; + discretes = nothing, variables = nothing, parameters = nothing, + independent_variables = nothing, interp = nothing, dense = false + ) where {T, N, VT <: AbstractArray{T, N}} + sys = SymbolCache( + something(variables, []), + something(parameters, []), + something(independent_variables, []) + ) + 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 + +function RaggedDiffEqArray( + vec::AbstractVector, ts::AbstractVector, p, sys; + discretes = nothing, interp = nothing, dense = false + ) + _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) +end + +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) +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 + +function DiffEqArray(r::AbstractRaggedDiffEqArray) + return DiffEqArray( + r.u, r.t, r.p, r.sys; discretes = r.discretes, + interp = r.interp, dense = r.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 + +# 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)) +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] +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)) +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.@propagate_inbounds function Base.getindex( + r::RaggedVectorOfArray{T, N}, i::Int + ) where {T, N} + return _ragged_linear_getindex(r, i) +end + +Base.@propagate_inbounds function Base.getindex( + r::RaggedDiffEqArray{T, N}, i::Int + ) where {T, N} + return _ragged_linear_getindex(r, i) +end + +Base.@propagate_inbounds function Base.setindex!( + r::RaggedVectorOfArray{T, 1}, v, i::Int + ) where {T} + r.u[i] = v + return v +end + +Base.@propagate_inbounds function Base.setindex!( + r::RaggedDiffEqArray{T, 1}, v, i::Int + ) where {T} + r.u[i] = v + return v +end + +Base.@propagate_inbounds function Base.setindex!( + r::RaggedVectorOfArray{T, N}, v, i::Int + ) where {T, N} + return _ragged_linear_setindex!(r, v, i) +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 + +# 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 + +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 + +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 + +# 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) + +# A[:, :] returns a copy of the ragged array +function Base.getindex(r::RaggedVectorOfArray, ::Colon, ::Colon) + return RaggedVectorOfArray(copy(r.u)) +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 + ) +end + +# A[:, idx_array] returns a subset +function Base.getindex( + r::RaggedVectorOfArray, ::Colon, + I::Union{AbstractArray{Int}, AbstractArray{Bool}} + ) + return RaggedVectorOfArray(r.u[I]) +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 + ) +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] +end +function Base.getindex(r::RaggedDiffEqArray, i::Int, ::Colon) + return [u[i] for u in r.u] +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] +end + +# 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] +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 + ) + 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], :) +end + +function Base.view(r::AbstractRaggedVectorOfArray, I, i::Int) + return view(r.u[i], 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... + ) + 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 + else + _ragged_getindex(r, symtype, elsymtype, _arg, args...) + end +end + +# Non-symbolic multi-arg: Colon + Int already handled above, handle remaining cases +Base.@propagate_inbounds function _ragged_getindex( + r::AbstractRaggedDiffEqArray, ::NotSymbolic, + I::Union{Int, AbstractArray{Int}, AbstractArray{Bool}, Colon}... + ) + if last(I) isa Int + return r.u[last(I)][Base.front(I)...] + else + col_idxs = last(I) isa Colon ? eachindex(r.u) : last(I) + if all(idx -> idx isa Colon, Base.front(I)) + u_slice = [r.u[col][Base.front(I)...] for col in col_idxs] + return RaggedDiffEqArray( + u_slice, r.t[col_idxs], r.p, r.sys; + discretes = r.discretes, interp = r.interp, dense = r.dense + ) + else + return [r.u[col][Base.front(I)...] for col in col_idxs] + end + end +end + +# ParameterIndexingError — reuse from RecursiveArrayTools +struct RaggedParameterIndexingError + sym::Any +end +function Base.showerror(io::IO, pie::RaggedParameterIndexingError) + return print( + io, + "Indexing with parameters is deprecated. Use `getp(A, $(pie.sym))` for parameter indexing." + ) +end + +# Symbolic indexing methods +for (symtype, elsymtype, valtype, errcheck) in [ + ( + ScalarSymbolic, SymbolicIndexingInterface.SymbolicTypeTrait, Any, + :(is_parameter(A, sym) && !is_timeseries_parameter(A, sym)), + ), + ( + ArraySymbolic, SymbolicIndexingInterface.SymbolicTypeTrait, Any, + :(is_parameter(A, sym) && !is_timeseries_parameter(A, sym)), + ), + ( + NotSymbolic, SymbolicIndexingInterface.SymbolicTypeTrait, + Union{<:Tuple, <:AbstractArray}, + :(all(x -> is_parameter(A, x) && !is_timeseries_parameter(A, x), sym)), + ), + ] + @eval Base.@propagate_inbounds function _ragged_getindex( + A::AbstractRaggedDiffEqArray, ::$symtype, + ::$elsymtype, sym::$valtype, arg... + ) + if $errcheck + throw(RaggedParameterIndexingError(sym)) + end + return getu(A, sym)(A, arg...) + end +end + +Base.@propagate_inbounds function _ragged_getindex( + A::AbstractRaggedDiffEqArray, ::ScalarSymbolic, + ::NotSymbolic, ::SymbolicIndexingInterface.SolvedVariables, args... + ) + return getindex(A, variable_symbols(A), args...) +end + +Base.@propagate_inbounds function _ragged_getindex( + A::AbstractRaggedDiffEqArray, ::ScalarSymbolic, + ::NotSymbolic, ::SymbolicIndexingInterface.AllVariables, args... + ) + return getindex(A, all_variable_symbols(A), args...) +end + +# ═══════════════════════════════════════════════════════════════════════════════ +# Observed functions +# ═══════════════════════════════════════════════════════════════════════════════ + +function _ragged_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} + 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} + return observed(A, sym).(A.u, (A.p,), A.t) +end + +# ═══════════════════════════════════════════════════════════════════════════════ +# Sizes of individual inner arrays +# ═══════════════════════════════════════════════════════════════════════════════ + +""" + inner_sizes(r::AbstractRaggedVectorOfArray) + +Return a vector of the sizes of each inner array. +""" +inner_sizes(r::AbstractRaggedVectorOfArray) = size.(r.u) + +""" + inner_lengths(r::AbstractRaggedVectorOfArray) + +Return a vector of the lengths of each inner array. +""" +inner_lengths(r::AbstractRaggedVectorOfArray) = length.(r.u) + +export inner_sizes, inner_lengths + +# ═══════════════════════════════════════════════════════════════════════════════ +# Copy, zero, similar, fill! +# ═══════════════════════════════════════════════════════════════════════════════ + +function _ragged_copyfield(r, fname) + return if fname == :u + copy(r.u) + elseif fname == :t + copy(r.t) + else + getfield(r, fname) + end +end + +function Base.copy(r::AbstractRaggedVectorOfArray) + return typeof(r)((_ragged_copyfield(r, fname) for fname in fieldnames(typeof(r)))...) +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...) +end + +function Base.similar(r::RaggedVectorOfArray) + return RaggedVectorOfArray(similar.(r.u)) +end + +function Base.similar(r::RaggedVectorOfArray, ::Type{T}) where {T} + return RaggedVectorOfArray(similar.(r.u, T)) +end + +function Base.fill!(r::AbstractRaggedVectorOfArray, x) + for u in r.u + fill!(u, x) + end + return r +end + +# ═══════════════════════════════════════════════════════════════════════════════ +# Mutation: push!, append!, resize! +# ═══════════════════════════════════════════════════════════════════════════════ + +function Base.push!(r::AbstractRaggedVectorOfArray, new_item::AbstractArray) + return push!(r.u, new_item) +end + +function Base.append!( + r::AbstractRaggedVectorOfArray, other::AbstractRaggedVectorOfArray + ) + for item in copy(other.u) + push!(r, item) + end + return r +end + +Base.sizehint!(r::AbstractRaggedVectorOfArray, i) = sizehint!(r.u, i) + +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 + ) +end + +function Base.copyto!( + dest::AbstractRaggedVectorOfArray, src::AbstractRaggedVectorOfArray + ) + for (i, j) in zip(eachindex(dest.u), eachindex(src.u)) + copyto!(dest.u[i], src.u[j]) + end + return dest +end + +# ═══════════════════════════════════════════════════════════════════════════════ +# Equality +# ═══════════════════════════════════════════════════════════════════════════════ + +function Base.:(==)(a::AbstractRaggedVectorOfArray, b::AbstractRaggedVectorOfArray) + return a.u == b.u +end + +# ═══════════════════════════════════════════════════════════════════════════════ +# Broadcasting — column-wise, preserving ragged structure +# ═══════════════════════════════════════════════════════════════════════════════ + +struct RaggedVectorOfArrayStyle{N} <: Broadcast.BroadcastStyle end +RaggedVectorOfArrayStyle{N}(::Val{N}) where {N} = RaggedVectorOfArrayStyle{N}() +RaggedVectorOfArrayStyle(::Val{N}) where {N} = RaggedVectorOfArrayStyle{N}() + +function Broadcast.BroadcastStyle(::Type{<:AbstractRaggedVectorOfArray{T, N}}) where {T, N} + return RaggedVectorOfArrayStyle{N}() +end + +# Make ragged arrays broadcastable without being collected +Broadcast.broadcastable(r::AbstractRaggedVectorOfArray) = r + +# 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))) +end +# RaggedVectorOfArrayStyle wins over itself at different dims +function Broadcast.BroadcastStyle( + ::RaggedVectorOfArrayStyle{M}, ::RaggedVectorOfArrayStyle{N} + ) where {M, N} + return RaggedVectorOfArrayStyle(Val(max(M, N))) +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")) + ) + ) +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)) +end +function _ragged_unpack(bc::Broadcast.Broadcasted{<:RaggedVectorOfArrayStyle}, i) + return Broadcast.Broadcasted(bc.f, _ragged_unpack_args(i, bc.args)) +end +function _ragged_unpack_args(i, args::Tuple) + return (_ragged_unpack(args[1], i), _ragged_unpack_args(i, Base.tail(args))...) +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)) + end + return RaggedVectorOfArray(u) +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) +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) +end + +@inline function Broadcast.materialize!( + dest::AbstractRaggedVectorOfArray, + bc::Broadcast.Broadcasted + ) + return copyto!(dest, bc) +end + +@inline function Broadcast.materialize!( + dest::AbstractRaggedVectorOfArray, + x::Any + ) + return Broadcast.materialize!(dest, Broadcast.broadcasted(identity, x)) +end + +# copyto!: in-place broadcast +@inline function Base.copyto!( + dest::AbstractRaggedVectorOfArray, + 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)) + end + return dest +end + +# ═══════════════════════════════════════════════════════════════════════════════ +# Show +# ═══════════════════════════════════════════════════════════════════════════════ + +function Base.show(io::IO, m::MIME"text/plain", r::AbstractRaggedVectorOfArray) + println(io, summary(r), ':') + return show(io, m, r.u) +end + +function Base.summary(r::AbstractRaggedVectorOfArray{T, N}) where {T, N} + return string("RaggedVectorOfArray{", T, ",", N, "}") +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) +end + +function Base.summary(r::AbstractRaggedDiffEqArray{T, N}) where {T, N} + return string("RaggedDiffEqArray{", T, ",", N, "}") +end + +# ═══════════════════════════════════════════════════════════════════════════════ +# Callable interface — dense interpolation +# ═══════════════════════════════════════════════════════════════════════════════ + +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 + +# ═══════════════════════════════════════════════════════════════════════════════ +# SymbolicIndexingInterface +# ═══════════════════════════════════════════════════════════════════════════════ + +SymbolicIndexingInterface.is_timeseries(::Type{<:AbstractRaggedVectorOfArray}) = Timeseries() + +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() +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 + +RecursiveArrayTools.has_discretes(::T) where {T <: AbstractRaggedDiffEqArray} = hasfield(T, :discretes) +RecursiveArrayTools.get_discretes(r::AbstractRaggedDiffEqArray) = r.discretes + +function SymbolicIndexingInterface.get_parameter_timeseries_collection(r::AbstractRaggedDiffEqArray) + return r.discretes +end + +end # module diff --git a/lib/RecursiveArrayToolsRaggedArrays/test/runtests.jl b/lib/RecursiveArrayToolsRaggedArrays/test/runtests.jl new file mode 100644 index 00000000..c3626499 --- /dev/null +++ b/lib/RecursiveArrayToolsRaggedArrays/test/runtests.jl @@ -0,0 +1,557 @@ +using RecursiveArrayTools +using RecursiveArrayToolsRaggedArrays +using SymbolicIndexingInterface +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 + + # From typed vector + r2 = RaggedVectorOfArray(Vector{Float64}[[1.0], [2.0, 3.0]]) + @test length(r2) == 2 + end + + @testset "RaggedVectorOfArray linear indexing (column-major)" begin + r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) + + # 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] + + # 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 + end + + @testset "RaggedVectorOfArray N-ary indexing — no zero-padding" begin + r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0]]) + + # 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 + + # 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 + + # A[j, :] returns time series of j-th component + @test r[1, :] == [1.0, 3.0, 6.0] + + # 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] + + # A[:, bool_array] + r_bool = r[:, [true, false, true]] + @test r_bool isa RaggedVectorOfArray + @test length(r_bool) == 2 + + # A[:, :] + r_all = r[:, :] + @test r_all isa RaggedVectorOfArray + @test r_all == r + @test r_all.u !== r.u # copy + 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 + + r[:, 2] = [30.0, 40.0, 50.0] + @test r[:, 2] == [30.0, 40.0, 50.0] + 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] + end + + @testset "RaggedVectorOfArray copy, zero, similar, fill!" begin + r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) + + r2 = copy(r) + @test r2 == r + @test r2.u !== r.u + + r0 = zero(r) + @test r0[:, 1] == [0.0, 0.0] + @test r0[:, 2] == [0.0, 0.0, 0.0] + + rs = similar(r) + @test length(rs[:, 1]) == 2 + @test length(rs[:, 2]) == 3 + + fill!(r0, 7.0) + @test r0[:, 1] == [7.0, 7.0] + @test r0[:, 2] == [7.0, 7.0, 7.0] + 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] + + r2 = RaggedVectorOfArray([[7.0]]) + append!(r, r2) + @test length(r) == 4 + @test r[:, 4] == [7.0] + end + + @testset "RaggedVectorOfArray broadcasting" begin + r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) + + # 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] + + # 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] + + # 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] + + # .= 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] + 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] + end + + @testset "RaggedVectorOfArray ↔ VectorOfArray conversion" begin + r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0]]) + + # To VectorOfArray (rectangular, zero-padded) + va = VectorOfArray(r) + @test va isa VectorOfArray + @test size(va) == (3, 2) + @test va[:, 1] == [1.0, 2.0, 0.0] # zero-padded + @test va[:, 2] == [3.0, 4.0, 5.0] + + # Back to RaggedVectorOfArray + r2 = RaggedVectorOfArray(va) + @test r2 isa RaggedVectorOfArray + # u is shared, inner arrays have original sizes + @test r2[:, 1] == [1.0, 2.0] + @test r2[:, 2] == [3.0, 4.0, 5.0] + end + + @testset "RaggedVectorOfArray equality" begin + r1 = RaggedVectorOfArray([[1, 2], [3, 4, 5]]) + r2 = RaggedVectorOfArray([[1, 2], [3, 4, 5]]) + r3 = RaggedVectorOfArray([[1, 2], [3, 4, 6]]) + @test r1 == r2 + @test r1 != r3 + end + + @testset "RaggedVectorOfArray reverse" begin + r = RaggedVectorOfArray([[1.0], [2.0, 3.0], [4.0, 5.0, 6.0]]) + rr = reverse(r) + @test rr[:, 1] == [4.0, 5.0, 6.0] + @test rr[:, 2] == [2.0, 3.0] + @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]] + @test r.t == [0.0, 1.0] + @test length(r) == 2 + + # With parameters + r2 = RaggedDiffEqArray([[1.0], [2.0, 3.0]], [0.0, 1.0], [0.5]) + @test r2.p == [0.5] + end + + @testset "RaggedDiffEqArray indexing — no zero-padding" begin + r = RaggedDiffEqArray([[1.0, 2.0], [3.0, 4.0, 5.0]], [0.0, 1.0]) + + @test r[:, 1] == [1.0, 2.0] + @test r[:, 2] == [3.0, 4.0, 5.0] + @test r[1, 1] == 1.0 + @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 + + # Colon-colon preserves DiffEqArray type + r_all = r[:, :] + @test r_all isa RaggedDiffEqArray + @test r_all.t == [0.0, 1.0] + + # Subset preserves time + r_sub = r[:, [2]] + @test r_sub isa RaggedDiffEqArray + @test r_sub.t == [1.0] + @test r_sub[:, 1] == [3.0, 4.0, 5.0] + end + + @testset "RaggedDiffEqArray ↔ DiffEqArray conversion" begin + r = RaggedDiffEqArray([[1.0, 2.0], [3.0, 4.0, 5.0]], [0.0, 1.0]) + + da = DiffEqArray(r) + @test da isa DiffEqArray + @test da.t == [0.0, 1.0] + + r2 = RaggedDiffEqArray(da) + @test r2 isa RaggedDiffEqArray + @test r2[:, 1] == [1.0, 2.0] + @test r2[:, 2] == [3.0, 4.0, 5.0] + @test r2.t == [0.0, 1.0] + end + + @testset "RaggedDiffEqArray copy/zero" begin + r = RaggedDiffEqArray([[1.0, 2.0], [3.0, 4.0, 5.0]], [0.0, 1.0]) + + r2 = copy(r) + @test r2 == r + @test r2.u !== r.u + + r0 = zero(r) + @test r0[:, 1] == [0.0, 0.0] + @test r0[:, 2] == [0.0, 0.0, 0.0] + @test r0.t == [0.0, 1.0] + end + + @testset "RaggedDiffEqArray SymbolicIndexingInterface" begin + r = RaggedDiffEqArray( + [[1.0, 2.0], [3.0, 4.0, 5.0]], [0.0, 1.0]; + variables = [:a, :b], parameters = [:p], independent_variables = [:t] + ) + + @test is_timeseries(r) == Timeseries() + @test state_values(r) === r.u + @test current_time(r) === r.t + @test symbolic_container(r) === r.sys + + # Variable/parameter query + @test is_variable(r, :a) + @test is_variable(r, :b) + @test is_parameter(r, :p) + @test is_independent_variable(r, :t) + @test variable_index(r, :a) == 1 + @test variable_index(r, :b) == 2 + @test parameter_index(r, :p) == 1 + end + + @testset "RaggedDiffEqArray symbolic getindex" begin + r = RaggedDiffEqArray( + [[1.0, 2.0], [3.0, 4.0, 5.0]], [0.0, 1.0], [0.5]; + variables = [:a, :b], parameters = [:p], independent_variables = [:t] + ) + + # Symbolic variable indexing + @test r[:a] == [1.0, 3.0] + @test r[:a, 1] == 1.0 + @test r[:a, 2] == 3.0 + @test r[:b] == [2.0, 4.0] + + # Parameter indexing should throw + @test_throws RecursiveArrayToolsRaggedArrays.RaggedParameterIndexingError r[:p] + end + + @testset "2D inner arrays" begin + r = RaggedVectorOfArray([zeros(2, 3), zeros(2, 4)]) + @test length(r) == 2 + @test ndims(r) == 3 + @test size(r[:, 1]) == (2, 3) + @test size(r[:, 2]) == (2, 4) + + # No zero-padding in column access + @test r[:, 1] == zeros(2, 3) + @test r[:, 2] == zeros(2, 4) + end + + @testset "RaggedDiffEqArray interp/dense fields" begin + r = RaggedDiffEqArray([[1.0, 2.0], [3.0, 4.0, 5.0]], [0.0, 1.0]) + @test r.interp === nothing + @test r.dense == false + + # With interp kwarg + r2 = RaggedDiffEqArray( + [[1.0, 2.0], [3.0, 4.0, 5.0]], [0.0, 1.0]; + interp = :test_interp, dense = true + ) + @test r2.interp == :test_interp + @test r2.dense == true + + # Callable errors when no interp + @test_throws ErrorException r(0.5) + + # Conversion preserves interp/dense + da = DiffEqArray(r2) + @test da.interp == :test_interp + @test da.dense == true + + r3 = RaggedDiffEqArray(da) + @test r3.interp == :test_interp + @test r3.dense == true + + # Copy preserves interp/dense + r4 = copy(r2) + @test r4.interp == :test_interp + @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]]) + + # zero preserves ragged structure + r0 = zero(r) + @test r0[:, 1] == [0.0, 0.0] + @test r0[:, 2] == [0.0, 0.0, 0.0] + @test length(r0[:, 1]) == 2 + @test length(r0[:, 2]) == 3 + + # fill! preserves ragged structure + fill!(r0, 42.0) + @test r0[:, 1] == [42.0, 42.0] + @test r0[:, 2] == [42.0, 42.0, 42.0] + + # .= zero works (the exact bug JoshuaLampert reported) + r2 = copy(r) + r2 .= zero(r2) + @test r2[:, 1] == [0.0, 0.0] + @test r2[:, 2] == [0.0, 0.0, 0.0] + end + + @testset "v3 component timeseries (ported)" begin + r = RaggedVectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0]]) + + # A[j, :] returns the j-th component across all inner arrays + @test r[1, :] == [1.0, 3.0, 6.0] + @test r[2, :] == [2.0, 4.0, 7.0] + end + + @testset "Type hierarchy" begin + r = RaggedVectorOfArray([[1, 2], [3, 4, 5]]) + @test r isa RecursiveArrayTools.AbstractRaggedVectorOfArray + @test !(r isa AbstractArray) + + rd = RaggedDiffEqArray([[1.0], [2.0, 3.0]], [0.0, 1.0]) + @test rd isa RecursiveArrayTools.AbstractRaggedDiffEqArray + @test rd isa RecursiveArrayTools.AbstractRaggedVectorOfArray + @test !(rd isa AbstractArray) + end +end diff --git a/lib/RecursiveArrayToolsShorthandConstructors/Project.toml b/lib/RecursiveArrayToolsShorthandConstructors/Project.toml new file mode 100644 index 00000000..a355d716 --- /dev/null +++ b/lib/RecursiveArrayToolsShorthandConstructors/Project.toml @@ -0,0 +1,16 @@ +name = "RecursiveArrayToolsShorthandConstructors" +uuid = "39fb7555-b4ad-4efd-8abe-30331df017d3" +version = "1.0.0" + +[deps] +RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" + +[compat] +RecursiveArrayTools = "4" +julia = "1.10" + +[extras] +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" + +[targets] +test = ["Test"] diff --git a/lib/RecursiveArrayToolsShorthandConstructors/README.md b/lib/RecursiveArrayToolsShorthandConstructors/README.md new file mode 100644 index 00000000..7b957a5d --- /dev/null +++ b/lib/RecursiveArrayToolsShorthandConstructors/README.md @@ -0,0 +1,45 @@ +# RecursiveArrayToolsShorthandConstructors.jl + +Shorthand constructor syntax for +[RecursiveArrayTools.jl](https://github.com/SciML/RecursiveArrayTools.jl) types. + +This subpackage provides `VA[...]` and `AP[...]` syntax for constructing +`VectorOfArray` and `ArrayPartition` objects. It is separated from the main +package because the `getindex(::Type, ...)` method definitions invalidate +compiled specializations of `Base.getindex(::Type{T}, vals...)`, increasing +load times for downstream packages that don't need the syntax. + +## Installation + +```julia +using Pkg +Pkg.add(url = "https://github.com/SciML/RecursiveArrayTools.jl", + subdir = "lib/RecursiveArrayToolsShorthandConstructors") +``` + +## Usage + +```julia +using RecursiveArrayTools +using RecursiveArrayToolsShorthandConstructors + +# VectorOfArray shorthand +u = VA[[1, 2, 3], [4, 5, 6], [7, 8, 9]] + +# ArrayPartition shorthand +x0, v0, a0 = rand(3, 3), rand(3, 3), rand(3, 3) +u0 = AP[x0, v0, a0] + +# Nesting works too +nested = VA[ + fill(1, 2, 3), + VA[3ones(3), zeros(3)], +] +``` + +Without this package, use the equivalent explicit constructors: + +```julia +u = VectorOfArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]]) +u0 = ArrayPartition(x0, v0, a0) +``` diff --git a/lib/RecursiveArrayToolsShorthandConstructors/src/RecursiveArrayToolsShorthandConstructors.jl b/lib/RecursiveArrayToolsShorthandConstructors/src/RecursiveArrayToolsShorthandConstructors.jl new file mode 100644 index 00000000..d11f6916 --- /dev/null +++ b/lib/RecursiveArrayToolsShorthandConstructors/src/RecursiveArrayToolsShorthandConstructors.jl @@ -0,0 +1,19 @@ +""" + RecursiveArrayToolsShorthandConstructors + +Shorthand `VA[a, b, c]` and `AP[a, b, c]` constructor syntax. Separated +into its own subpackage because `getindex(::Type{VA}, xs...)` invalidates +`getindex(::Type{T}, vals...)` from Base. + +```julia +using RecursiveArrayToolsShorthandConstructors +``` +""" +module RecursiveArrayToolsShorthandConstructors + +using RecursiveArrayTools: VA, VectorOfArray, AP, ArrayPartition + +Base.getindex(::Type{VA}, xs...) = VectorOfArray(collect(xs)) +Base.getindex(::Type{AP}, xs...) = ArrayPartition(xs...) + +end diff --git a/lib/RecursiveArrayToolsShorthandConstructors/test/runtests.jl b/lib/RecursiveArrayToolsShorthandConstructors/test/runtests.jl new file mode 100644 index 00000000..006061ea --- /dev/null +++ b/lib/RecursiveArrayToolsShorthandConstructors/test/runtests.jl @@ -0,0 +1,29 @@ +using RecursiveArrayTools, RecursiveArrayToolsShorthandConstructors, Test + +@testset "VA[...] shorthand" begin + recs = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] + testva = VA[[1, 2, 3], [4, 5, 6], [7, 8, 9]] + @test testva isa VectorOfArray + @test testva.u == recs + @test Array(testva) == [1 4 7; 2 5 8; 3 6 9] + + # Nesting + nested = VA[ + fill(1, 2, 3), + VA[3ones(3), zeros(3)], + ] + @test nested isa VectorOfArray + @test nested.u[1] == fill(1, 2, 3) + @test nested.u[2] isa VectorOfArray +end + +@testset "AP[...] shorthand" begin + x = AP[1:5, 1:6] + @test x isa ArrayPartition + @test length(x) == 11 + @test x[1:5] == 1:5 + @test x[6:11] == 1:6 + + @test length(AP[]) == 0 + @test isempty(AP[]) +end diff --git a/src/RecursiveArrayTools.jl b/src/RecursiveArrayTools.jl index f6f813fc..681e7474 100644 --- a/src/RecursiveArrayTools.jl +++ b/src/RecursiveArrayTools.jl @@ -26,9 +26,13 @@ module RecursiveArrayTools !!! note - In 2023 the linear indexing `A[i]` was deprecated. It previously had the behavior that `A[i] = A.u[i]`. However, this is incompatible with standard `AbstractArray` interfaces, Since if `A = VectorOfArray([[1,2],[3,4]])` and `A` is supposed to act like `[1 3; 2 4]`, then there is a difference `A[1] = [1,2]` for the VectorOfArray while `A[1] = 1` for the matrix. This causes many issues if `AbstractVectorOfArray <: AbstractArray`. Thus we plan in 2026 to complete the deprecation and thus have a breaking update where `A[i]` matches the linear indexing of an`AbstractArray`, and then making `AbstractVectorOfArray <: AbstractArray`. Until then, `AbstractVectorOfArray` due to - this interface break but manually implements an `AbstractArray`-like interface for - future compatibility. + As of v4.0, `AbstractVectorOfArray <: AbstractArray`. Linear indexing `A[i]` + now returns the `i`th element in column-major order, matching standard Julia + `AbstractArray` behavior. To access the `i`th inner array, use `A.u[i]` or + `A[:, i]`. For ragged arrays (inner arrays of different sizes), `size(A)` + reports the maximum size in each dimension and out-of-bounds elements are + interpreted as zero (sparse representation). This means all standard linear + algebra operations work out of the box. ## Fields @@ -99,7 +103,7 @@ module RecursiveArrayTools to make the array type match the internal array type (for example, if `A` is an array of GPU arrays, `stack(A)` will be a GPU array). """ - abstract type AbstractVectorOfArray{T, N, A} end + abstract type AbstractVectorOfArray{T, N, A} <: AbstractArray{T, N} end """ AbstractDiffEqArray{T, N, A} <: AbstractVectorOfArray{T, N, A} @@ -117,25 +121,74 @@ module RecursiveArrayTools An AbstractDiffEqArray adds the following fields: - `t` which holds the times of each timestep. + - `p` which holds the parameter values. + - `sys` which holds the symbolic system (e.g. `SymbolCache`). + - `discretes` which holds discrete parameter timeseries. + - `interp` which holds an interpolation object for dense output (default `nothing`). + - `dense` which indicates whether dense interpolation is available (default `false`). + + ## Callable Interface + + When `interp` is not `nothing`, the array supports callable syntax for interpolation: + + ```julia + da(t) # interpolate at time t + da(t, Val{1}) # first derivative at time t + da(t; idxs=1) # interpolate component 1 + da(t; idxs=[1,2]) # interpolate components 1 and 2 + da(t; continuity=:right) # right-continuity at discontinuities + ``` + + The interpolation object is called as `interp(t, idxs, deriv, p, continuity)`. """ abstract type AbstractDiffEqArray{T, N, A} <: AbstractVectorOfArray{T, N, A} end + """ + AbstractRaggedVectorOfArray{T, N, A} + + Abstract supertype for ragged (non-rectangular) vector-of-array types that + preserve the true ragged structure without zero-padding. Unlike + `AbstractVectorOfArray`, this does **not** subtype `AbstractArray` — indexing + returns actual stored data and `A[:, i]` gives the `i`-th inner array with + its original size. + + Concrete subtypes live in the `RecursiveArrayToolsRaggedArrays` subpackage + to avoid method invalidations on the hot path. + """ + abstract type AbstractRaggedVectorOfArray{T, N, A} end + + """ + AbstractRaggedDiffEqArray{T, N, A} <: AbstractRaggedVectorOfArray{T, N, A} + + Abstract supertype for ragged diff-eq arrays that carry a time vector `t`, + parameters `p`, and symbolic system `sys` alongside ragged solution data. + """ + abstract type AbstractRaggedDiffEqArray{T, N, A} <: AbstractRaggedVectorOfArray{T, N, A} end + include("utils.jl") include("vector_of_array.jl") include("array_partition.jl") include("named_array_partition.jl") - function Base.show(io::IO, x::Union{ArrayPartition, AbstractVectorOfArray}) + function Base.show(io::IO, x::ArrayPartition) return invoke(show, Tuple{typeof(io), Any}, io, x) end + # AbstractVectorOfArray uses AbstractArray's show import GPUArraysCore - Base.convert(T::Type{<:GPUArraysCore.AnyGPUArray}, VA::AbstractVectorOfArray) = stack(VA.u) - (T::Type{<:GPUArraysCore.AnyGPUArray})(VA::AbstractVectorOfArray) = T(Array(VA)) + # GPU conversion via stack (stays on device, avoids materializing dense Array). + # Only define convert — do NOT define a constructor (::Type{<:AbstractGPUArray})(::AbstractVectorOfArray) + # because it creates an unresolvable ambiguity with CUDA.jl's CuArray(::AbstractArray{T,N}). + # Instead, the KernelAbstractions extension defines the concrete CuArray dispatch. + Base.convert(::Type{T}, VA::AbstractVectorOfArray) where {T <: GPUArraysCore.AnyGPUArray} = T(stack(VA.u)) export VectorOfArray, VA, DiffEqArray, AbstractVectorOfArray, AbstractDiffEqArray, AllObserved, vecarr_to_vectors, tuples + # Plotting helpers (used by SciMLBase recipe delegation) + export DEFAULT_PLOT_FUNC, plottable_indices, plot_indices, getindepsym_defaultt, + interpret_vars, add_labels!, diffeq_to_arrays, solplot_vecs_and_labels + export recursivecopy, recursivecopy!, recursivefill!, vecvecapply, copyat_or_push!, vecvec_to_mat, recursive_one, recursive_mean, recursive_bottom_eltype, recursive_unitless_bottom_eltype, recursive_unitless_eltype diff --git a/src/array_partition.jl b/src/array_partition.jl index 8e2a97be..0f2eba96 100644 --- a/src/array_partition.jl +++ b/src/array_partition.jl @@ -223,12 +223,6 @@ end end end Base.filter(f, A::ArrayPartition) = ArrayPartition(map(x -> filter(f, x), A.x)) -Base.any(f, A::ArrayPartition) = any((any(f, x) for x in A.x)) -Base.any(f::Function, A::ArrayPartition) = any((any(f, x) for x in A.x)) -Base.any(A::ArrayPartition) = any(identity, A) -Base.all(f, A::ArrayPartition) = all((all(f, x) for x in A.x)) -Base.all(f::Function, A::ArrayPartition) = all((all(f, x) for x in A.x)) -Base.all(A::ArrayPartition) = all(identity, A) for type in [AbstractArray, PermutedDimsArray] @eval function Base.copyto!(dest::$(type), A::ArrayPartition) @@ -746,4 +740,10 @@ ODEProblem(func, AP[ [1.,2.,3.], [1. 2.;3. 4.] ], (0, 1)) |> solve """ struct AP end -Base.getindex(::Type{AP}, xs...) = ArrayPartition(xs...) + +# Optimized partition-level any/all for ArrayPartition lives in +# RecursiveArrayToolsArrayPartitionAnyAll to avoid ~780 invalidations. +# Without the extension, any/all uses the AbstractArray element-by-element +# fallback, which triggers scalar indexing errors on GPU arrays. +# Load the subpackage to fix: +# using RecursiveArrayToolsArrayPartitionAnyAll diff --git a/src/named_array_partition.jl b/src/named_array_partition.jl index ea11379d..b2546334 100644 --- a/src/named_array_partition.jl +++ b/src/named_array_partition.jl @@ -92,9 +92,10 @@ end Base.size(x::NamedArrayPartition) = size(ArrayPartition(x)) Base.length(x::NamedArrayPartition) = length(ArrayPartition(x)) -Base.getindex(x::NamedArrayPartition, args...) = getindex(ArrayPartition(x), args...) - -Base.setindex!(x::NamedArrayPartition, args...) = setindex!(ArrayPartition(x), args...) +# Delegate indexing to the underlying ArrayPartition. +# Use concrete index types to avoid invalidating AbstractArray's generic setindex!. +Base.@propagate_inbounds Base.getindex(x::NamedArrayPartition, i::Int) = ArrayPartition(x)[i] +Base.@propagate_inbounds Base.setindex!(x::NamedArrayPartition, v, i::Int) = (ArrayPartition(x)[i] = v) function Base.map(f, x::NamedArrayPartition) return NamedArrayPartition(map(f, ArrayPartition(x)), getfield(x, :names_to_indices)) end diff --git a/src/utils.jl b/src/utils.jl index 4f8b942e..e9d10295 100644 --- a/src/utils.jl +++ b/src/utils.jl @@ -133,7 +133,7 @@ function recursivefill!( T <: StaticArraysCore.StaticArray, T2 <: StaticArraysCore.StaticArray, N, } - return @inbounds for b in bs, i in eachindex(b) + return @inbounds for b in bs.u, i in eachindex(b) b[i] = copy(a) end @@ -159,7 +159,7 @@ function recursivefill!( T <: StaticArraysCore.SArray, T2 <: Union{Number, Bool}, N, } - return @inbounds for b in bs, i in eachindex(b) + return @inbounds for b in bs.u, i in eachindex(b) # Preserve static array shape while replacing all entries with the scalar b[i] = map(_ -> a, b[i]) diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index 8e416675..be0041d5 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -62,7 +62,7 @@ A.t ``` """ mutable struct DiffEqArray{ - T, N, A, B, F, S, D <: Union{Nothing, ParameterTimeseriesCollection}, + T, N, A, B, F, S, D <: Union{Nothing, ParameterTimeseriesCollection}, I, } <: AbstractDiffEqArray{T, N, A} u::A # A <: AbstractVector{<: AbstractArray{T, N - 1}} @@ -70,6 +70,8 @@ mutable struct DiffEqArray{ p::F sys::S discretes::D + interp::I + dense::Bool end ### Abstract Interface struct AllObserved @@ -87,7 +89,20 @@ function Base.Array( <:AbstractVector, }, } - return reduce(hcat, VA.u) + if allequal(length.(VA.u)) + return reduce(hcat, VA.u) + else + # Ragged: zero-padded + s = size(VA) + result = zeros(T, s) + for j in 1:length(VA.u) + u_j = VA.u[j] + for i in 1:length(u_j) + result[i, j] = u_j[i] + end + end + return result + end end function Base.Array( VA::AbstractVectorOfArray{ @@ -156,16 +171,40 @@ function Base.Vector( } return VA.u end -function Base.Array(VA::AbstractVectorOfArray) - vecs = vec.(VA.u) - return Array(reshape(reduce(hcat, vecs), size(VA.u[1])..., length(VA.u))) +function Base.Array(VA::AbstractVectorOfArray{T, N}) where {T, N} + if allequal(size.(VA.u)) + vecs = vec.(VA.u) + return Array(reshape(reduce(hcat, vecs), size(VA.u[1])..., length(VA.u))) + else + # Ragged: create zero-padded dense array + s = size(VA) + result = zeros(T, s) + for j in 1:length(VA.u) + u_j = VA.u[j] + for ci in CartesianIndices(size(u_j)) + result[ci, j] = u_j[ci] + end + end + return result + end end -function Base.Array{U}(VA::AbstractVectorOfArray) where {U} - vecs = vec.(VA.u) - return Array(reshape(reduce(hcat, vecs), size(VA.u[1])..., length(VA.u))) +function Base.Array{U}(VA::AbstractVectorOfArray{T, N}) where {U, T, N} + if allequal(size.(VA.u)) + vecs = vec.(VA.u) + return Array{U}(reshape(reduce(hcat, vecs), size(VA.u[1])..., length(VA.u))) + else + s = size(VA) + result = zeros(U, s) + for j in 1:length(VA.u) + u_j = VA.u[j] + for ci in CartesianIndices(size(u_j)) + result[ci, j] = U(u_j[ci]) + end + end + return result + end end -Base.convert(::Type{AbstractArray}, VA::AbstractVectorOfArray) = stack(VA.u) function Adapt.adapt_structure(to, VA::AbstractVectorOfArray) return VectorOfArray(Adapt.adapt.((to,), VA.u)) @@ -208,7 +247,8 @@ Base.parent(vec::VectorOfArray) = vec.u # first element representative function DiffEqArray( vec::AbstractVector, ts::AbstractVector; discretes = nothing, - variables = nothing, parameters = nothing, independent_variables = nothing + variables = nothing, parameters = nothing, independent_variables = nothing, + interp = nothing, dense = false ) sys = SymbolCache( something(variables, []), @@ -225,12 +265,15 @@ function DiffEqArray( Nothing, typeof(sys), typeof(discretes), + typeof(interp), }( vec, ts, nothing, sys, - discretes + discretes, + interp, + dense ) end @@ -238,7 +281,7 @@ end function DiffEqArray( vec::AbstractVector{VT}, ts::AbstractVector; discretes = nothing, variables = nothing, parameters = nothing, - independent_variables = nothing + independent_variables = nothing, interp = nothing, dense = false ) where {T, N, VT <: AbstractArray{T, N}} sys = SymbolCache( something(variables, []), @@ -253,12 +296,15 @@ function DiffEqArray( Nothing, typeof(sys), typeof(discretes), + typeof(interp), }( vec, ts, nothing, sys, - discretes + discretes, + interp, + dense ) end @@ -267,39 +313,46 @@ end # NTuple, T from type function DiffEqArray( vec::AbstractVector{T}, ts::AbstractVector, - ::NTuple{N, Int}; discretes = nothing + ::NTuple{N, Int}; discretes = nothing, interp = nothing, dense = false ) where {T, N} return DiffEqArray{ eltype(T), N, typeof(vec), typeof(ts), Nothing, Nothing, typeof(discretes), + typeof(interp), }( vec, ts, nothing, nothing, - discretes + discretes, + interp, + dense ) end # NTuple parameter function DiffEqArray( vec::AbstractVector{VT}, ts::AbstractVector, p::NTuple{N2, Int}; - discretes = nothing + discretes = nothing, interp = nothing, dense = false ) where {T, N, VT <: AbstractArray{T, N}, N2} return DiffEqArray{ eltype(T), N + 1, typeof(vec), typeof(ts), typeof(p), Nothing, typeof(discretes), + typeof(interp), }( vec, ts, p, nothing, - discretes + discretes, + interp, + dense ) end # first element representative function DiffEqArray( vec::AbstractVector, ts::AbstractVector, p; discretes = nothing, - variables = nothing, parameters = nothing, independent_variables = nothing + variables = nothing, parameters = nothing, independent_variables = nothing, + interp = nothing, dense = false ) sys = SymbolCache( something(variables, []), @@ -316,12 +369,15 @@ function DiffEqArray( typeof(p), typeof(sys), typeof(discretes), + typeof(interp), }( vec, ts, p, sys, - discretes + discretes, + interp, + dense ) end @@ -329,7 +385,7 @@ end function DiffEqArray( vec::AbstractVector{VT}, ts::AbstractVector, p; discretes = nothing, variables = nothing, parameters = nothing, - independent_variables = nothing + independent_variables = nothing, interp = nothing, dense = false ) where {T, N, VT <: AbstractArray{T, N}} sys = SymbolCache( something(variables, []), @@ -339,12 +395,15 @@ function DiffEqArray( return DiffEqArray{ eltype(T), N + 1, typeof(vec), typeof(ts), typeof(p), typeof(sys), typeof(discretes), + typeof(interp), }( vec, ts, p, sys, - discretes + discretes, + interp, + dense ) end @@ -353,38 +412,47 @@ end # NTuple, T from type function DiffEqArray( vec::AbstractVector{T}, ts::AbstractVector, - ::NTuple{N, Int}, p; discretes = nothing + ::NTuple{N, Int}, p; discretes = nothing, interp = nothing, dense = false ) where {T, N} return DiffEqArray{ eltype(T), N, typeof(vec), typeof(ts), typeof(p), Nothing, typeof(discretes), + typeof(interp), }( vec, ts, p, nothing, - discretes + discretes, + interp, + dense ) end # NTuple parameter function DiffEqArray( vec::AbstractVector{VT}, ts::AbstractVector, p::NTuple{N2, Int}, sys; - discretes = nothing + discretes = nothing, interp = nothing, dense = false ) where {T, N, VT <: AbstractArray{T, N}, N2} return DiffEqArray{ eltype(T), N + 1, typeof(vec), typeof(ts), typeof(p), typeof(sys), typeof(discretes), + typeof(interp), }( vec, ts, p, sys, - discretes + discretes, + interp, + dense ) end # first element representative -function DiffEqArray(vec::AbstractVector, ts::AbstractVector, p, sys; discretes = nothing) +function DiffEqArray( + vec::AbstractVector, ts::AbstractVector, p, sys; + discretes = nothing, interp = nothing, dense = false + ) _size = size(vec[1]) T = eltype(vec[1]) return DiffEqArray{ @@ -395,29 +463,35 @@ function DiffEqArray(vec::AbstractVector, ts::AbstractVector, p, sys; discretes typeof(p), typeof(sys), typeof(discretes), + typeof(interp), }( vec, ts, p, sys, - discretes + discretes, + interp, + dense ) end # T and N from type function DiffEqArray( vec::AbstractVector{VT}, ts::AbstractVector, p, sys; - discretes = nothing + discretes = nothing, interp = nothing, dense = false ) where {T, N, VT <: AbstractArray{T, N}} return DiffEqArray{ eltype(T), N + 1, typeof(vec), typeof(ts), typeof(p), typeof(sys), typeof(discretes), + typeof(interp), }( vec, ts, p, sys, - discretes + discretes, + interp, + dense ) end @@ -426,16 +500,19 @@ end # NTuple, T from type function DiffEqArray( vec::AbstractVector{T}, ts::AbstractVector, - ::NTuple{N, Int}, p, sys; discretes = nothing + ::NTuple{N, Int}, p, sys; discretes = nothing, interp = nothing, dense = false ) where {T, N} return DiffEqArray{ eltype(T), N, typeof(vec), typeof(ts), typeof(p), typeof(sys), typeof(discretes), + typeof(interp), }( vec, ts, p, sys, - discretes + discretes, + interp, + dense ) end @@ -447,154 +524,103 @@ function SymbolicIndexingInterface.is_parameter_timeseries( ::Type{ DiffEqArray{ T, N, A, B, - F, S, D, + F, S, D, I, }, } - ) where {T, N, A, B, F, S, D <: ParameterIndexingProxy} + ) where {T, N, A, B, F, S, D <: ParameterIndexingProxy, I} return Timeseries() end SymbolicIndexingInterface.state_values(A::AbstractDiffEqArray) = A.u SymbolicIndexingInterface.current_time(A::AbstractDiffEqArray) = A.t SymbolicIndexingInterface.parameter_values(A::AbstractDiffEqArray) = A.p +# Need explicit 2-arg method since AbstractDiffEqArray <: AbstractArray +# and SymbolicIndexingInterface defines parameter_values(::AbstractArray, i) = arr[i] +function SymbolicIndexingInterface.parameter_values(A::AbstractDiffEqArray, i) + return parameter_values(A.p, i) +end SymbolicIndexingInterface.symbolic_container(A::AbstractDiffEqArray) = A.sys function SymbolicIndexingInterface.get_parameter_timeseries_collection(A::AbstractDiffEqArray) return get_discretes(A) end -Base.IndexStyle(A::AbstractVectorOfArray) = Base.IndexStyle(typeof(A)) -Base.IndexStyle(::Type{<:AbstractVectorOfArray}) = IndexCartesian() - -@inline Base.length(VA::AbstractVectorOfArray) = length(VA.u) -@inline function Base.eachindex(VA::AbstractVectorOfArray) - return eachindex(VA.u) -end -@inline function Base.eachindex( - ::IndexLinear, VA::AbstractVectorOfArray{T, N, <:AbstractVector{T}} - ) where {T, N} - return eachindex(IndexLinear(), VA.u) -end -@inline Base.IteratorSize(::Type{<:AbstractVectorOfArray}) = Base.HasLength() -@inline Base.first(VA::AbstractVectorOfArray) = first(VA.u) -@inline Base.last(VA::AbstractVectorOfArray) = last(VA.u) -function Base.firstindex(VA::AbstractVectorOfArray{T, N, A}) where {T, N, A} - N > 1 && Base.depwarn( - "Linear indexing of `AbstractVectorOfArray` is deprecated. Change `A[i]` to `A.u[i]` ", - :firstindex - ) - return firstindex(VA.u) -end - -function Base.lastindex(VA::AbstractVectorOfArray{T, N, A}) where {T, N, A} - N > 1 && Base.depwarn( - "Linear indexing of `AbstractVectorOfArray` is deprecated. Change `A[i]` to `A.u[i]` ", - :lastindex - ) - return lastindex(VA.u) -end +## Callable interface for interpolation +# +# Any AbstractDiffEqArray with a non-nothing `interp` field supports `da(t)`. +# The interpolation object is called as `interp(t, idxs, deriv, p, continuity)`. +# SciMLBase's more-specific `(::AbstractODESolution)(t,...)` methods win dispatch +# for solution objects and handle symbolic idxs, discrete params, etc. -# 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::AbstractVectorOfArray, 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 +function (da::AbstractDiffEqArray)( + t, ::Type{deriv} = Val{0}; + idxs = nothing, continuity = :left + ) where {deriv} + da.interp === nothing && + error("No interpolation data is available. Provide an interpolation object via the `interp` keyword.") + return da.interp(t, idxs, deriv, da.p, continuity) end -Base.getindex(A::AbstractVectorOfArray, I::Int) = A.u[I] -Base.getindex(A::AbstractVectorOfArray, I::AbstractArray{Int}) = A.u[I] -Base.getindex(A::AbstractDiffEqArray, I::Int) = A.u[I] -Base.getindex(A::AbstractDiffEqArray, I::AbstractArray{Int}) = A.u[I] - -@deprecate Base.getindex( - VA::AbstractVectorOfArray{T, N, A}, - I::Int -) where {T, N, A <: Union{AbstractArray, AbstractVectorOfArray}} VA.u[I] false - -@deprecate Base.getindex( - VA::AbstractVectorOfArray{T, N, A}, - I::AbstractArray{Int} -) where {T, N, A <: Union{AbstractArray, AbstractVectorOfArray}} VA.u[I] false - -@deprecate Base.getindex( - VA::AbstractDiffEqArray{T, N, A}, - I::AbstractArray{Int} -) where {T, N, A <: Union{AbstractArray, AbstractVectorOfArray}} VA.u[I] false - -@deprecate Base.getindex( - VA::AbstractDiffEqArray{T, N, A}, - i::Int -) where {T, N, A <: Union{AbstractArray, AbstractVectorOfArray}} VA.u[i] false - -__parameterless_type(T) = Base.typename(T).wrapper - -# `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) - -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 - -# 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) +Base.IndexStyle(::Type{<:AbstractVectorOfArray}) = IndexCartesian() -struct RaggedRange - dim::Int - start::Int - step::Int - offset::Int -end -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 -function Base.:(:)(start::Integer, step::Integer, stop::RaggedEnd) - return RaggedRange(stop.dim, Int(start), Int(step), stop.offset) -end -function Base.:(:)(start::RaggedEnd, stop::RaggedEnd) - return RaggedRange(stop.dim, start.offset, 1, stop.offset) -end -function Base.:(:)(start::RaggedEnd, step::Integer, stop::RaggedEnd) - return RaggedRange(stop.dim, start.offset, Int(step), stop.offset) -end -function Base.:(:)(start::RaggedEnd, stop::Integer) - return RaggedRange(start.dim, start.offset, 1, Int(stop)) -end -function Base.:(:)(start::RaggedEnd, step::Integer, stop::Integer) - return RaggedRange(start.dim, start.offset, Int(step), Int(stop)) +## Linear indexing: convert to Cartesian and dispatch to the N-ary getindex +Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray{T, N}, i::Int) where {T, N} + @boundscheck checkbounds(A, i) + if N == 1 + return @inbounds A.u[i] + end + return @inbounds A[CartesianIndices(size(A))[i]] end -Base.broadcastable(x::RaggedRange) = Ref(x) -@inline function _is_ragged_dim(VA::AbstractVectorOfArray, 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 +Base.@propagate_inbounds function Base.setindex!(A::AbstractVectorOfArray{T, N}, v, i::Int) where {T, N} + @boundscheck checkbounds(A, i) + if N == 1 + A.u[i] = v + return v end - return false + ci = CartesianIndices(size(A))[i] + return @inbounds A[ci] = v end + Base.@propagate_inbounds function _getindex( - A::AbstractVectorOfArray, ::NotSymbolic, ::Colon, I::Int - ) - return A.u[I] + A::AbstractVectorOfArray{T, N}, ::NotSymbolic, ::Colon, I::Int + ) where {T, N} + u_col = A.u[I] + s = size(A) + leading_size = Base.front(s) + # If inner array matches the rectangular size, return directly + if size(u_col) == leading_size + return u_col + end + # If inner array has different ndims, return as-is (can't meaningfully reshape) + if ndims(u_col) != N - 1 + return u_col + end + # Zero-padded for ragged arrays with same ndims but different sizes + result = zeros(T, leading_size) + for ci in CartesianIndices(size(u_col)) + result[ci] = u_col[ci] + end + return result end Base.@propagate_inbounds function _getindex( - A::AbstractDiffEqArray, ::NotSymbolic, ::Colon, I::Int - ) - return A.u[I] + A::AbstractDiffEqArray{T, N}, ::NotSymbolic, ::Colon, I::Int + ) where {T, N} + u_col = A.u[I] + s = size(A) + leading_size = Base.front(s) + if size(u_col) == leading_size + return u_col + end + if ndims(u_col) != N - 1 + return u_col + end + result = zeros(T, leading_size) + for ci in CartesianIndices(size(u_col)) + result[ci] = u_col[ci] + end + return result end Base.@propagate_inbounds function _getindex( @@ -635,12 +661,20 @@ Base.@propagate_inbounds function _getindex( end end Base.@propagate_inbounds function _getindex( - VA::AbstractVectorOfArray, ::NotSymbolic, ii::CartesianIndex - ) + VA::AbstractVectorOfArray{T}, ::NotSymbolic, ii::CartesianIndex + ) where {T} ti = Tuple(ii) - i = last(ti) - jj = CartesianIndex(Base.front(ti)) - return VA.u[i][jj] + col = last(ti) + inner_I = Base.front(ti) + u_col = VA.u[col] + # Return zero for indices outside ragged storage + for d in 1:length(inner_I) + if inner_I[d] > size(u_col, d) + return zero(T) + end + end + jj = CartesianIndex(inner_I) + return u_col[jj] end Base.@propagate_inbounds function _getindex( @@ -709,296 +743,62 @@ Base.@propagate_inbounds function _getindex( return getindex(A, all_variable_symbols(A), args...) end -@inline _column_indices(VA::AbstractVectorOfArray, idx) = idx -@inline _column_indices(VA::AbstractVectorOfArray, idx::Colon) = eachindex(VA.u) -@inline function _column_indices(VA::AbstractVectorOfArray, idx::AbstractArray{Bool}) - return findall(idx) -end -@inline function _column_indices(VA::AbstractVectorOfArray, 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::AbstractVectorOfArray, 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, ::AbstractVectorOfArray, ::Any) = idx -@inline function _resolve_ragged_index(idx::RaggedEnd, VA::AbstractVectorOfArray, 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::AbstractVectorOfArray, 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::AbstractVectorOfArray, 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::AbstractVectorOfArray, col) - return Base.Slice(_resolve_ragged_index(idx.indices, VA, col)) -end -@inline function _resolve_ragged_index(idx::CartesianIndex, VA::AbstractVectorOfArray, col) - return CartesianIndex(_resolve_ragged_indices(Tuple(idx), VA, col)...) -end -@inline function _resolve_ragged_index( - idx::AbstractArray{<:RaggedEnd}, VA::AbstractVectorOfArray, col - ) - return map(i -> _resolve_ragged_index(i, VA, col), idx) -end -@inline function _resolve_ragged_index( - idx::AbstractArray{<:RaggedRange}, VA::AbstractVectorOfArray, col - ) - return map(i -> _resolve_ragged_index(i, VA, col), idx) -end -@inline function _resolve_ragged_index(idx::AbstractArray, VA::AbstractVectorOfArray, 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::AbstractVectorOfArray, 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::AbstractVectorOfArray, args::Tuple) - # Handle empty tuple case - length(args) == 0 && return args - if !_has_ragged_end(args...) - return args +# CartesianIndex with more dimensions than ndims(A) — for heterogeneous inner arrays +# where (inner_indices..., column_index) may have more entries than ndims(A) +Base.@propagate_inbounds function Base.getindex( + A::AbstractVectorOfArray{T, N}, ii::CartesianIndex + ) where {T, N} + ti = Tuple(ii) + if length(ti) == N + # Standard case: let AbstractArray handle via the N-ary method + return A[ti...] 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) + # Heterogeneous case: last element is column, rest are inner indices + col = last(ti) + inner_I = Base.front(ti) + u_col = A.u[col] + for d in 1:length(inner_I) + if inner_I[d] > size(u_col, d) + return zero(T) 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 DiffEqArray type when slicing -@inline function _preserve_array_type(A::AbstractVectorOfArray, u_slice, col_idxs) - return VectorOfArray(u_slice) -end - -@inline function _preserve_array_type(A::AbstractDiffEqArray, u_slice, col_idxs) - return DiffEqArray(u_slice, A.t[col_idxs], parameter_values(A), symbolic_container(A)) -end - -@inline function _ragged_getindex(A::AbstractVectorOfArray, 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 + return u_col[CartesianIndex(inner_I)] end -@inline function _ragged_getindex_nm1dims(A::AbstractVectorOfArray, 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) +Base.@propagate_inbounds function Base.setindex!( + A::AbstractVectorOfArray{T, N}, x, ii::CartesianIndex + ) where {T, N} + ti = Tuple(ii) + if length(ti) == N + return A[ti...] = x 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)..., + col = last(ti) + inner_I = Base.front(ti) + u_col = A.u[col] + for d in 1:length(inner_I) + if inner_I[d] > size(u_col, d) + iszero(x) && return x + throw( + ArgumentError( + "Cannot set non-zero value at index $ii: outside ragged storage bounds." ) - 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 DiffEqArray type if we're selecting actual columns, not inner dimensions - if is_column_selection - return _preserve_array_type(A, u_slice, cols) - else - return VectorOfArray(u_slice) - end - end -end - -@inline function _padded_resolved_indices(prefix, A::AbstractVectorOfArray, 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::AbstractVectorOfArray, 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 VectorOfArray 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::AbstractVectorOfArray, 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 + return u_col[CartesianIndex(inner_I)] = x end Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray, _arg, args...) + # Flatten CartesianIndex arguments (e.g. from sum(A; dims=d)) to plain Ints + # so they hit the N-ary getindex method instead of the symbolic dispatch. + if _arg isa Int && length(args) == 1 && args[1] isa CartesianIndex + return A[_arg, Tuple(args[1])...] + end 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...) @@ -1006,17 +806,10 @@ Base.@propagate_inbounds function Base.getindex(A::AbstractVectorOfArray, _arg, _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...) + _getindex(A, symtype, elsymtype, _arg, args...) end end -Base.@propagate_inbounds function Base.getindex( - A::Adjoint{T, <:AbstractVectorOfArray}, idxs... - ) where {T} - return getindex(A.parent, reverse(to_indices(A, idxs))...) -end function _observed(A::AbstractDiffEqArray{T, N}, sym, i::Int) where {T, N} return observed(A, sym)(A.u[i], A.p, A.t[i]) @@ -1035,15 +828,6 @@ Base.@propagate_inbounds function Base.setindex!( return VA.u[I] = v end -Base.@propagate_inbounds Base.setindex!(VA::AbstractVectorOfArray, v, I::Int) = Base.setindex!( - VA.u, v, I -) -@deprecate Base.setindex!( - VA::AbstractVectorOfArray{T, N, A}, v, - I::Int -) where {T, N, A <: Union{AbstractArray, AbstractVectorOfArray}} Base.setindex!( - VA.u, v, I -) false Base.@propagate_inbounds function Base.setindex!( VA::AbstractVectorOfArray{T, N}, v, @@ -1052,15 +836,6 @@ Base.@propagate_inbounds function Base.setindex!( return VA.u[I] = v end -Base.@propagate_inbounds Base.setindex!(VA::AbstractVectorOfArray, v, I::Colon) = Base.setindex!( - VA.u, v, I -) -@deprecate Base.setindex!( - VA::AbstractVectorOfArray{T, N, A}, v, - I::Colon -) where {T, N, A <: Union{AbstractArray, AbstractVectorOfArray}} Base.setindex!( - VA.u, v, I -) false Base.@propagate_inbounds function Base.setindex!( VA::AbstractVectorOfArray{T, N}, v, @@ -1069,15 +844,6 @@ Base.@propagate_inbounds function Base.setindex!( return VA.u[I] = v end -Base.@propagate_inbounds Base.setindex!(VA::AbstractVectorOfArray, v, I::AbstractArray{Int}) = Base.setindex!( - VA.u, v, I -) -@deprecate Base.setindex!( - VA::AbstractVectorOfArray{T, N, A}, v, - I::AbstractArray{Int} -) where {T, N, A <: Union{AbstractArray, AbstractVectorOfArray}} Base.setindex!( - VA, v, :, I -) false Base.@propagate_inbounds function Base.setindex!( VA::AbstractVectorOfArray{T, N}, v, i::Int, @@ -1088,20 +854,11 @@ Base.@propagate_inbounds function Base.setindex!( end return v end -Base.@propagate_inbounds function Base.setindex!( - VA::AbstractVectorOfArray{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::AbstractVectorOfArray{T, N}, x, - idxs::Union{Int, Colon, CartesianIndex, AbstractArray{Int}, AbstractArray{Bool}}... + idxs::Union{Int, Colon, AbstractArray{Int}, AbstractArray{Bool}}... ) where { T, N, } @@ -1116,34 +873,60 @@ Base.@propagate_inbounds function Base.setindex!( return x end -# Interface for the two-dimensional indexing, a more standard AbstractArray interface -@inline Base.size(VA::AbstractVectorOfArray) = (size(VA.u[1])..., length(VA.u)) -@inline Base.size(VA::AbstractVectorOfArray, i) = size(VA)[i] -@inline Base.size(A::Adjoint{T, <:AbstractVectorOfArray}) where {T} = reverse(size(A.parent)) -@inline Base.size(A::Adjoint{T, <:AbstractVectorOfArray}, i) where {T} = size(A)[i] -Base.axes(VA::AbstractVectorOfArray) = Base.OneTo.(size(VA)) -Base.axes(VA::AbstractVectorOfArray, d::Int) = Base.OneTo(size(VA)[d]) +# Interface for the AbstractArray interface +@inline function Base.size(VA::AbstractVectorOfArray{T, N}) where {T, N} + isempty(VA.u) && return ntuple(_ -> 0, Val(N)) + leading = ntuple(Val(N - 1)) do d + maximum(size(u, d) for u in VA.u) + end + return (leading..., length(VA.u)) +end Base.@propagate_inbounds function Base.setindex!( VA::AbstractVectorOfArray{T, N}, v, I::Int... ) where {T, N} - return VA.u[I[end]][Base.front(I)...] = v + col = I[end] + inner_I = Base.front(I) + u_col = VA.u[col] + # Check if within ragged storage bounds + for d in 1:length(inner_I) + if inner_I[d] > size(u_col, d) + iszero(v) && return v + throw( + ArgumentError( + "Cannot set non-zero value at index $I: outside ragged storage bounds. " * + "Inner array $col has size $(size(u_col)) but index requires $(inner_I)." + ) + ) + end + end + return u_col[inner_I...] = v +end + +# Core N-dimensional getindex for AbstractArray interface +# Handles ragged arrays by returning zero for out-of-bounds inner indices +Base.@propagate_inbounds function Base.getindex( + A::AbstractVectorOfArray{T, N}, I::Vararg{Int, N} + ) where {T, N} + @boundscheck checkbounds(A, I...) + col = I[N] + inner_I = Base.front(I) + u_col = A.u[col] + # Return zero for indices outside ragged storage + for d in 1:(N - 1) + if inner_I[d] > size(u_col, d) + return zero(T) + end + end + return @inbounds u_col[inner_I...] end function Base.:(==)(A::AbstractVectorOfArray, B::AbstractVectorOfArray) return A.u == B.u end -function Base.:(==)(A::AbstractVectorOfArray, B::AbstractArray) - return A.u == B -end -Base.:(==)(A::AbstractArray, B::AbstractVectorOfArray) = B == A +# Comparison with plain arrays uses AbstractArray element-wise comparison via default -# The iterator will be over the subarrays of the container, not the individual elements -# unlike an true AbstractArray -function Base.iterate(VA::AbstractVectorOfArray, state = 1) - return state >= length(VA.u) + 1 ? nothing : (VA[:, state], state + 1) -end tuples(VA::DiffEqArray) = tuple.(VA.t, VA.u) # Growing the array simply adds to the container vector @@ -1161,9 +944,10 @@ function Base.copy(VA::AbstractVectorOfArray) end function Base.zero(VA::AbstractVectorOfArray) - val = copy(VA) - val.u .= zero.(VA.u) - return val + T = typeof(VA) + u_zero = [zero(u) for u in VA.u] + fields = [fname == :u ? u_zero : _copyfield(VA, fname) for fname in fieldnames(T)] + return T(fields...) end Base.sizehint!(VA::AbstractVectorOfArray{T, N}, i) where {T, N} = sizehint!(VA.u, i) @@ -1196,7 +980,7 @@ function Base.append!( VA::AbstractVectorOfArray{T, N}, new_item::AbstractVectorOfArray{T, N} ) where {T, N} - for item in copy(new_item) + for item in copy(new_item.u) push!(VA, item) end return VA @@ -1243,51 +1027,7 @@ function Base.view(A::AbstractVectorOfArray, I::Vararg{Any, M}) where {M} @boundscheck checkbounds(A, J...) return SubArray(A, J) end -function Base.SubArray(parent::AbstractVectorOfArray, indices::Tuple) - @inline - return SubArray( - IndexStyle(Base.viewindexing(indices), IndexStyle(parent)), parent, - Base.ensure_indexable(indices), Base.index_dimsum(indices...) - ) -end -Base.isassigned(VA::AbstractVectorOfArray, idxs...) = checkbounds(Bool, VA, idxs...) -function Base.check_parent_index_match( - ::RecursiveArrayTools.AbstractVectorOfArray{T, N}, ::NTuple{N, Bool} - ) where {T, N} - return nothing -end -Base.ndims(::AbstractVectorOfArray{T, N}) where {T, N} = N -Base.ndims(::Type{<:AbstractVectorOfArray{T, N}}) where {T, N} = N -function Base.checkbounds( - ::Type{Bool}, VA::AbstractVectorOfArray{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::AbstractVectorOfArray, 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::AbstractVectorOfArray, idx...) - return checkbounds(Bool, VA, idx...) || throw(BoundsError(VA, idx)) -end function Base.copyto!( dest::AbstractVectorOfArray{T, N}, src::AbstractVectorOfArray{T2, N} @@ -1320,72 +1060,10 @@ function Base.copyto!( copyto!(dest.u, src) return dest end -# Required for broadcasted setindex! when slicing across subarrays -# E.g. if `va = VectorOfArray([rand(3, 3) for i in 1:5])` -# Need this method for `va[2, :, :] .= 3.0` -Base.@propagate_inbounds function Base.maybeview(A::AbstractVectorOfArray, I...) - return view(A, I...) -end - -# Operations -function Base.isapprox( - A::AbstractVectorOfArray, - B::Union{AbstractVectorOfArray, AbstractArray}; - kwargs... - ) - return all(isapprox.(A, B; kwargs...)) -end - -function Base.isapprox(A::AbstractArray, B::AbstractVectorOfArray; kwargs...) - return all(isapprox.(A, B; kwargs...)) -end - -for op in [:(Base.:-), :(Base.:+)] - @eval function ($op)(A::AbstractVectorOfArray, B::AbstractVectorOfArray) - return ($op).(A, B) - end - @eval Base.@propagate_inbounds function ($op)( - A::AbstractVectorOfArray, - B::AbstractArray - ) - @boundscheck length(A) == length(B) - return VectorOfArray([($op).(a, b) for (a, b) in zip(A, B)]) - end - @eval Base.@propagate_inbounds function ($op)( - A::AbstractArray, B::AbstractVectorOfArray - ) - @boundscheck length(A) == length(B) - return VectorOfArray([($op).(a, b) for (a, b) in zip(A, B)]) - end -end - -for op in [:(Base.:/), :(Base.:\), :(Base.:*)] - if op !== :(Base.:/) - @eval ($op)(A::Number, B::AbstractVectorOfArray) = ($op).(A, B) - end - if op !== :(Base.:\) - @eval ($op)(A::AbstractVectorOfArray, B::Number) = ($op).(A, B) - end -end - -function Base.CartesianIndices(VA::AbstractVectorOfArray) - if !allequal(size.(VA.u)) - error("CartesianIndices only valid for non-ragged arrays") - end - return CartesianIndices((size(VA.u[1])..., length(VA.u))) -end # Tools for creating similar objects -Base.eltype(::Type{<:AbstractVectorOfArray{T}}) where {T} = T - -@inline function Base.similar(VA::AbstractVectorOfArray, args...) - if args[end] isa Type - return Base.similar(eltype(VA)[], args..., size(VA)) - else - return Base.similar(eltype(VA)[], args...) - end -end +# similar(VA) - same type and size function Base.similar( vec::VectorOfArray{ T, N, AT, @@ -1403,7 +1081,8 @@ function Base.similar( return VectorOfArray(similar(Base.parent(vec))) end -@inline function Base.similar(VA::VectorOfArray, ::Type{T} = eltype(VA)) where {T} +# similar(VA, T) - same structure, different element type +@inline function Base.similar(VA::VectorOfArray, ::Type{T}) where {T} if eltype(VA.u) <: Union{AbstractArray, AbstractVectorOfArray} return VectorOfArray(similar.(VA.u, T)) else @@ -1411,15 +1090,22 @@ end end end -@inline function Base.similar(VA::VectorOfArray, dims::N) where {N <: Number} - l = length(VA) - return if dims <= l - VectorOfArray(similar.(VA.u[1:dims])) - else - VectorOfArray([similar.(VA.u); [similar(VA.u[end]) for _ in (l + 1):dims]]) - end +# similar(VA, T, dims) - return a regular Array (standard AbstractArray behavior) +@inline function Base.similar( + VA::AbstractVectorOfArray, ::Type{T}, dims::Tuple{Vararg{Int}} + ) where {T} + return Array{T}(undef, dims...) +end +@inline function Base.similar( + VA::AbstractVectorOfArray, ::Type{T}, dims::Tuple{ + Union{Integer, Base.OneTo}, + Vararg{Union{Integer, Base.OneTo}}, + } + ) where {T} + return similar(Array{T}, dims) end + # fill! # For DiffEqArray it ignores ts and fills only u function Base.fill!(VA::AbstractVectorOfArray, x) @@ -1437,38 +1123,20 @@ function Base.fill!(VA::AbstractVectorOfArray, x) end return VA end - -Base.reshape(A::AbstractVectorOfArray, dims...) = Base.reshape(Array(A), dims...) - -# Need this for ODE_DEFAULT_UNSTABLE_CHECK from DiffEqBase to work properly -@inline Base.any(f, VA::AbstractVectorOfArray) = any(any(f, u) for u in VA.u) -@inline Base.all(f, VA::AbstractVectorOfArray) = all(all(f, u) for u in VA.u) - # conversion tools vecarr_to_vectors(VA::AbstractVectorOfArray) = [VA[i, :] for i in eachindex(VA.u[1])] -Base.vec(VA::AbstractVectorOfArray) = vec(convert(Array, VA)) # Allocates -# stack non-ragged arrays to convert them -function Base.convert(::Type{Array}, VA::AbstractVectorOfArray) - if !allequal(size.(VA.u)) - error("Can only convert non-ragged VectorOfArray to Array") - end - return Array(VA) -end +# linear algebra +ArrayInterface.issingular(va::AbstractVectorOfArray) = ArrayInterface.issingular(Matrix(va)) -# statistics -@inline Base.sum(VA::AbstractVectorOfArray; kwargs...) = sum(identity, VA; kwargs...) -@inline function Base.sum(f, VA::AbstractVectorOfArray; kwargs...) - return mapreduce(f, Base.add_sum, VA; kwargs...) +# Type-stable sum/mapreduce that avoids inference issues on Julia 1.10 +# with deeply nested VectorOfArray type parameters +function Base.sum(VA::AbstractVectorOfArray{T}) where {T} + return sum(sum, VA.u)::T end -@inline Base.prod(VA::AbstractVectorOfArray; kwargs...) = prod(identity, VA; kwargs...) -@inline function Base.prod(f, VA::AbstractVectorOfArray; kwargs...) - return mapreduce(f, Base.mul_prod, VA; kwargs...) -end - -@inline Base.adjoint(VA::AbstractVectorOfArray) = Adjoint(VA) -# linear algebra -ArrayInterface.issingular(va::AbstractVectorOfArray) = ArrayInterface.issingular(Matrix(va)) +function Base.sum(f::F, VA::AbstractVectorOfArray{T}) where {F, T} + return sum(u -> sum(f, u), VA.u) +end # make it show just like its data function Base.show(io::IO, m::MIME"text/plain", x::AbstractVectorOfArray) @@ -1486,28 +1154,285 @@ end @recipe function f(VA::AbstractVectorOfArray) convert(Array, VA) end -@recipe function f(VA::AbstractDiffEqArray) - xguide --> isempty(independent_variable_symbols(VA)) ? "" : - independent_variable_symbols(VA)[1] - label --> isempty(variable_symbols(VA)) ? "" : - reshape(string.(variable_symbols(VA)), 1, :) - VA.t, VA' + +## Plotting helper functions (shared with SciMLBase) + +DEFAULT_PLOT_FUNC(x, y) = (x, y) +DEFAULT_PLOT_FUNC(x, y, z) = (x, y, z) + +plottable_indices(x::AbstractArray) = 1:length(x) +plottable_indices(x::Number) = 1 +plot_indices(A::AbstractArray) = eachindex(A) + +""" + getindepsym_defaultt(A) + +Return the independent variable symbol for `A`, defaulting to `:t`. +""" +function getindepsym_defaultt(A) + syms = independent_variable_symbols(A) + return isempty(syms) ? :t : first(syms) end -@recipe function f(VA::DiffEqArray{T, 1}) where {T} - VA.t, VA.u + +""" + interpret_vars(vars, A) + +Normalize user-provided variable specifications into a standard internal format: +a list of tuples `(func, xvar, yvar[, zvar])`. Index `0` represents the +independent variable (time). +""" +function interpret_vars(vars, A) + if vars === nothing + if A[:, 1] isa Union{Tuple, AbstractArray} + vars = collect((DEFAULT_PLOT_FUNC, 0, i) for i in plot_indices(A[:, 1])) + else + vars = [(DEFAULT_PLOT_FUNC, 0, 1)] + end + end + + if vars isa Base.Integer + vars = [(DEFAULT_PLOT_FUNC, 0, vars)] + end + + if vars isa AbstractArray + tmp = Tuple[] + for x in vars + if x isa Tuple + if x[1] isa Int + push!(tmp, tuple(DEFAULT_PLOT_FUNC, x...)) + else + push!(tmp, x) + end + else + push!(tmp, (DEFAULT_PLOT_FUNC, 0, x)) + end + end + vars = tmp + end + + if vars isa Tuple + if vars[end - 1] isa AbstractArray + if vars[end] isa AbstractArray + vars = collect( + zip( + [DEFAULT_PLOT_FUNC for i in eachindex(vars[end - 1])], + vars[end - 1], vars[end] + ) + ) + else + vars = [(DEFAULT_PLOT_FUNC, x, vars[end]) for x in vars[end - 1]] + end + else + if vars[2] isa AbstractArray + vars = [(DEFAULT_PLOT_FUNC, vars[end - 1], y) for y in vars[end]] + else + if vars[1] isa Int || symbolic_type(vars[1]) != NotSymbolic() + vars = [tuple(DEFAULT_PLOT_FUNC, vars...)] + else + vars = [vars] + end + end + end + end + + @assert(typeof(vars) <: AbstractArray) + @assert(eltype(vars) <: Tuple) + return vars end -Base.map(f, A::RecursiveArrayTools.AbstractVectorOfArray) = map(f, A.u) +function _var_label(A, x) + varsyms = variable_symbols(A) + if symbolic_type(x) != NotSymbolic() + return string(x) + elseif !isempty(varsyms) && x isa Integer && x > 0 && x <= length(varsyms) + return string(varsyms[x]) + elseif x isa Integer + return x == 0 ? "t" : "u[$x]" + else + return string(x) + end +end -function Base.mapreduce(f, op, A::AbstractVectorOfArray; kwargs...) - return mapreduce(f, op, view(A, ntuple(_ -> :, ndims(A))...); kwargs...) +function add_labels!(labels, x, dims, A, strs) + if ((x[2] isa Integer && x[2] == 0) || isequal(x[2], getindepsym_defaultt(A))) && + dims == 2 + push!(labels, strs[end]) + elseif x[1] !== DEFAULT_PLOT_FUNC + push!(labels, "f($(join(strs, ',')))") + else + push!(labels, "($(join(strs, ',')))") + end + return labels end -function Base.mapreduce( - f, op, A::AbstractVectorOfArray{T, 1, <:AbstractVector{T}}; kwargs... - ) where {T} - return mapreduce(f, op, A.u; kwargs...) + +""" + diffeq_to_arrays(A, denseplot, plotdensity, tspan, vars, tscale, plotat) + +Convert an `AbstractDiffEqArray` into plot-ready arrays. Returns `(plot_vecs, labels)`. +""" +function diffeq_to_arrays( + A, denseplot, plotdensity, tspan, vars, tscale, plotat + ) + if tspan === nothing + start_idx = 1 + end_idx = length(A.u) + else + start_idx = searchsortedfirst(A.t, tspan[1]) + end_idx = searchsortedlast(A.t, tspan[end]) + end + + densetspacer = if tscale in [:ln, :log10, :log2] + (start, stop, n) -> exp10.(range(log10(start), stop = log10(stop), length = n)) + else + (start, stop, n) -> range(start; stop = stop, length = n) + end + + if plotat !== nothing + plott = plotat + elseif denseplot + if tspan === nothing + plott = collect(densetspacer(A.t[start_idx], A.t[end_idx], plotdensity)) + else + plott = collect(densetspacer(tspan[1], tspan[end], plotdensity)) + end + else + plott = A.t[start_idx:end_idx] + end + + dims = length(vars[1]) - 1 + for var in vars + @assert length(var) - 1 == dims + end + return solplot_vecs_and_labels(dims, vars, plott, A) +end + +function solplot_vecs_and_labels(dims, vars, plott, A) + plot_vecs = [] + labels = String[] + batch_symbolic_vars = [] + for x in vars + for j in 2:length(x) + if (x[j] isa Integer && x[j] == 0) || isequal(x[j], getindepsym_defaultt(A)) + else + push!(batch_symbolic_vars, x[j]) + end + end + end + batch_symbolic_vars = identity.(batch_symbolic_vars) + + # Use callable if available, otherwise index directly + has_interp = hasproperty(A, :interp) && A.interp !== nothing + if has_interp + indexed_solution = A(plott; idxs = batch_symbolic_vars) + else + # For non-interpolating DiffEqArrays, find matching time indices + # plott should be a subset of A.t in this case + indexed_solution = A + end + + idxx = 0 + for x in vars + tmp = [] + strs = String[] + for j in 2:length(x) + if (x[j] isa Integer && x[j] == 0) || isequal(x[j], getindepsym_defaultt(A)) + push!(tmp, plott) + push!(strs, "t") + else + idxx += 1 + if has_interp + push!(tmp, indexed_solution[idxx, :]) + else + # Direct indexing: extract component from each time point + idx = batch_symbolic_vars[idxx] + if idx isa Integer + push!(tmp, [A.u[ti][idx] for ti in eachindex(plott)]) + else + push!(tmp, [A[idx, ti] for ti in eachindex(plott)]) + end + end + push!(strs, _var_label(A, x[j])) + end + end + + f = x[1] + tmp = map(f, tmp...) + tmp = tuple((getindex.(tmp, i) for i in eachindex(tmp[1]))...) + for i in eachindex(tmp) + if length(plot_vecs) < i + push!(plot_vecs, []) + end + push!(plot_vecs[i], tmp[i]) + end + add_labels!(labels, x, dims, A, strs) + end + + plot_vecs = [hcat(x...) for x in plot_vecs] + return plot_vecs, labels end +@recipe function f( + VA::AbstractDiffEqArray; + denseplot = ( + hasproperty(VA, :dense) && VA.dense && + hasproperty(VA, :interp) && VA.interp !== nothing + ), + plotdensity = max(1000, 10 * length(VA.u)), + tspan = nothing, plotat = nothing, + idxs = nothing + ) + + idxs_input = idxs === nothing ? plottable_indices(VA.u[1]) : idxs + if !(idxs_input isa Union{Tuple, AbstractArray}) + vars = interpret_vars([idxs_input], VA) + else + vars = interpret_vars(idxs_input, VA) + end + + tdir = sign(VA.t[end] - VA.t[1]) + xflip --> tdir < 0 + seriestype --> :path + + tscale = get(plotattributes, :xscale, :identity) + plot_vecs, labels = diffeq_to_arrays( + VA, denseplot, plotdensity, tspan, vars, tscale, plotat + ) + + # Axis labels for tuple-style idxs + if idxs_input isa Tuple && vars[1][1] === DEFAULT_PLOT_FUNC + for (guide, idx) in [(:xguide, 2), (:yguide, 3)] + if idx <= length(vars[1]) + guide --> _var_label(VA, vars[1][idx]) + end + end + if length(vars[1]) > 3 + zguide --> _var_label(VA, vars[1][4]) + end + end + + # Default xguide for time-vs-variable plots + if all( + x -> (x[2] isa Integer && x[2] == 0) || + isequal(x[2], getindepsym_defaultt(VA)), vars + ) + xguide --> "$(getindepsym_defaultt(VA))" + if tspan === nothing + if tdir > 0 + xlims --> (VA.t[1], VA.t[end]) + else + xlims --> (VA.t[end], VA.t[1]) + end + else + xlims --> (tspan[1], tspan[end]) + end + end + + label --> reshape(labels, 1, length(labels)) + (plot_vecs...,) +end +@recipe function f(VA::DiffEqArray{T, 1}) where {T} + VA.t, VA.u +end ## broadcasting struct VectorOfArrayStyle{N} <: Broadcast.AbstractArrayStyle{N} end # N is only used when voa sees other abstract arrays @@ -1537,8 +1462,6 @@ end function Broadcast.BroadcastStyle(::Type{<:AbstractVectorOfArray{T, N}}) where {T, N} return VectorOfArrayStyle{N}() end -# make vectorofarrays broadcastable so they aren't collected -Broadcast.broadcastable(x::AbstractVectorOfArray) = x # recurse through broadcast arguments and return a parent array for # the first VoA or DiffEqArray in the bc arguments @@ -1587,13 +1510,17 @@ for (type, N_expr) in [ 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)) + # Use dest.u[i] directly to get the actual inner array without + # zero-padding (dest[:, i] zero-pads ragged arrays, causing + # DimensionMismatch when the source inner array has a different size) + inner = dest.u[i] + if inner isa AbstractArray + if ArrayInterface.ismutable(inner) + copyto!(inner, 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 = StaticArraysCore.similar_type(inner) + dest.u[i] = if length(unpacked) == 1 && length(inner) == 1 arr_type(unpacked[1]) elseif length(unpacked) == 1 fill(copy(unpacked), arr_type) @@ -1602,7 +1529,7 @@ for (type, N_expr) in [ end end else - dest[:, i] = copy(unpack_voa(bc, i)) + dest.u[i] = copy(unpack_voa(bc, i)) end end return dest @@ -1682,4 +1609,3 @@ nested = VA[ """ struct VA end -Base.getindex(::Type{VA}, xs...) = VectorOfArray(collect(xs)) diff --git a/test/adjoints.jl b/test/adjoints.jl index 6eefb906..fec7f548 100644 --- a/test/adjoints.jl +++ b/test/adjoints.jl @@ -1,5 +1,4 @@ using RecursiveArrayTools, Zygote, ForwardDiff, Test -using SciMLBase function loss(x) return sum(abs2, Array(VectorOfArray([x .* i for i in 1:5]))) @@ -33,12 +32,6 @@ function loss5(x) return sum(abs2, Array(ArrayPartition([x .* i for i in 1:5]...))) end -function loss6(x) - _x = ArrayPartition([x .* i for i in 1:5]...) - _prob = ODEProblem((u, p, t) -> u, _x, (0, 1)) - return sum(abs2, Array(_prob.u0)) -end - function loss7(x) _x = VectorOfArray([x .* i for i in 1:5]) return sum(abs2, _x .- 1) @@ -85,7 +78,6 @@ loss(x) @test Zygote.gradient(loss3, x)[1] == ForwardDiff.gradient(loss3, x) @test Zygote.gradient(loss4, x)[1] == ForwardDiff.gradient(loss4, x) @test Zygote.gradient(loss5, x)[1] == ForwardDiff.gradient(loss5, x) -@test Zygote.gradient(loss6, x)[1] == ForwardDiff.gradient(loss6, x) @test Zygote.gradient(loss7, x)[1] == ForwardDiff.gradient(loss7, x) @test Zygote.gradient(loss8, x)[1] == ForwardDiff.gradient(loss8, x) @test ForwardDiff.derivative(loss9, 0.0) == diff --git a/test/basic_indexing.jl b/test/basic_indexing.jl index 4ead8ee8..f3b05cbd 100644 --- a/test/basic_indexing.jl +++ b/test/basic_indexing.jl @@ -1,4 +1,5 @@ using RecursiveArrayTools, Test +using RecursiveArrayToolsShorthandConstructors # Example Problem recs = [[1, 2, 3], [4, 5, 6], [7, 8, 9]] @@ -12,7 +13,7 @@ testva = VA[[1, 2, 3], [4, 5, 6], [7, 8, 9]] # broadcast with array X = rand(3, 3) mulX = sqrt.(abs.(testva .* X)) -ref = mapreduce((x, y) -> sqrt.(abs.(x .* y)), hcat, testva, eachcol(X)) +ref = sqrt.(abs.(Array(testva) .* X)) @test mulX == ref fill!(mulX, 0) mulX .= sqrt.(abs.(testva .* X)) @@ -55,28 +56,24 @@ testvasim = similar(testva, Float32) @test size(testvasim) == size(testva) @test eltype(testvasim) == Float32 testvb = deepcopy(testva) -@test testva == testvb == recs +@test testva == testvb +@test testva.u == 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 - -# ## Linear indexing -@test_deprecated testva[1] -@test_deprecated testva[1:2] -@test_deprecated testva[begin] -@test_deprecated testva[end] -@test testva[:, begin] == first(testva) -@test testva[:, end] == last(testva) -@test testa[:, 1] == recs[1] +@test Array(testva + testvb) == Array(testva) + Array(testvb) +@test Array(2testva) == 2 .* Array(testva) +@test Array(testva / 2) == Array(testva) ./ 2 + +# ## Linear indexing (now matches AbstractArray behavior) +@test testva[1] == testa[1] # first element, column-major +@test testva[end] == testa[end] # last element +@test testva[begin] == testa[begin] # first element +@test testva[:, 1] == recs[1] @test testva.u == recs @test testva[:, 2:end] == VectorOfArray([recs[i] for i in 2:length(recs)]) diffeq = DiffEqArray(recs, t) -@test_deprecated diffeq[1] -@test_deprecated diffeq[1:2] +@test diffeq[1] == testa[1] # linear indexing now matches AbstractArray @test diffeq[:, 1] == testa[:, 1] @test diffeq.u == recs @test diffeq[:, end] == testa[:, end] @@ -145,158 +142,117 @@ diffeq = DiffEqArray(recs, t) @test testa[1:1, 2:3, 1:2] == testva[1:1, 2:3, 1:2] @test testa[1:1, 2:3, 1:2] == diffeq[1:1, 2:3, 1:2] -# ## Test ragged arrays work, or give errors as needed -#TODO: I am not really sure what the behavior of this is, what does Mathematica do? +# ## Test ragged arrays with zero-padded rectangular interpretation t = 1:3 recs = [[1, 2, 3], [3, 5, 6, 7], [8, 9, 10, 11]] -testva = VectorOfArray(recs) #TODO: clearly this printed form is nonsense +testva = VectorOfArray(recs) diffeq = DiffEqArray(recs, t) -@test testva[:, 1] == recs[1] +# size uses maximum inner length +@test size(testva) == (4, 3) +# [:, 1] is zero-padded to length 4 +@test testva[:, 1] == [1, 2, 3, 0] @test testva[1:2, 1:2] == [1 3; 2 5] -@test diffeq[:, 1] == recs[1] +@test diffeq[:, 1] == [1, 2, 3, 0] @test diffeq[1:2, 1:2] == [1 3; 2 5] -@test diffeq[:, 1:2] == DiffEqArray([recs[i] for i in 1:2], t[1:2]) -@test diffeq[:, 1:2].t == t[1:2] -@test diffeq[:, 2:end] == DiffEqArray([recs[i] for i in 2:3], t[2:end]) -@test diffeq[:, 2:end].t == t[2:end] -@test diffeq[:, (end - 1):end] == DiffEqArray([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 views of heterogeneous arrays (issue #453) +# Test views of heterogeneous arrays - now use rectangular size with zero padding f = VA[[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] +# size(f) = (2, 2), view(:, 1) has length 2 +@test size(f) == (2, 2) +@test f[:, 1] == [1.0, 0.0] # zero-padded +@test f[:, 2] == [2.0, 3.0] f2 = VA[[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] - -# Test `end` with ragged arrays +@test size(f2) == (2, 2) +@test f2[:, 1] == [1.0, 2.0] +@test f2[:, 2] == [3.0, 0.0] # zero-padded + +# Test `end` with ragged arrays - now `end` in dim 1 = max size across all inner arrays ragged = VA[[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 +# size(ragged) = (4, 3) since max inner length is 4 +@test size(ragged) == (4, 3) +# end in dim 1 = 4 (max size), so ragged[end, 1] = ragged[4, 1] = 0.0 (ragged zero) +@test ragged[end, 1] == 0.0 +@test ragged[end, 2] == 0.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[2, 1] == 2.0 +@test ragged[3, 2] == 5.0 +@test ragged[4, 3] == 9.0 +# ragged[:, 1] now returns all 4 elements (with zero padding) +@test ragged[:, 1] == [1.0, 2.0, 0.0, 0.0] +@test ragged[:, 2] == [3.0, 4.0, 5.0, 0.0] +@test ragged[:, 3] == [6.0, 7.0, 8.0, 9.0] @test ragged[:, end] == [6.0, 7.0, 8.0, 9.0] -@test ragged[:, 2:end] == VectorOfArray(ragged.u[2:end]) -@test ragged[:, (end - 1):end] == VectorOfArray(ragged.u[(end - 1):end]) +# Array conversion with zero-padding +@test Array(ragged) == [1.0 3.0 6.0; 2.0 4.0 7.0; 0.0 5.0 8.0; 0.0 0.0 9.0] ragged2 = VA[[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] == VectorOfArray(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] == VectorOfArray(ragged2.u[(end - 1):end]) - -# Test that RaggedEnd and RaggedRange broadcast as scalars -# (fixes issue with SymbolicIndexingInterface where broadcasting over RaggedEnd would fail) +# size(ragged2) = (4, 3) since max inner length is 4 +@test size(ragged2) == (4, 3) +@test ragged2[end, 1] == 4.0 # ragged2[4, 1] = 4.0 +@test ragged2[end, 2] == 0.0 # ragged2[4, 2] = 0.0 (ragged zero) +@test ragged2[end, 3] == 0.0 # ragged2[4, 3] = 0.0 (ragged zero) +@test ragged2[:, 1] == [1.0, 2.0, 3.0, 4.0] +@test ragged2[:, 2] == [5.0, 6.0, 0.0, 0.0] +@test ragged2[:, 3] == [7.0, 8.0, 9.0, 0.0] +@test ragged2[:, end] == [7.0, 8.0, 9.0, 0.0] + +# lastindex now returns plain integers since size uses max sizes ragged_idx = lastindex(ragged, 1) -@test ragged_idx isa RecursiveArrayTools.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 RecursiveArrayTools.RaggedRange -@test Base.broadcastable(ragged_range_idx) isa Ref -# Broadcasting with RaggedRange should not error -@test identity.(ragged_range_idx) === ragged_range_idx +@test ragged_idx isa Int +@test ragged_idx == 4 # max inner array length # Broadcasting of heterogeneous arrays (issue #454) u = VA[[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] +# Now size is (2, 2), views use rectangular size +# Assignment into actual storage +u.u[2] .= [10.0, 11.0] @test u.u[2] == [10.0, 11.0] -# Test DiffEqArray with 2D inner arrays (matrices) +# Test DiffEqArray with 2D inner arrays (matrices) - uniform sizes t = 1:2 -recs_2d = [rand(2, 3), rand(2, 4)] +recs_2d = [rand(2, 3), rand(2, 3)] diffeq_2d = DiffEqArray(recs_2d, t) -@test diffeq_2d[:, 1] == recs_2d[1] -@test diffeq_2d[:, 2] == recs_2d[2] -@test diffeq_2d[:, 1:2] == DiffEqArray(recs_2d[1:2], t[1:2]) -@test diffeq_2d[:, 1:2].t == t[1:2] -@test diffeq_2d[:, 2:end] == DiffEqArray(recs_2d[2:end], t[2:end]) -@test diffeq_2d[:, 2:end].t == t[2:end] -@test diffeq_2d[:, (end - 1):end] == DiffEqArray(recs_2d[(end - 1):end], t[(end - 1):end]) -@test diffeq_2d[:, (end - 1):end].t == t[(end - 1):end] - -# Test DiffEqArray with 3D inner arrays (tensors) -recs_3d = [rand(2, 3, 4), rand(2, 3, 5)] +@test diffeq_2d[:, :, 1] == recs_2d[1] +@test diffeq_2d[:, :, 2] == recs_2d[2] + +# Test DiffEqArray with 3D inner arrays (tensors) - uniform sizes +recs_3d = [rand(2, 3, 4), rand(2, 3, 4)] diffeq_3d = DiffEqArray(recs_3d, t) @test diffeq_3d[:, :, :, 1] == recs_3d[1] @test diffeq_3d[:, :, :, 2] == recs_3d[2] -@test diffeq_3d[:, :, :, 1:2] == DiffEqArray(recs_3d[1:2], t[1:2]) -@test diffeq_3d[:, :, :, 1:2].t == t[1:2] -@test diffeq_3d[:, :, :, 2:end] == DiffEqArray(recs_3d[2:end], t[2:end]) -@test diffeq_3d[:, :, :, 2:end].t == t[2:end] -@test diffeq_3d[:, :, :, (end - 1):end] == DiffEqArray(recs_3d[(end - 1):end], t[(end - 1):end]) -@test diffeq_3d[:, :, :, (end - 1):end].t == t[(end - 1):end] # 2D inner arrays (matrices) with ragged second dimension u = VectorOfArray([zeros(1, n) for n in (2, 3)]) -@test length(view(u, 1, :, 1)) == 2 -@test length(view(u, 1, :, 2)) == 3 -u[1, :, 2] .= [1.0, 2.0, 3.0] +# size(u) = (1, 3, 2) since max columns = 3 +@test size(u) == (1, 3, 2) +# Direct storage manipulation for ragged inner arrays +u.u[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] +u.u[2] .= [7.0 2.0 9.0] @test u.u[2] == [7.0 2.0 9.0] -# test scalar indexing with end +# test scalar indexing with end (end in dim 2 = 3, end in dim 3 = 2) @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]) # 3D inner arrays (tensors) with ragged third dimension u = VectorOfArray([zeros(2, 1, n) for n in (2, 3)]) -@test size(view(u, :, :, :, 1)) == (2, 1, 2) -@test size(view(u, :, :, :, 2)) == (2, 1, 3) -# assign into a slice of the second inner array using last index Int -u[2, 1, :, 2] .= [7.0, 8.0, 9.0] +# size(u) = (2, 1, 3, 2) since max in dim 3 = 3 +@test size(u) == (2, 1, 3, 2) +# Direct manipulation of storage +u.u[2][2, 1, :] .= [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] +u.u[2][1, 1, 1] = 1.0 +u.u[2][2, 1, 1] = 2.0 +u.u[2][1, 1, 3] = 3.0 +u.u[2][2, 1, 3] = 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] == VectorOfArray(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] -# Test that views can be modified -f3 = VA[[1.0, 2.0], [3.0, 4.0, 5.0]] +# Test that views can be modified (non-ragged) +f3 = VA[[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]] v = view(f3, :, 2) @test length(v) == 3 v[1] = 10.0 diff --git a/test/downstream/adjoints.jl b/test/downstream/adjoints.jl new file mode 100644 index 00000000..f6672de9 --- /dev/null +++ b/test/downstream/adjoints.jl @@ -0,0 +1,13 @@ +using RecursiveArrayTools, Zygote, ForwardDiff, Test +using SciMLBase + +# Test that ArrayPartition works through ODEProblem construction +# (requires SciMLBase, so this is a downstream test) +function loss_odeproblem(x) + _x = ArrayPartition([x .* i for i in 1:5]...) + _prob = ODEProblem((u, p, t) -> u, _x, (0, 1)) + return sum(abs2, Array(_prob.u0)) +end + +x = float.(6:10) +@test Zygote.gradient(loss_odeproblem, x)[1] == ForwardDiff.gradient(loss_odeproblem, x) diff --git a/test/gpu/Project.toml b/test/gpu/Project.toml index a38a49e0..79c3e5ea 100644 --- a/test/gpu/Project.toml +++ b/test/gpu/Project.toml @@ -1,6 +1,7 @@ [deps] CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" +RecursiveArrayToolsArrayPartitionAnyAll = "172d604e-c495-4f00-97bf-d70957099afa" [compat] CUDA = "3.12, 4, 5" diff --git a/test/gpu/arraypartition_gpu.jl b/test/gpu/arraypartition_gpu.jl index 384ac30b..51e8f9f1 100644 --- a/test/gpu/arraypartition_gpu.jl +++ b/test/gpu/arraypartition_gpu.jl @@ -1,4 +1,4 @@ -using RecursiveArrayTools, ArrayInterface, CUDA, Adapt, Test +using RecursiveArrayTools, RecursiveArrayToolsArrayPartitionAnyAll, ArrayInterface, CUDA, Adapt, Test CUDA.allowscalar(false) # Test indexing with colon diff --git a/test/gpu/vectorofarray_gpu.jl b/test/gpu/vectorofarray_gpu.jl index bb5ee76a..b5bf351f 100644 --- a/test/gpu/vectorofarray_gpu.jl +++ b/test/gpu/vectorofarray_gpu.jl @@ -36,8 +36,8 @@ function f(p) end Zygote.gradient(f, p) -# Check conversion preserves device -va_cu = convert(AbstractArray, va) +# Check conversion to dense GPU array +va_cu = stack(va.u) @test va_cu isa CuArray @test size(va_cu) == size(x) diff --git a/test/interface_tests.jl b/test/interface_tests.jl index 77b75b4e..c8da9bc9 100644 --- a/test/interface_tests.jl +++ b/test/interface_tests.jl @@ -1,4 +1,5 @@ using RecursiveArrayTools, StaticArrays, Test +using RecursiveArrayToolsShorthandConstructors using FastBroadcast using SymbolicIndexingInterface: SymbolCache @@ -6,11 +7,13 @@ t = 1:3 testva = VA[[1, 2, 3], [4, 5, 6], [7, 8, 9]] testda = DiffEqArray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], t) -for (i, elem) in enumerate(testva) +# Iteration now goes over elements (AbstractArray behavior) +# To iterate over columns, use .u +for (i, elem) in enumerate(testva.u) @test elem == testva[:, i] end -for (i, elem) in enumerate(testda) +for (i, elem) in enumerate(testda.u) @test elem == testda[:, i] end @@ -158,8 +161,9 @@ testva2 = similar(testva) @test typeof(testva2) == typeof(testva) @test size(testva2) == size(testva) +# similar(VA, dims) returns a regular Array (AbstractArray behavior) testva3 = similar(testva, 10) -@test typeof(testva3) == typeof(testva) +@test testva3 isa Vector{Float64} @test length(testva3) == 10 # Fill AbstractVectorOfArray and check all @@ -225,11 +229,13 @@ emptyva = VectorOfArray(Array{Vector{Float64}}([])) emptyda = DiffEqArray(Array{Vector{Float64}}([]), Vector{Float64}()) @test isempty(emptyda) +# map now works element-wise (AbstractArray behavior) +# To map over inner arrays, use .u A = VectorOfArray(map(i -> rand(2, 4), 1:7)) -@test map(x -> maximum(x), A) isa Vector +@test map(x -> maximum(x), A.u) isa Vector DA = DiffEqArray(map(i -> rand(2, 4), 1:7), 1:7) -@test map(x -> maximum(x), DA) isa Vector +@test map(x -> maximum(x), DA.u) isa Vector u = VectorOfArray([fill(2, SVector{2, Float64}), ones(SVector{2, Float64})]) @test typeof(zero(u)) <: typeof(u) diff --git a/test/jet_tests.jl b/test/jet_tests.jl index 494b9fb9..f064ad2c 100644 --- a/test/jet_tests.jl +++ b/test/jet_tests.jl @@ -4,10 +4,14 @@ using JET, Test, RecursiveArrayTools result = JET.report_package(RecursiveArrayTools; target_modules = (RecursiveArrayTools,)) reports = JET.get_reports(result) -# Filter out similar_type inference errors from StaticArraysCore +# Filter out known false positives filtered_reports = filter(reports) do report s = string(report) - !(occursin("similar_type", s) && occursin("StaticArraysCore", s)) + # StaticArraysCore similar_type inference + occursin("similar_type", s) && occursin("StaticArraysCore", s) && return false + # RecipesBase user recipe keywords (denseplot, plotdensity, etc.) are dynamic + occursin("is_key_supported", s) && occursin("RecipesBase", s) && return false + return true end # Check if there are any non-filtered errors diff --git a/test/linalg.jl b/test/linalg.jl index 78d49581..494f1cbc 100644 --- a/test/linalg.jl +++ b/test/linalg.jl @@ -1,4 +1,5 @@ using RecursiveArrayTools, Test, Random +using RecursiveArrayToolsShorthandConstructors using LinearAlgebra, ArrayInterface n, m = 5, 6 diff --git a/test/named_array_partition_tests.jl b/test/named_array_partition_tests.jl index bec8cd5c..9a012635 100644 --- a/test/named_array_partition_tests.jl +++ b/test/named_array_partition_tests.jl @@ -7,7 +7,7 @@ using RecursiveArrayTools, ArrayInterface, Test @test typeof(similar(x)) <: NamedArrayPartition @test typeof(similar(x, Int)) <: NamedArrayPartition @test x.a ≈ ones(10) - @test typeof(x .+ x[1:end]) <: Vector # test broadcast precedence + @test typeof(x .+ x[1:end]) <: NamedArrayPartition # x[1:end] preserves type @test all(x .== x[1:end]) @test ArrayInterface.zeromatrix(x) isa Matrix @test size(ArrayInterface.zeromatrix(x)) == (30, 30) diff --git a/test/partitions_test.jl b/test/partitions_test.jl index 0d708f5f..44bf10c9 100644 --- a/test/partitions_test.jl +++ b/test/partitions_test.jl @@ -1,4 +1,5 @@ using RecursiveArrayTools, Test, Statistics, ArrayInterface, Adapt +using RecursiveArrayToolsShorthandConstructors @test length(ArrayPartition()) == 0 @test isempty(ArrayPartition()) @@ -177,7 +178,13 @@ recursivecopy!(dest_ap, src_ap) @inferred mapreduce(string, *, x) @test mapreduce(i -> string(i) * "q", *, x) == "1q2q3.0q4.0q" -# any +# any/all — optimized partition-level iteration requires RecursiveArrayToolsArrayPartitionAnyAll +# to avoid ~780 invalidations. Without the extension, Base's AbstractArray fallback is used. +# On GPU arrays, the fallback triggers scalar indexing errors — load the subpackage to fix. +@test which(any, Tuple{Function, ArrayPartition}).module === Base +@test which(all, Tuple{Function, ArrayPartition}).module === Base + +# Correctness tests (work via AbstractArray fallback on CPU) @test !any(isnan, AP[[1, 2], [3.0, 4.0]]) @test !any(isnan, AP[[3.0, 4.0]]) @test any(isnan, AP[[NaN], [3.0, 4.0]]) @@ -186,7 +193,6 @@ recursivecopy!(dest_ap, src_ap) @test any(isnan, AP[[2], [NaN]]) @test any(isnan, AP[[2], AP[[NaN]]]) -# all @test !all(isnan, AP[[1, 2], [3.0, 4.0]]) @test !all(isnan, AP[[3.0, 4.0]]) @test !all(isnan, AP[[NaN], [3.0, 4.0]]) @@ -381,7 +387,8 @@ for i in 1:length(part_a.x) @test typeof(sub_a) === typeof(sub_b) # Test type equality end -# ArrayPartition `all` with a functor +# ArrayPartition `all` with a functor (works via AbstractArray fallback, +# optimized partition-level iteration requires RecursiveArrayToolsArrayPartitionAnyAll) struct TestIsnanFunctor end (::TestIsnanFunctor)(x) = isnan(x) @test all(TestIsnanFunctor(), AP[[NaN], [NaN]]) diff --git a/test/runtests.jl b/test/runtests.jl index 683b8309..b486f64a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,4 +1,10 @@ using Pkg +# Install the ShorthandConstructors subpackage for tests that need VA[...]/AP[...] syntax +Pkg.develop( + PackageSpec( + path = joinpath(dirname(@__DIR__), "lib", "RecursiveArrayToolsShorthandConstructors") + ) +) using RecursiveArrayTools using Test using SafeTestsets @@ -14,6 +20,11 @@ end function activate_gpu_env() Pkg.activate("gpu") Pkg.develop(PackageSpec(path = dirname(@__DIR__))) + Pkg.develop( + PackageSpec( + path = joinpath(dirname(@__DIR__), "lib", "RecursiveArrayToolsArrayPartitionAnyAll") + ) + ) return Pkg.instantiate() end @@ -42,12 +53,36 @@ 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 + end + if GROUP == "Downstream" activate_downstream_env() @time @safetestset "ODE Solve Tests" include("downstream/odesolve.jl") @time @safetestset "Event Tests with ArrayPartition" include("downstream/downstream_events.jl") @time @safetestset "Measurements and Units" include("downstream/measurements_and_units.jl") @time @safetestset "TrackerExt" include("downstream/TrackerExt.jl") + # TODO: re-enable after SciMLBase compat bump for RAT v4 (SciML/SciMLBase.jl#1297) + # @time @safetestset "Downstream Adjoint Tests" include("downstream/adjoints.jl") end if GROUP == "SymbolicIndexingInterface" || GROUP == "Downstream" diff --git a/test/utils_test.jl b/test/utils_test.jl index 15807ea1..96446d00 100644 --- a/test/utils_test.jl +++ b/test/utils_test.jl @@ -1,4 +1,5 @@ using RecursiveArrayTools, StaticArrays +using RecursiveArrayToolsShorthandConstructors using Test t = collect(range(0, stop = 10, length = 200)) @@ -112,15 +113,15 @@ recursivefill!(x, true) x_voa = VectorOfArray(x) @test eltype(x_voa) === Vec3 - @test first(x_voa) isa AbstractVector{Vec3} + @test first(x_voa.u) isa AbstractVector{Vec3} y_voa = recursivecopy(x_voa) recursivefill!(y_voa, true) - @test all(y_voa[:, n] == fill(ones(Vec3), n) for n in 1:4) + @test all(y_voa.u[n] == fill(ones(Vec3), n) for n in 1:4) y_voa = recursivecopy(x_voa) recursivefill!(y_voa, ones(Vec3)) - @test all(y_voa[:, n] == fill(ones(Vec3), n) for n in 1:4) + @test all(y_voa.u[n] == fill(ones(Vec3), n) for n in 1:4) end @testset "VectorOfArray recursivecopy!" begin