Merge pull request #1048 from AayushSabharwal/as/fix-downstream

test: make ensemble multi prob test independent of variable order

Author: ChrisRackauckas

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]

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[]