Merge pull request #1045 from AayushSabharwal/as/fix-mtk-downstream

test: update tests to MTKv10

Author: ChrisRackauckas

docs/Project.toml

@@ -8,7 +8,7 @@ StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [compat]
 Documenter = "1"
-ModelingToolkit = "8.35, 9"
+ModelingToolkit = "10"
 OrdinaryDiffEq = "6"
 Plots = "1"
 SciMLBase = "1.74, 2"

test/downstream/Project.toml

@@ -36,7 +36,7 @@ DelayDiffEq = "5"
 DiffEqCallbacks = "3, 4"
 ForwardDiff = "0.10"
 JumpProcesses = "9.10"
-ModelingToolkit = "9.64.3"
+ModelingToolkit = "10.0"
 ModelingToolkitStandardLibrary = "2.7"
 NonlinearSolve = "2, 3, 4"
 Optimization = "4"

test/downstream/adjoints.jl

@@ -8,14 +8,14 @@ eqs = [D(x) ~ σ * (y - x),
     D(y) ~ x * (ρ - z) - y,
     D(z) ~ x * y - β * z]
 
-@named lorenz1 = ODESystem(eqs, t)
-@named lorenz2 = ODESystem(eqs, t)
+@named lorenz1 = System(eqs, t)
+@named lorenz2 = System(eqs, t)
 
 @parameters γ
 @variables a(t) α(t)
 connections = [0 ~ lorenz1.x + lorenz2.y + a * γ,
     α ~ 2lorenz1.x + a * γ]
