test: update tests to account for new tearing

Author: AayushSabharwal

test/bvproblem.jl

@@ -30,7 +30,7 @@ daesolvers = [Ascher2, Ascher4, Ascher6]
     for solver in solvers
         sol = solve(bvp, solver(), dt = 0.01)
         @test isapprox(sol.u[end], osol.u[end]; atol = 0.01)
-        @test sol.u[1] == [1.0, 2.0]
+        @test sol[[x, y], 1] == [1.0, 2.0]
     end
 
     # Test out of place
@@ -40,7 +40,7 @@ daesolvers = [Ascher2, Ascher4, Ascher6]
     for solver in solvers
         sol = solve(bvp2, solver(), dt = 0.01)
         @test isapprox(sol.u[end], osol.u[end]; atol = 0.01)
-        @test sol.u[1] == [1.0, 2.0]
+        @test sol[[x, y], 1] == [1.0, 2.0]
     end
 end
 
@@ -129,7 +129,9 @@ end
     sol2 = solve(bvpi2, MIRK4(), dt = 0.01)
     sol3 = solve(bvpi3, MIRK4(), dt = 0.01)
     sol4 = solve(bvpi4, MIRK4(), dt = 0.01)
-    @test sol1 ≈ sol2 ≈ sol3 ≈ sol4 # don't get true equality here, not sure why
+    @test sol1.t ≈ sol2.t ≈ sol3.t ≈ sol4.t
+    @test sol1.u ≈ sol2.u ≈ sol3[[x(t), y(t)]] ≈ sol4[[x(t), y(t)]]
+    # @test sol1 ≈ sol2 ≈ sol3 ≈ sol4 # don't get true equality here, not sure why
 end
 
 function test_solvers(
@@ -296,7 +298,7 @@ end
     costfn = ModelingToolkit.generate_cost(
         lksys; expression = Val{false}, wrap_gfw = Val{true})
     _t = tspan[2]
-    @test costfn(sol, prob.p, _t) ≈ (sol(0.6)[1] + 3)^2 + sol(0.3)[1]^2
+    @test costfn(sol, prob.p, _t) ≈ (sol(0.6; idxs = x(t)) + 3)^2 + sol(0.3; idxs = x(t))^2
 
     ### With a parameter
     @parameters t_c
@@ -309,5 +311,6 @@ end
     sol = solve(prob, Tsit5())
     costfn = ModelingToolkit.generate_cost(
         lksys; expression = Val{false}, wrap_gfw = Val{true})
-    @test costfn(sol, prob.p, _t) ≈ log(sol(0.56)[2] + sol(0.0)[1]) - sol(0.4)[1]^2
+    @test costfn(sol, prob.p, _t) ≈
+          log(sol(0.56; idxs = y(t)) + sol(0.0; idxs = x(t))) - sol(0.4; idxs = x(t))^2
 end

test/dde.jl

@@ -46,13 +46,13 @@ prob = DDEProblem(sys,
     tspan,
     constant_lags = [tau])
 sol_mtk = solve(prob, alg, reltol = 1e-7, abstol = 1e-10)
-@test sol_mtk.u[end] ≈ sol.u[end]
+@test sol_mtk[[x₀, x₁, x₂(t)]][end] ≈ sol.u[end]
 prob2 = DDEProblem(sys,
     [x₀ => 1.0 - t * q1 * 10, x₁ => 1.0 - t * q1 * 10, x₂(t) => 1.0 - t * q1 * 10],
     tspan,
     constant_lags = [tau])
 sol2_mtk = solve(prob2, alg, reltol = 1e-7, abstol = 1e-10)
-@test sol2_mtk.u[end] ≈ sol2.u[end]
+@test sol2_mtk[[x₀, x₁, x₂(t)]][end] ≈ sol2.u[end]
 @test_nowarn sol2_mtk[[x₀, x₁, x₂(t)]]
 @test_nowarn sol2_mtk[[x₀, x₁, x₂(t - 0.1)]]
 

test/extensions/dynamic_optimization.jl

@@ -58,22 +58,22 @@ const M = ModelingToolkit
     jprob = JuMPDynamicOptProblem(lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     @test JuMP.num_constraints(jprob.model) == 2
     jsol = solve(jprob, Ipopt.Optimizer, constructTsitouras5, silent = true) # 12.190 s, 9.68 GiB
-    @test jsol.sol(0.6)[1] ≈ 3.5
-    @test jsol.sol(0.3)[1] ≈ 7.0
+    @test jsol.sol(0.6; idxs = x) ≈ 3.5
+    @test jsol.sol(0.3; idxs = x) ≈ 7.0
 
     cprob = CasADiDynamicOptProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     csol = solve(cprob, "ipopt", constructTsitouras5, silent = true)
-    @test csol.sol(0.6)[1] ≈ 3.5
-    @test csol.sol(0.3)[1] ≈ 7.0
+    @test csol.sol(0.6; idxs = x) ≈ 3.5
+    @test csol.sol(0.3; idxs = x) ≈ 7.0
 
     iprob = InfiniteOptDynamicOptProblem(
         lksys, u0map, tspan, parammap; guesses = guess, dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer,
         derivative_method = InfiniteOpt.OrthogonalCollocation(3), silent = true) # 48.564 ms, 9.58 MiB
     sol = isol.sol
-    @test sol(0.6)[1] ≈ 3.5
-    @test sol(0.3)[1] ≈ 7.0
+    @test sol(0.6; idxs = x) ≈ 3.5
+    @test sol(0.3; idxs = x) ≈ 7.0
 
     # Test whole-interval constraints
     constr = [x(t) ≳ 1, y(t) ≳ 1]
@@ -128,14 +128,14 @@ end
     # Linear systems have bang-bang controls
     @test is_bangbang(jsol.input_sol, [-1.0], [1.0])
     # Test reached final position.
-    @test ≈(jsol.sol.u[end][2], 0.25, rtol = 1e-5)
+    @test ≈(jsol.sol[x][end], 0.25, rtol = 1e-5)
 
     cprob = CasADiDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01)
     csol = solve(cprob, "ipopt", constructVerner8, silent = true)
     # Linear systems have bang-bang controls
     @test is_bangbang(csol.input_sol, [-1.0], [1.0])
     # Test reached final position.
-    @test ≈(csol.sol.u[end][2], 0.25, rtol = 1e-5)
+    @test ≈(csol.sol[x][end], 0.25, rtol = 1e-5)
 
     # Test dynamics
     @parameters (u_interp::ConstantInterpolation)(..)
@@ -149,7 +149,7 @@ end
     iprob = InfiniteOptDynamicOptProblem(block, u0map, tspan, parammap; dt = 0.01)
     isol = solve(iprob, Ipopt.Optimizer; silent = true)
     @test is_bangbang(isol.input_sol, [-1.0], [1.0])
-    @test ≈(isol.sol.u[end][2], 0.25, rtol = 1e-5)
+    @test ≈(isol.sol[x][end], 0.25, rtol = 1e-5)
     osol = solve(oprob, ImplicitEuler(); dt = 0.01, adaptive = false)
     @test ≈(isol.sol.u, osol.u, rtol = 0.05)
 
@@ -220,16 +220,16 @@ end
         Tₘ => 3.5 * g₀ * m₀, T(t) => 0.0, h₀ => 1, m_c => 0.6]
     jprob = JuMPDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
     jsol = solve(jprob, Ipopt.Optimizer, constructRadauIIA5, silent = true)
