diff --git a/test/downstream/ensemble_multi_prob.jl b/test/downstream/ensemble_multi_prob.jl index bd45a26b4..d79aeaac8 100644 --- a/test/downstream/ensemble_multi_prob.jl +++ b/test/downstream/ensemble_multi_prob.jl @@ -1,4 +1,5 @@ using ModelingToolkit, OrdinaryDiffEq, Test +using SymbolicIndexingInterface using ModelingToolkit: t_nounits as t, D_nounits as D @variables x(t), y(t) @@ -16,10 +17,12 @@ prob3 = ODEProblem(sys3, [3.0, 3.0], (0.0, 1.0)) # test that when passing a vector of problems, trajectories and the prob_func are chosen appropriately ensemble_prob = EnsembleProblem([prob1, prob2, prob3]) sol = solve(ensemble_prob, Tsit5(), EnsembleThreads()) +xidx = variable_index(sys1, x) +yidx = variable_index(sys1, y) for i in 1:3 - @test sol[1, :, i] == sol.u[i][x] - @test sol[2, :, i] == sol.u[i][y] + @test sol[xidx, :, i] == sol.u[i][x] + @test sol[yidx, :, i] == sol.u[i][y] end # Ensemble is a recursive array -@test only.(sol(0.0, idxs = [x])) == sol[1, 1, :] -@test only.(sol(1.0, idxs = [x])) ≈ [sol[i][1, end] for i in 1:3] +@test only.(sol(0.0, idxs = [x])) == sol[xidx, 1, :] +@test only.(sol(1.0, idxs = [x])) ≈ [sol[i][xidx, end] for i in 1:3] diff --git a/test/downstream/solution_interface.jl b/test/downstream/solution_interface.jl index 887bb44fa..c31957db9 100644 --- a/test/downstream/solution_interface.jl +++ b/test/downstream/solution_interface.jl @@ -103,7 +103,7 @@ prob = ODEProblem(sys, [u0; p], tspan) sol = solve(prob, Rodas4()) @test_throws ArgumentError sol[x] -@test in(sol[lorenz1.x], [getindex.(sol.u, 1) for i in 1:length(unknowns(sol.prob.f.sys))]) +@test in(sol[lorenz1.x], [getindex.(sol.u, i) for i in 1:length(unknowns(sol.prob.f.sys))]) @test_throws KeyError sol[:x] ### Non-symbolic indexing tests @@ -331,13 +331,14 @@ end @test _idxs == [1] @test _ss isa SciMLBase.SavedSubsystem _idxs, _ss = SciMLBase.get_save_idxs_and_saved_subsystem(prob, x) - @test _idxs == 1 + xidx = variable_index(prob, x) + @test _idxs == xidx @test _ss isa SciMLBase.SavedSubsystem _idxs, _ss = SciMLBase.get_save_idxs_and_saved_subsystem(prob, [x]) - @test _idxs == [1] + @test _idxs == [xidx] @test _ss isa SciMLBase.SavedSubsystem _idxs, _ss = SciMLBase.get_save_idxs_and_saved_subsystem(prob, [x, q]) - @test _idxs == [1] + @test _idxs == [xidx] @test _ss isa SciMLBase.SavedSubsystem _idxs, _ss = SciMLBase.get_save_idxs_and_saved_subsystem(prob, [q]) @test _idxs == Int[]