-@mtkbuild sys = ODESystem(connections, t, [a, α], [γ], systems = [lorenz1, lorenz2])
+@mtkcompile sys = System(connections, t, [a, α], [γ], systems = [lorenz1, lorenz2])
 
 u0 = [lorenz1.x => 1.0,
     lorenz1.y => 0.0,
@@ -33,7 +33,7 @@ p = [lorenz1.σ => 10.0,
     γ => 2.0]
 
 tspan = (0.0, 100.0)
-prob = ODEProblem(sys, u0, tspan, p)
+prob = ODEProblem(sys, [u0; p], tspan)
 sol = solve(prob, Rodas4())
 
 gs_sym, = Zygote.gradient(sol) do sol
@@ -71,10 +71,10 @@ end
 
 # BatchedInterface AD
 @variables x(t)=1.0 y(t)=1.0 z(t)=1.0 w(t)=1.0
-@named sys1 = ODESystem([D(x) ~ x + y, D(y) ~ y * z, D(z) ~ z * t * x], t)
+@named sys1 = System([D(x) ~ x + y, D(y) ~ y * z, D(z) ~ z * t * x], t)
 sys1 = complete(sys1)
 prob1 = ODEProblem(sys1, [], (0.0, 10.0))
-@named sys2 = ODESystem([D(x) ~ x + w, D(y) ~ w * t, D(w) ~ x + y + w], t)
+@named sys2 = System([D(x) ~ x + w, D(y) ~ w * t, D(w) ~ x + y + w], t)
 sys2 = complete(sys2)
 prob2 = ODEProblem(sys2, [], (0.0, 10.0))
 

test/downstream/comprehensive_indexing.jl

@@ -40,11 +40,11 @@ begin
     loss = kd * (k1 - X)^2 + k2 * (kp * Y - X^2)^2
 
     # Create systems (without structural_simplify, since that might modify systems to affect intended tests).
-    osys = complete(ODESystem(diff_eqs, t; observed, name = :osys))
-    ssys = complete(SDESystem(
-        diff_eqs, noise_eqs, t, [X, Y], [kp, kd, k1, k2]; observed, name = :ssys))
+    osys = complete(System(diff_eqs, t; observed, name = :osys))
+    ssys = complete(System(
+        diff_eqs, t, [X, Y], [kp, kd, k1, k2]; noise_eqs, observed, name = :ssys))
     jsys = complete(JumpSystem(jumps, t, [X, Y], [kp, kd, k1, k2]; observed, name = :jsys))
-    nsys = complete(NonlinearSystem(alg_eqs; observed, name = :nsys))
+    nsys = complete(System(alg_eqs; observed, name = :nsys))
     optsys = complete(OptimizationSystem(
         loss, [X, Y], [kp, kd, k1, k2]; observed, name = :optsys))
 end
@@ -57,14 +57,14 @@ begin
     p_vals = [kp => 1.0, kd => 0.1, k1 => 0.25, k2 => 0.5]
 
     # Creates problems.
-    oprob = ODEProblem(osys, u0_vals, tspan, p_vals)
-    sprob = SDEProblem(ssys, u0_vals, tspan, p_vals)
-    dprob = DiscreteProblem(jsys, u0_vals, tspan, p_vals)
-    jprob = JumpProblem(jsys, deepcopy(dprob), Direct(); rng)
-    nprob = NonlinearProblem(nsys, u0_vals, p_vals)
+    oprob = ODEProblem(osys, [u0_vals; p_vals], tspan)
+    sprob = SDEProblem(ssys, [u0_vals; p_vals], tspan)
+    dprob = DiscreteProblem(jsys, )
+    jprob = JumpProblem(jsys, [u0_vals; p_vals], tspan; aggregator = Direct(), rng)
+    nprob = NonlinearProblem(nsys, [u0_vals; p_vals])
     hcprob = NonlinearProblem(HomotopyNonlinearFunction(nprob.f), nprob.u0, nprob.p)
-    ssprob = SteadyStateProblem(osys, u0_vals, p_vals)
-    optprob = OptimizationProblem(optsys, u0_vals, p_vals, grad = true, hess = true)
+    ssprob = SteadyStateProblem(osys, [u0_vals; p_vals])
+    optprob = OptimizationProblem(optsys, [u0_vals; p_vals], grad = true, hess = true)
     problems = [oprob, sprob, dprob, jprob, nprob, hcprob, ssprob, optprob]
     systems = [osys, ssys, jsys, jsys, nsys, nsys, osys, optsys]
 
@@ -358,7 +358,7 @@ end
     ps = @parameters p[1:3] = [1, 2, 3]
     eqs = [collect(D.(x) .~ x)
            D(y) ~ norm(x) * y - x[1]]
-    @named sys = ODESystem(eqs, t, [sts...;], ps)
+    @named sys = System(eqs, t, [sts...;], ps)
     sys = complete(sys)
     prob = ODEProblem(sys, [], (0, 1.0))
     sol = solve(prob, Tsit5())
@@ -510,9 +510,9 @@ end
 # Issue https://github.com/SciML/ModelingToolkit.jl/issues/2697
 @testset "Interpolation of derivative of observed variables" begin
     @variables x(t) y(t) z(t) w(t)[1:2]
-    @named sys = ODESystem(
+    @named sys = System(
         [D(x) ~ 1, y ~ x^2, z ~ 2y^2 + 3x, w[1] ~ x + y + z, w[2] ~ z * x * y], t)
-    sys = structural_simplify(sys)
+    sys = mtkcompile(sys)
     prob = ODEProblem(sys, [x => 0.0], (0.0, 1.0))
     sol = solve(prob, Tsit5())
     @test_throws ErrorException sol(1.0, Val{1}, idxs = y)
@@ -912,9 +912,9 @@ end
 @testset "Continuous interpolation before discrete save" begin
     @variables x(t)
     @parameters c(t)
-    @mtkbuild sys = ODESystem(
-        D(x) ~ c * cos(x), t, [x], [c]; discrete_events = [1.0 => [c ~ c + 1]])
-    prob = ODEProblem(sys, [x => 0.0], (0.0, 2pi), [c => 1.0])
+    @mtkcompile sys = System(
+        D(x) ~ c * cos(x), t, [x], [c]; discrete_events = [1.0 => [c ~ Pre(c) + 1]])
+    prob = ODEProblem(sys, [x => 0.0, c => 1.0], (0.0, 2pi))
     sol = solve(prob, Tsit5())
     @test_nowarn sol(-0.1; idxs = sys.x)
     @test_nowarn sol(-0.1; idxs = [sys.x, 2sys.x])
@@ -939,7 +939,7 @@ end
         osc2.jcn ~ osc1.delx]
     @named coupledOsc = System(eqs, t)
     @named coupledOsc = compose(coupledOsc, systems)
-    sys = structural_simplify(coupledOsc)
+    sys = mtkcompile(coupledOsc)
     prob = DDEProblem(sys, [], (0.0, 10.0); constant_lags = [sys.osc1.τ, sys.osc2.τ])
     sym = sys.osc1.delx
     delay = sys.osc1.τ
@@ -970,7 +970,7 @@ end
         osc2.jcn ~ osc1.delx]
     @named coupledOsc = System(eqs, t)
     @named coupledOsc = compose(coupledOsc, systems)
-    sys = structural_simplify(coupledOsc)
+    sys = mtkcompile(coupledOsc)
     prob = SDDEProblem(sys, [], (0.0, 10.0); constant_lags = [sys.osc1.τ, sys.osc2.τ])
     sym = sys.osc1.delx
     delay = sys.osc1.τ
