fix: fix interpolation of array symbolics

SciML · May 6, 2024 · 289d9e2 · 289d9e2
1 parent a0fab7a
commit 289d9e2
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 1 deletion.
diff --git a/src/solutions/ode_solutions.jl b/src/solutions/ode_solutions.jl
@@ -206,8 +206,10 @@ end
 
 function (sol::AbstractODESolution)(t::Number, ::Type{deriv}, idxs::AbstractVector,
         continuity) where {deriv}
-    all(!isequal(NotSymbolic()), symbolic_type.(idxs)) ||
+    if symbolic_type(idxs) == NotSymbolic() &&
+       any(isequal(NotSymbolic()), symbolic_type.(idxs))
         error("Incorrect specification of `idxs`")
+    end
     interp_sol = augment(sol.interp([t], nothing, deriv, sol.prob.p, continuity), sol)
     first(interp_sol[idxs])
 end

diff --git a/test/downstream/symbol_indexing.jl b/test/downstream/symbol_indexing.jl
@@ -226,6 +226,8 @@ sol = solve(prob, Tsit5())
 @test sol[x] isa Vector{<:Vector}
 @test sol[@nonamespace sys.x] isa Vector{<:Vector}
 @test sol.ps[p] == [1, 2, 3]
+# interpolation of array variables
+@test sol(1.0, idxs = x) == [sol(1.0, idxs = x[i]) for i in 1:3]
 
 x_idx = variable_index.((sys,), [x[1], x[2], x[3]])
 y_idx = variable_index(sys, y)