-    @test jsol.sol.u[end][1] > 1.012
+    @test jsol.sol[h][end] > 1.012
 
     cprob = CasADiDynamicOptProblem(rocket, u0map, (ts, te), pmap; dt = 0.001, cse = false)
     csol = solve(cprob, "ipopt"; silent = true)
-    @test csol.sol.u[end][1] > 1.012
+    @test csol.sol[h][end] > 1.012
 
     iprob = InfiniteOptDynamicOptProblem(
         rocket, u0map, (ts, te), pmap; dt = 0.001)
     isol = solve(iprob, Ipopt.Optimizer, silent = true)
-    @test isol.sol.u[end][1] > 1.012
+    @test isol.sol[h][end] > 1.012
 
     # Test solution
     @parameters (T_interp::CubicSpline)(..)

test/initial_values.jl

@@ -8,9 +8,10 @@ using SymbolicIndexingInterface
 @variables x(t)[1:3]=[1.0, 2.0, 3.0] y(t) z(t)[1:2]
 
 @mtkcompile sys=System([D(x) ~ t * x], t) simplify=false
-@test get_u0(sys, []) == [1.0, 2.0, 3.0]
-@test get_u0(sys, [x => [2.0, 3.0, 4.0]]) == [2.0, 3.0, 4.0]
-@test get_u0(sys, [x[1] => 2.0, x[2] => 3.0, x[3] => 4.0]) == [2.0, 3.0, 4.0]
+reorderer = getsym(sys, x)
+@test reorderer(get_u0(sys, [])) == [1.0, 2.0, 3.0]
+@test reorderer(get_u0(sys, [x => [2.0, 3.0, 4.0]])) == [2.0, 3.0, 4.0]
+@test reorderer(get_u0(sys, [x[1] => 2.0, x[2] => 3.0, x[3] => 4.0])) == [2.0, 3.0, 4.0]
 @test get_u0(sys, [2.0, 3.0, 4.0]) == [2.0, 3.0, 4.0]
 
 @mtkcompile sys=System([
@@ -182,7 +183,7 @@ end
     prob = ODEProblem(sys, [], (0.0, 1.0))
     sol = solve(prob)
     @test SciMLBase.successful_retcode(sol)
-    @test sol.u[1] ≈ [1.0, 1.0, 0.5, 0.5]
+    @test sol[x, 1] ≈ [1.0, 1.0, 0.5, 0.5]
 end
 
 @testset "Missing/cyclic guesses throws error" begin

test/initializationsystem.jl

@@ -459,7 +459,7 @@ sol = solve(prob, Tsit5())
 
 prob = ODEProblem(simpsys, [z => 1.0, y => 1.0], tspan, guesses = [x => 2.0])
 sol = solve(prob, Tsit5())
-@test sol.u[1] == [0.0, 1.0]
+@test sol[[x, y], 1] == [0.0, 1.0]
 
 # This should warn, but logging tests can't be marked as broken
 @test_logs prob = ODEProblem(simpsys, [], tspan, guesses = [x => 2.0])

test/input_output_handling.jl

@@ -369,7 +369,7 @@ eqs = [D(y₁) ~ -k₁ * y₁ + k₃ * y₂ * y₃ + u1
 m_inputs = [u[1], u[2]]
 m_outputs = [y₂]
 sys_simp = mtkcompile(sys, inputs = m_inputs, outputs = m_outputs)
-@test isequal(unknowns(sys_simp), collect(x[1:2]))
+@test issetequal(unknowns(sys_simp), collect(x[1:2]))
 @test length(inputs(sys_simp)) == 2
 
 # https://github.com/SciML/ModelingToolkit.jl/issues/1577
@@ -417,7 +417,7 @@ matrices, ssys = linearize(augmented_sys,
     ], outs;
     op = [augmented_sys.u => 0.0, augmented_sys.input.u[2] => 0.0, augmented_sys.d => 0.0])
 matrices = ModelingToolkit.reorder_unknowns(
-    matrices, unknowns(ssys), [ssys.x[2], ssys.integrator.x[1], ssys.x[1]])
+    matrices, unknowns(ssys), [ssys.x[1], ssys.x[2], ssys.integrator.x[1]])
 @test matrices.A ≈ [A [1; 0]; zeros(1, 2) -0.001]
 @test matrices.B == I
 @test matrices.C == [C zeros(2)]

test/odesystem.jl

@@ -568,10 +568,10 @@ let
     @named sys = System(eqs, t, vcat(x, [y]), [k], defaults = Dict(x .=> 0))
     sys = mtkcompile(sys)
 
-    u0 = unknowns(sys) .=> [0.5, 0]
-    du0 = D.(unknowns(sys)) .=> 0.0
+    u0 = x .=> [0.5, 0]
+    du0 = D.(x) .=> 0.0
     prob = DAEProblem(sys, [du0; u0], (0, 50))
-    @test prob.u0 ≈ [0.5, 0.0]
+    @test prob[x] ≈ [0.5, 0.0]
     @test prob.du0 ≈ [0.0, 0.0]
     @test prob.p isa MTKParameters
     @test prob.ps[k] ≈ 1
@@ -581,7 +581,8 @@ let
 
     prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0, x[1] => 0.5],
         (0, 50))