@@ -983,14 +983,14 @@ end
 @testset "RODESolutions save discretes" begin
     @parameters k(t)
     @variables A(t)
-    function affect2!(integ, u, p, ctx)
-        integ.ps[p.k] += 1.0
+    function affect2!(m, o, ctx, integ)
+        return (; k = m.k + 1.0)
     end
-    db = 1.0 => (affect2!, [], [k], [k], nothing)
+    db = 1.0 => ModelingToolkit.ImperativeAffect(affect2!; modified = (; k))
 
-    @named ssys = SDESystem(D(A) ~ k * A, [0.0], t, [A], [k], discrete_events = db)
+    @named ssys = System(D(A) ~ k * A, t, [A], [k]; noise_eqs = [0.0] discrete_events = db)
     ssys = complete(ssys)
-    prob = SDEProblem(ssys, [A => 1.0], (0.0, 4.0), [k => 1.0])
+    prob = SDEProblem(ssys, [A => 1.0, k => 1.0], (0.0, 4.0))
     sol = solve(prob, RI5())
     @test sol[k] isa AbstractVector
     @test sol[k] == [1.0, 2.0, 3.0, 4.0]

test/downstream/ensemble_multi_prob.jl

@@ -2,11 +2,11 @@ using ModelingToolkit, OrdinaryDiffEq, Test
 using ModelingToolkit: t_nounits as t, D_nounits as D
 @variables x(t), y(t)
 
