update indexing tests

SciML · Jul 30, 2024 · 310fa9e · 310fa9e
1 parent e7a8df5
commit 310fa9e
Show file tree

Hide file tree

Showing 3 changed files with 9 additions and 9 deletions.
diff --git a/test/downstream/comprehensive_indexing.jl b/test/downstream/comprehensive_indexing.jl
@@ -314,11 +314,11 @@ end
     @testset "Interpolation" begin
         sol = valp
         interpolated_sol = sol(0.0:1.0:10.0)
-        @test interpolated_sol[XY] isa Vector
+        @test interpolated_sol.u[XY] isa Vector
         @test interpolated_sol[XY, :] isa Vector
         @test interpolated_sol[XY, 2] isa Float64
         @test length(interpolated_sol[XY, 1:5]) == 5
-        @test interpolated_sol[XY] ≈ interpolated_sol[X] .+ interpolated_sol[Y]
+        @test interpolated_sol.u[XY] ≈ interpolated_sol[X] .+ interpolated_sol[Y]
         @test collect(interpolated_sol[t]) isa Vector
         @test collect(interpolated_sol[t, :]) isa Vector
         @test interpolated_sol[t, 2] isa Float64
@@ -891,7 +891,7 @@ end
         plotfn(t, u) = (t, 2u)
         all_idxs = [ud1, 2ud1, ud2, (plotfn, 0, ud1), (plotfn, t, ud1)]
         sym_idxs = [:ud1, :ud2, (plotfn, 0, :ud1), (plotfn, 0, :ud1)]
-        
+
         for idx in Iterators.flatten((all_idxs, sym_idxs))
             @test_nowarn plot(sol; idxs = idx)
             @test_nowarn plot(sol; idxs = [idx])

diff --git a/test/downstream/ensemble_multi_prob.jl b/test/downstream/ensemble_multi_prob.jl
@@ -17,8 +17,8 @@ prob3 = ODEProblem(sys3, [3.0, 3.0], (0.0, 1.0))
 ensemble_prob = EnsembleProblem([prob1, prob2, prob3])
 sol = solve(ensemble_prob, Tsit5(), EnsembleThreads())
 for i in 1:3
-    @test sol[x, :][i] == sol[i][x]
-    @test sol[y, :][i] == sol[i][y]
+    @test sol[x, :][i] == sol.u[i][x]
+    @test sol[y, :][i] == sol.u[i][y]
 end
 # Ensemble is a recursive array
 @test only.(sol(0.0, idxs = [x])) == sol[1, 1, :] == first.(sol[x, :])

diff --git a/test/downstream/ensemble_nondes.jl b/test/downstream/ensemble_nondes.jl
@@ -13,23 +13,23 @@ ensembleprob = Optimization.EnsembleProblem(
 
 sol = Optimization.solve(ensembleprob, OptimizationOptimJL.BFGS(),
     EnsembleThreads(), trajectories = 4, maxiters = 5)
-@test findmin(i -> sol[i].objective, 1:4)[1] < sol1.objective
+@test findmin(i -> sol.u[i].objective, 1:4)[1] < sol1.objective
 
 sol = Optimization.solve(ensembleprob, OptimizationOptimJL.BFGS(),
     EnsembleDistributed(), trajectories = 4, maxiters = 5)
-@test findmin(i -> sol[i].objective, 1:4)[1] < sol1.objective
+@test findmin(i -> sol.u[i].objective, 1:4)[1] < sol1.objective
 
 prob = OptimizationProblem(optf, x0, lb = [-0.5, -0.5], ub = [0.5, 0.5])
 ensembleprob = Optimization.EnsembleProblem(
     prob, prob_func = (prob, i, repeat) -> remake(prob, u0 = rand(-0.5:0.001:0.5, 2)))
 
 sol = Optimization.solve(ensembleprob, OptimizationOptimJL.BFGS(),
     EnsembleThreads(), trajectories = 5, maxiters = 5)
-@test findmin(i -> sol[i].objective, 1:4)[1] < sol1.objective
+@test findmin(i -> sol.u[i].objective, 1:4)[1] < sol1.objective
 
 sol = Optimization.solve(ensembleprob, OptimizationOptimJL.BFGS(),
     EnsembleDistributed(), trajectories = 5, maxiters = 5)
-@test findmin(i -> sol[i].objective, 1:4)[1] < sol1.objective
+@test findmin(i -> sol.u[i].objective, 1:4)[1] < sol1.objective
 
 using NonlinearSolve