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/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))