test: fix SII test group

Author: AayushSabharwal

test/downstream/comprehensive_indexing.jl

@@ -29,24 +29,24 @@ begin
         sqrt(k1 + Y)
     ]
     jumps = [
-        ConstantRateJump(kp, [X ~ X + 1]),
-        ConstantRateJump(kd * X, [X ~ X - 1]),
-        ConstantRateJump(k1 * X, [X ~ X - 1, Y ~ Y + 1]),
-        ConstantRateJump(k2 * Y, [X ~ X + 1, Y ~ Y - 1]),
-        ConstantRateJump(1, [Y ~ Y + 1]),
-        ConstantRateJump(Y, [Y ~ Y - 1])
+        ConstantRateJump(kp, [X ~ Pre(X) + 1]),
+        ConstantRateJump(kd * X, [X ~ Pre(X) - 1]),
+        ConstantRateJump(k1 * X, [X ~ Pre(X) - 1, Y ~ Pre(Y) + 1]),
+        ConstantRateJump(k2 * Y, [X ~ Pre(X) + 1, Y ~ Pre(Y) - 1]),
+        ConstantRateJump(1, [Y ~ Pre(Y) + 1]),
+        ConstantRateJump(Y, [Y ~ Pre(Y) - 1])
     ]
-    observed = [XY ~ X + Y]
+    obs = [XY ~ X + Y]
     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(System(diff_eqs, t; observed, name = :osys))
+    osys = complete(System(diff_eqs, t; observed = obs, 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(System(alg_eqs; observed, name = :nsys))
+        diff_eqs, t, [X, Y], [kp, kd, k1, k2]; noise_eqs, observed = obs, name = :ssys))
+    jsys = complete(JumpSystem(jumps, t, [X, Y], [kp, kd, k1, k2]; observed = obs, name = :jsys))
+    nsys = complete(System(alg_eqs; observed = obs, name = :nsys))
     optsys = complete(OptimizationSystem(
-        loss, [X, Y], [kp, kd, k1, k2]; observed, name = :optsys))
+        loss, [X, Y], [kp, kd, k1, k2]; observed = obs, name = :optsys))
 end
 
 # Prepares problems, integrators, and solutions.
@@ -59,25 +59,23 @@ begin
     # Creates problems.
     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)
-    problems = [oprob, sprob, dprob, jprob, nprob, hcprob, ssprob, optprob]
-    systems = [osys, ssys, jsys, jsys, nsys, nsys, osys, optsys]
+    problems = [oprob, sprob, jprob, nprob, hcprob, ssprob, optprob]
+    systems = [osys, ssys, jsys, nsys, nsys, osys, optsys]
 
     # Creates an `EnsembleProblem` for each problem.
     eoprob = EnsembleProblem(oprob)
     esprob = EnsembleProblem(sprob)
-    edprob = EnsembleProblem(dprob)
     ejprob = EnsembleProblem(jprob)
     enprob = EnsembleProblem(nprob)
     essprob = EnsembleProblem(ssprob)
     eoptprob = EnsembleProblem(optprob)
-    eproblems = [eoprob, esprob, edprob, ejprob, enprob, essprob, optprob]
-    esystems = [osys, ssys, jsys, jsys, nsys, osys, optsys]
+    eproblems = [eoprob, esprob, ejprob, enprob, essprob, optprob]
+    esystems = [osys, ssys, jsys, nsys, osys, optsys]
 
     # Creates integrators.
     oint = init(oprob, Tsit5(); save_everystep = false)
@@ -912,8 +910,9 @@ end
 @testset "Continuous interpolation before discrete save" begin
     @variables x(t)
     @parameters c(t)
+    devt = ModelingToolkit.SymbolicDiscreteCallback(1.0, [c ~ Pre(c) + 1]; discrete_parameters = [c], iv = t)
     @mtkcompile sys = System(
-        D(x) ~ c * cos(x), t, [x], [c]; discrete_events = [1.0 => [c ~ Pre(c) + 1]])
+        D(x) ~ c * cos(x), t, [x], [c]; discrete_events = [devt])
     prob = ODEProblem(sys, [x => 0.0, c => 1.0], (0.0, 2pi))
     sol = solve(prob, Tsit5())
     @test_nowarn sol(-0.1; idxs = sys.x)
@@ -925,7 +924,7 @@ end
 @testset "DDEs" begin
     function oscillator(; name, k = 1.0, τ = 0.01)
         @parameters k=k τ=τ
-        @variables x(..)=0.1+t y(t)=0.1+t jcn(t)=0.0+t delx(t)
+        @variables x(..)=0.1+t y(t)=0.1+t jcn(t) delx(t)
         eqs = [D(x(t)) ~ y,
             D(y) ~ -k * x(t - τ) + jcn,
             delx ~ x(t - τ)]
@@ -956,7 +955,7 @@ end
     function oscillator(; name, k = 1.0, τ = 0.01)
         @parameters k=k τ=τ
         @brownian a
