diff --git a/src/vector_of_array.jl b/src/vector_of_array.jl index 81e7c9c2..42261fa3 100644 --- a/src/vector_of_array.jl +++ b/src/vector_of_array.jl @@ -562,6 +562,18 @@ 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)) +end Base.broadcastable(x::RaggedRange) = Ref(x) @inline function _is_ragged_dim(VA::AbstractVectorOfArray, d::Integer) @@ -579,6 +591,12 @@ Base.@propagate_inbounds function _getindex( return A.u[I] end +Base.@propagate_inbounds function _getindex( + A::AbstractDiffEqArray, ::NotSymbolic, ::Colon, I::Int + ) + return A.u[I] +end + Base.@propagate_inbounds function _getindex( A::AbstractVectorOfArray, ::NotSymbolic, I::Union{Int, AbstractArray{Int}, AbstractArray{Bool}, Colon}... @@ -589,6 +607,33 @@ Base.@propagate_inbounds function _getindex( stack(getindex.(A.u[last(I)], tuple.(Base.front(I))...)) end end + +Base.@propagate_inbounds function _getindex( + A::AbstractDiffEqArray, ::NotSymbolic, + I::Union{Int, AbstractArray{Int}, AbstractArray{Bool}, Colon}... + ) + return if last(I) isa Int + A.u[last(I)][Base.front(I)...] + else + col_idxs = last(I) + # Only preserve DiffEqArray type if all prefix indices are Colons (selecting whole inner arrays) + if all(idx -> idx isa Colon, Base.front(I)) + # For Colon, select all columns + if col_idxs isa Colon + col_idxs = eachindex(A.u) + end + # For DiffEqArray, we need to preserve the time values and type + # Create a vector of sliced arrays instead of stacking into higher-dim array + u_slice = [A.u[col][Base.front(I)...] for col in col_idxs] + # Return as DiffEqArray with sliced time values + return DiffEqArray(u_slice, A.t[col_idxs], parameter_values(A), symbolic_container(A)) + else + # Prefix indices are not all Colons - do the same as VectorOfArray + # (stack the results into a higher-dimensional array) + return stack(getindex.(A.u[col_idxs], tuple.(Base.front(I))...)) + end + end +end Base.@propagate_inbounds function _getindex( VA::AbstractVectorOfArray, ::NotSymbolic, ii::CartesianIndex ) @@ -674,6 +719,17 @@ end 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 @@ -757,27 +813,54 @@ end return (Base.front(args)..., resolved_last) end elseif args[end] isa RaggedRange - resolved_last = _resolve_ragged_index(args[end], A, 1) - if length(args) == 1 - return (resolved_last,) + # 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 - return (Base.front(args)..., resolved_last) + 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 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(raw_cols.start; step = raw_cols.step, stop = stop_val) + Base.range(start_val; step = raw_cols.step, stop = stop_val) else _column_indices(A, raw_cols) end @@ -800,37 +883,41 @@ end end return A.u[cols][padded...] else - return VectorOfArray( - [ - 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 + 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 - end - A.u[col][padded...] + ( + resolved_prefix..., + ( + lastindex( + A.u[col], + length(resolved_prefix) + i + ) for i in 1:n_missing + )..., + ) end - for col in cols - ] - ) + 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 @@ -864,7 +951,7 @@ end if col_idxs isa Int return A.u[col_idxs] else - return VectorOfArray(A.u[col_idxs]) + return _preserve_array_type(A, A.u[col_idxs], col_idxs) end end # If col_idxs resolved to a single Int, handle it directly diff --git a/test/basic_indexing.jl b/test/basic_indexing.jl index dc5e606b..df6b3914 100644 --- a/test/basic_indexing.jl +++ b/test/basic_indexing.jl @@ -77,7 +77,12 @@ diffeq = DiffEqArray(recs, t) @test diffeq[:, 1] == testa[:, 1] @test diffeq.u == recs @test diffeq[:, end] == testa[:, end] -@test diffeq[:, 2:end] == DiffEqArray([recs[i] for i in 2:length(recs)], t) +@test diffeq[:, 2:end] == DiffEqArray([recs[i] for i in 2:length(recs)], 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 diffeq[:, (end - 5):8] == DiffEqArray([recs[i] for i in (length(t) - 5):8], t[(length(t) - 5):8]) +@test diffeq[:, (end - 5):8].t == t[(length(t) - 5):8] # ## (Int, Int) @test testa[5, 4] == testva[5, 4] @@ -148,6 +153,12 @@ diffeq = DiffEqArray(recs, t) @test testva[1:2, 1:2] == [1 3; 2 5] @test diffeq[:, 1] == recs[1] @test diffeq[1:2, 1:2] == [1 3; 2 5] +@test diffeq[:, 1:2] == 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) f = VectorOfArray([[1.0], [2.0, 3.0]]) @@ -179,6 +190,7 @@ ragged = VectorOfArray([[1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0, 9.0]]) @test ragged[1:end, 3] == [6.0, 7.0, 8.0, 9.0] @test ragged[:, end] == [6.0, 7.0, 8.0, 9.0] @test ragged[:, 2:end] == VectorOfArray(ragged.u[2:end]) +@test ragged[:, (end - 1):end] == VectorOfArray(ragged.u[(end - 1):end]) ragged2 = VectorOfArray([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]]) @test ragged2[end, 1] == 4.0 @@ -199,6 +211,7 @@ ragged2 = VectorOfArray([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0], [7.0, 8.0, 9.0]]) @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) @@ -222,6 +235,31 @@ u = VectorOfArray([[1.0], [2.0, 3.0]]) u[:, 2] .= [10.0, 11.0] @test u.u[2] == [10.0, 11.0] +# Test DiffEqArray with 2D inner arrays (matrices) +t = 1:2 +recs_2d = [rand(2, 3), rand(2, 4)] +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)] +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