-    @test prob.u0 ≈ [0.5, 0]
+
+    @test prob[x] ≈ [0.5, 0]
     @test prob.du0 ≈ [0, 0]
     @test prob.p isa MTKParameters
     @test prob.ps[k] ≈ 1
@@ -590,7 +591,7 @@ let
 
     prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0, x[1] => 0.5, k => 2],
         (0, 50))
-    @test prob.u0 ≈ [0.5, 0]
+    @test prob[x] ≈ [0.5, 0]
     @test prob.du0 ≈ [0, 0]
     @test prob.p isa MTKParameters
     @test prob.ps[k] ≈ 2
@@ -769,7 +770,7 @@ let
     @named sys4 = System([D(x) ~ -y; D(y) ~ 1 + pp * y + x], t)
     sys4s = mtkcompile(sys4)
     prob = ODEProblem(sys4s, [x => 1.0, D(x) => 1.0], (0, 1.0))
-    @test string.(unknowns(prob.f.sys)) == ["x(t)", "y(t)"]
+    @test issetequal(string.(unknowns(prob.f.sys)), ["x(t)", "y(t)"])
     @test string.(parameters(prob.f.sys)) == ["pp"]
     @test string.(independent_variables(prob.f.sys)) == ["t"]
 end

test/reduction.jl