-        @variables x(..)=0.1+t y(t)=0.1+t jcn(t)=0.0+t delx(t)
+        @variables x(..)=0.1+t y(t)=0.1+t jcn(t) delx(t)
         eqs = [D(x(t)) ~ y + a,
             D(y) ~ -k * x(t - τ) + jcn,
             delx ~ x(t - τ)]
@@ -988,7 +987,7 @@ end
     end
     db = 1.0 => ModelingToolkit.ImperativeAffect(affect2!; modified = (; k))
 
-    @named ssys = System(D(A) ~ k * A, t, [A], [k]; noise_eqs = [0.0] 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, k => 1.0], (0.0, 4.0))
     sol = solve(prob, RI5())

test/downstream/modelingtoolkit_remake.jl

@@ -18,7 +18,7 @@ eqs = [D(x) ~ σ * (y - x),
     D(y) ~ x * (ρ - z) - y,
     D(z) ~ x * y - β * z]
 
-@named sys = System(eqs, t; parameter_dependencies = [q ~ 3β])
+@named sys = System([eqs; q ~ 3β], t)
 sys = complete(sys)
 u0 = [x => 1.0,
     y => 0.0,
@@ -65,20 +65,7 @@ push!(probs, JumpProblem(js, [u0; p], tspan; aggregator = Direct()))
 @named optsys = OptimizationSystem(sum(eq.lhs for eq in eqs), [x, y, z], [σ, ρ, β])
 optsys = complete(optsys)
 push!(syss, optsys)
-push!(probs, OptimizationProblem(optsys, u0, p))
-
-k = ShiftIndex(t)
-@mtkcompile discsys = System(
-    [x ~ x(k - 1) * ρ + y(k - 2), y ~ y(k - 1) * σ - z(k - 2), z ~ z(k - 1) * β + x(k - 2)],
-    t; defaults = [x => 1.0, y => 1.0, z => 1.0, x(k-1) => 0.0, y(k-1) => 0.0, z(k-1) => 0.0])
-# Roundabout method to avoid having to specify values for previous timestep
-discprob = DiscreteProblem(discsys, p, (0, 10))
-for (var, v) in u0
-    discprob[var] = v
-    discprob[var(k - 1)] = 0.0
-end
-push!(syss, discsys)
-push!(probs, discprob)
+push!(probs, OptimizationProblem(optsys, [u0; p]))
 
 @mtkcompile sys = System(
     [0 ~ x^3 * β + y^3 * ρ - σ, 0 ~ x^2 + 2x * y + y^2, 0 ~ z^2 - 4z + 4],
@@ -167,6 +154,94 @@ for (sys, prob) in zip(syss, probs)
     # Not testing `Symbol => expr` since nested substitution doesn't work with that
 end
 
+@testset "DiscreteProblem" begin
+    k = ShiftIndex(t)
+    @mtkcompile discsys = System(
+        [x ~ x(k - 1) * ρ + y(k - 2), y ~ y(k - 1) * σ - z(k - 2), z ~ z(k - 1) * β + x(k - 2)],
+        t; defaults = [x => 1.0, y => 1.0, z => 1.0, x(k-1) => 0.0, y(k-1) => 0.0, z(k-1) => 0.0])
+    prob = DiscreteProblem(discsys, p, (0, 10))
+    prob[x(k-1)] = 1.0
+    prob[y(k-1)] = prob[z(k-1)] = 0.0
+
+    @test parameter_values(prob) isa ModelingToolkit.MTKParameters
+    @inferred typeof(prob) remake(prob)
+
+    baseType = Base.typename(typeof(prob)).wrapper
+    ugetter = getsym(prob, [x(k-1), y(k-1), z(k-1)])
+    prob2 = @inferred baseType remake(prob; u0 = [x(k-1) => 2.0, y(k-1) => 3.0, z(k-1) => 4.0])
+    @test ugetter(prob2) == [2.0, 3.0, 4.0]
+    prob2 = @inferred baseType remake(prob; u0 = [sys.x(k-1) => 2.0, sys.y(k-1) => 3.0, sys.z(k-1) => 4.0])
+    @test ugetter(prob2) == [2.0, 3.0, 4.0]
+    prob2 = @inferred baseType remake(prob; u0 = [:xₜ₋₁ => 2.0, :yₜ₋₁ => 3.0, :zₜ₋₁ => 4.0])
+    @test ugetter(prob2) == [2.0, 3.0, 4.0]
+    prob2 = @inferred baseType remake(prob; u0 = [x(k-1) => 2.0, sys.y(k-1) => 3.0, :zₜ₋₁ => 4.0])
+    @test ugetter(prob2) == [2.0, 3.0, 4.0]
+
+    prob2 = @inferred baseType remake(prob; u0 = [x(k-1) => 12.0])
+    @test ugetter(prob2) == [12.0, 0.0, 0.0]
+    prob2 = @inferred baseType remake(prob; u0 = [sys.x(k-1) => 12.0])
+    @test ugetter(prob2) == [12.0, 0.0, 0.0]
+    prob2 = @inferred baseType remake(prob; u0 = [:xₜ₋₁ => 12.0])
+    @test ugetter(prob2) == [12.0, 0.0, 0.0]
+
+    pgetter = getp(prob, [σ, β, ρ])
+    prob2 = @inferred baseType remake(prob; p = [σ => 0.1, β => 0.2, ρ => 0.3])
+    @test pgetter(prob2) == [0.1, 0.2, 0.3]
+    if prob isa ODEProblem
+        @test prob2.ps[q] ≈ 0.6
+    end
+    prob2 = @inferred baseType remake(prob; p = [sys.σ => 0.1, sys.β => 0.2, sys.ρ => 0.3])
+    @test pgetter(prob2) == [0.1, 0.2, 0.3]
+    if prob isa ODEProblem
+        @test prob2.ps[q] ≈ 0.6
+    end
+    prob2 = @inferred baseType remake(prob; p = [:σ => 0.1, :β => 0.2, :ρ => 0.3])
+    @test pgetter(prob2) == [0.1, 0.2, 0.3]
+    if prob isa ODEProblem
+        @test prob2.ps[q] ≈ 0.6
+    end
+    prob2 = @inferred baseType remake(prob; p = [σ => 0.1, sys.β => 0.2, :ρ => 0.3])
+    @test pgetter(prob2) == [0.1, 0.2, 0.3]
+    if prob isa ODEProblem
+        @test prob2.ps[q] ≈ 0.6
+    end
+
+    prob2 = @inferred baseType remake(prob; p = [σ => 0.5])
+    @test pgetter(prob2) == [0.5, 8 / 3, 10.0]
+    prob2 = @inferred baseType remake(prob; p = [sys.σ => 0.5])
+    @test pgetter(prob2) == [0.5, 8 / 3, 10.0]
+    prob2 = @inferred baseType remake(prob; p = [:σ => 0.5])
+    @test pgetter(prob2) == [0.5, 8 / 3, 10.0]
+
+    # Test p dependent on u0
+    @test_broken begin
+        prob2 = @inferred baseType remake(prob; p = [σ => 0.5x(k-1) + 1])
+        @test pgetter(prob2) ≈ [1.5, 8 / 3, 10.0]
+        prob2 = @inferred baseType remake(prob; p = [sys.σ => 0.5x(k-1) + 1])
+        @test pgetter(prob2) ≈ [1.5, 8 / 3, 10.0]
+        prob2 = @inferred baseType remake(prob; p = [:σ => 0.5x(k-1) + 1])
+        @test pgetter(prob2) ≈ [1.5, 8 / 3, 10.0]
+    end
+
+    # Test u0 dependent on p
+    prob2 = @inferred baseType remake(prob; u0 = [x(k-1) => 0.5σ + 1])
+    @test ugetter(prob2) ≈ [15.0, 0.0, 0.0]
+    prob2 = @inferred baseType remake(prob; u0 = [sys.x(k-1) => 0.5σ + 1])
+    @test ugetter(prob2) ≈ [15.0, 0.0, 0.0]
+    prob2 = @inferred baseType remake(prob; u0 = [:xₜ₋₁ => 0.5σ + 1])
+    @test ugetter(prob2) ≈ [15.0, 0.0, 0.0]
+
+    # Test u0 dependent on p and p dependent on u0
+    prob2 = @inferred baseType remake(prob; u0 = [x(k-1) => 0.5σ + 1], p = [β => 0.5x(k-1) + 1])
+    @test ugetter(prob2) ≈ [15.0, 0.0, 0.0]
+    @test_broken pgetter(prob2) ≈ [28.0, 8.5, 10.0]
+    prob2 = @inferred baseType remake(
+        prob; u0 = [sys.x(k-1) => 0.5σ + 1], p = [sys.β => 0.5x(k-1) + 1])
+    @test ugetter(prob2) ≈ [15.0, 0.0, 0.0]
+    @test_broken pgetter(prob2) ≈ [28.0, 8.5, 10.0]
+    # Not testing `Symbol => expr` since nested substitution doesn't work with that
+end
+
 # Optimization
 @parameters p
 @mtkcompile sys = System([D(x) ~ -p * x], t)

test/downstream/problem_interface.jl

@@ -214,12 +214,12 @@ eqs = [
     D(X) ~ p - d * X,
     X2 ~ 2 * X
 ]
-@mtkpile osys = System(eqs, t)
+@mtkcompile osys = System(eqs, t)
 
 u0 = [X => 0.1]
 tspan = (0.0, 10.0)
 ps = [p => 1.0, d => 0.2]
-oprob = ODEProblem(osys, u0, tspan, ps)
+oprob = ODEProblem(osys, [u0; ps], tspan)
 eprob = EnsembleProblem(oprob)
 
 @test eprob[X] == 0.1
@@ -345,7 +345,7 @@ end
         b[1] = p[3]
         b[2] = -8p[2] - p[1]
     end
-    f = SciMLBase.SymbolicLinearInterface(update_A!, update_b!, sys, nothing)
+    f = SciMLBase.SymbolicLinearInterface(update_A!, update_b!, sys, nothing, nothing)
     A = Float64[1 1 1; 6 -4 5; 5 2 2]
     b = Float64[2, 31, 13]
     p = Float64[1, -4, 2]