fix returning DiffEqArray

Author: JoshuaLampert

src/vector_of_array.jl

@@ -585,6 +585,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}...
@@ -595,6 +601,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
     )
@@ -680,6 +713,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
@@ -763,27 +807,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
@@ -806,37 +877,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
 
@@ -870,7 +945,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

test/basic_indexing.jl

@@ -77,7 +77,10 @@ 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)]
 
 # ## (Int, Int)
 @test testa[5, 4] == testva[5, 4]
@@ -148,6 +151,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]])
@@ -224,6 +233,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