-@mtkbuild sys1 = ODESystem([D(x) ~ x,
+@mtkcompile sys1 = System([D(x) ~ x,
         D(y) ~ -y], t)
-@mtkbuild sys2 = ODESystem([D(x) ~ 2x,
+@mtkcompile sys2 = System([D(x) ~ 2x,
         D(y) ~ -2y], t)
-@mtkbuild sys3 = ODESystem([D(x) ~ 3x,
+@mtkcompile sys3 = System([D(x) ~ 3x,
         D(y) ~ -3y], t)
 
 prob1 = ODEProblem(sys1, [1.0, 1.0], (0.0, 1.0))

test/downstream/initialization.jl

@@ -44,10 +44,10 @@ end
 
     @variables x(t) [guess = 1.0] y(t) [guess = 1.0]
     @parameters p=missing [guess=1.0] q=missing [guess=1.0]
-    @mtkbuild sys = ODESystem([D(x) ~ p * y + q * t, D(y) ~ 5x + q], t;
+    @mtkcompile sys = System([D(x) ~ p * y + q * t, D(y) ~ 5x + q], t;
         initialization_eqs = [p^2 + q^2 ~ 3, x^3 + y^3 ~ 5])
     prob = ODEProblem(
-        sys, [x => 1.0], (0.0, 1.0), [p => 1.0]; initializealg = SciMLBase.NoInit())
+        sys, [x => 1.0, p => 1.0], (0.0, 1.0); initializealg = SciMLBase.NoInit())
 
     @test prob.f.initialization_data isa SciMLBase.OverrideInitData
     integ = init(prob, Tsit5())
@@ -68,15 +68,15 @@ end
     eqs = [D(D(x)) ~ λ * x
            D(D(y)) ~ λ * y - g
            x^2 + y^2 ~ 1]
-    @mtkbuild pend = ODESystem(eqs, t)
+    @mtkcompile pend = System(eqs, t)
 
-    prob = ODEProblem(pend, [x => 1, y => 0], (0.0, 1.5), [g => 1], guesses = [λ => 1])
+    prob = ODEProblem(pend, [x => 1, y => 0, g => 1], (0.0, 1.5), guesses = [λ => 1])
     @test occursin("Initialization status: FULLY_DETERMINED", sprint(summary, prob))
 
-    prob = ODEProblem(pend, [], (0.0, 1.5), [g => 1], guesses = [λ => 1, x => 1, y => 0])
+    prob = ODEProblem(pend, [g => 1], (0.0, 1.5), guesses = [λ => 1, x => 1, y => 0])
     @test occursin("Initialization status: UNDERDETERMINED", sprint(summary, prob))
 
-    prob = ODEProblem(pend, [x => 1, y => 0, λ => 2], (0.0, 1.5), [g => 1])
+    prob = ODEProblem(pend, [x => 1, y => 0, λ => 2, g => 1], (0.0, 1.5))
     @test occursin("Initialization status: OVERDETERMINED", sprint(summary, prob))
 end
 

test/downstream/integrator_indexing.jl

@@ -10,13 +10,13 @@ using SciMLStructures: canonicalize, Tunable
 eqs = [D(s1) ~ a * s1 / (1 + s1 + s2) - b * s1,
     D(s2) ~ +c * s2 / (1 + s1 + s2) - d * s2]
 
-@named pop_model = ODESystem(eqs, t)
+@named pop_model = System(eqs, t)
 population_model = complete(pop_model)
 # Tests on ODEProblem.
 u0 = [s1 => 2.0, s2 => 1.0]
 p = [a => 2.0, b => 1.0, c => 1.0, d => 1.0]
 tspan = (0.0, 1000000.0)
-oprob = ODEProblem(population_model, u0, tspan, p)
+oprob = ODEProblem(population_model, [u0; p], tspan)
 integrator = init(oprob, Rodas4())
 
 @test_throws Exception integrator[a]
@@ -84,7 +84,7 @@ noiseeqs = [0.1 * s1,
     0.1 * s2]
 @named noisy_population_model = SDESystem(pop_model, noiseeqs)
 noisy_population_model = complete(noisy_population_model)
-sprob = SDEProblem(noisy_population_model, u0, (0.0, 100.0), p)
+sprob = SDEProblem(noisy_population_model, [u0; p], (0.0, 100.0))
 integrator = init(sprob, ImplicitEM())
 
 step!(integrator, 100.0, true)
@@ -127,14 +127,14 @@ eqs = [D(x) ~ σ * (y - x),
     D(y) ~ x * (ρ - z) - y,
     D(z) ~ x * y - β * z]
 
-@named lorenz1 = ODESystem(eqs, t)
-@named lorenz2 = ODESystem(eqs, t)
+@named lorenz1 = System(eqs, t)
+@named lorenz2 = System(eqs, t)
 
 @parameters γ
 @variables a(t) α(t)
 connections = [0 ~ lorenz1.x + lorenz2.y + a * γ,
     α ~ 2lorenz1.x + a * γ]
-@mtkbuild sys = ODESystem(
+@mtkbuild sys = System(
     connections, t, [a, α], [γ], systems = [lorenz1, lorenz2])
 
 u0 = [lorenz1.x => 1.0,
@@ -154,7 +154,7 @@ p = [lorenz1.σ => 10.0,
     γ => 2.0]
 
 tspan = (0.0, 100.0)
-prob = ODEProblem(sys, u0, tspan, p)
+prob = ODEProblem(sys, [u0; p], tspan)
 integrator = init(prob, Rodas4())
 sol = solve(prob, Rodas4())
 step!(integrator, 100.0, true)
@@ -185,8 +185,8 @@ step!(integrator, 100.0, true)
 @variables q(t)[1:2] = [1.0, 2.0]
 eqs = [D(q[1]) ~ 2q[1]
        D(q[2]) ~ 2.0]
-@named sys2 = ODESystem(eqs, t, [q...], [])
-sys2_simplified = structural_simplify(sys2)
+@named sys2 = System(eqs, t, [q...], [])
+sys2_simplified = mtkcompile(sys2)
 prob2 = ODEProblem(sys2_simplified, [], (0.0, 5.0))
 integrator2 = init(prob2, Tsit5())
 
@@ -199,7 +199,7 @@ integrator2 = init(prob2, Tsit5())
     @variables u(t)
     eqs = [D(u) ~ u]
 
-    @mtkbuild sys2 = ODESystem(eqs, t)
+    @mtkcompile sys2 = System(eqs, t)
 
     tspan = (0.0, 5.0)
 
@@ -324,7 +324,7 @@ sts = @variables x(t)[1:3]=[1, 2, 3.0] y(t)=1.0
 ps = @parameters p[1:3] = [1, 2, 3]
 eqs = [collect(D.(x) .~ x)
        D(y) ~ norm(x) * y - x[1]]
-@mtkbuild sys = ODESystem(eqs, t, [sts...;], [ps...;])
+@mtkcompile sys = System(eqs, t, [sts...;], [ps...;])
 prob = ODEProblem(sys, [], (0, 1.0))
 integrator = init(prob, Tsit5(), save_everystep = false)
 @test integrator[x] isa Vector{Float64}
@@ -365,5 +365,5 @@ integrator.ps[p] = [7, 8, 9]
 @variables X(t)
 
 eq = D(X) ~ p - X
-@mtkbuild osys = ODESystem([eq], t)
-oprob = ODEProblem(osys, [X => 0.1], (0.0, 1.0), [p => 1.0])
+@mtkcompile osys = System([eq], t)
+oprob = ODEProblem(osys, [X => 0.1, p => 1.0], (0.0, 1.0))