@@ -245,7 +245,7 @@ eqs = [D(x) ~ x
 sys = mtkcompile(model)
 Js = ModelingToolkit.jacobian_sparsity(sys)
 @test size(Js) == (3, 3)
-@test Js == Diagonal([1, 1, 0])
+@test Js == Diagonal([0, 1, 1])
 
 # MWE for #1722
 vars = @variables a(t) w(t) phi(t)

test/structural_transformation/tearing.jl

@@ -5,6 +5,7 @@ using ModelingToolkit.StructuralTransformations: SystemStructure, find_solvables
 using NonlinearSolve
 using LinearAlgebra
 using UnPack
+using SymbolicIndexingInterface
 using ModelingToolkit: t_nounits as t, D_nounits as D
 ###
 ### Nonlinear system
@@ -147,17 +148,21 @@ eqs = [D(x) ~ z * h
        0 ~ sin(z) + y - p * t]
 @named daesys = System(eqs, t)
 newdaesys = mtkcompile(daesys)
-@test equations(newdaesys) == [D(x) ~ h * z; 0 ~ y + sin(z) - p * t]
-@test equations(tearing_substitution(newdaesys)) == [D(x) ~ h * z; 0 ~ x + sin(z) - p * t]
-@test isequal(unknowns(newdaesys), [x, z])
+@test issetequal(equations(newdaesys), [D(x) ~ h * z; 0 ~ y + sin(z) - p * t])
+@test issetequal(
+    equations(tearing_substitution(newdaesys)), [D(x) ~ h * z; 0 ~ x + sin(z) - p * t])
+@test issetequal(unknowns(newdaesys), [x, z])
 prob = ODEProblem(newdaesys, [x => 1.0, z => -0.5π, p => 0.2], (0, 1.0))
 du = [0.0, 0.0];
 u = [1.0, -0.5π];
 pr = prob.p;
 tt = 0.1;
 @test (@ballocated $(prob.f)($du, $u, $pr, $tt)) == 0
 prob.f(du, u, pr, tt)
-@test du≈[u[2], u[1] + sin(u[2]) - prob.ps[p] * tt] atol=1e-5
+xgetter = getsym(prob, x)
+zgetter = getsym(prob, z)
+@test xgetter(du)≈zgetter(u) atol=1e-5
+@test zgetter(du)≈xgetter(u) + sin(zgetter(u)) - prob.ps[p] * tt atol=1e-5
 
 # test the initial guess is respected
 @named sys = System(eqs, t, defaults = Dict(z => NaN))