Merge pull request #1055 from DhairyaLGandhi/dg/v10

Update observables + AD tests with MTK v10

Author: ChrisRackauckas

.github/workflows/Tests.yml

@@ -32,6 +32,7 @@ jobs:
         group:
           - "Core"
           - "Downstream"
+          - "SymbolicIndexingInterface"
           - "QA"
           - "Python"
     uses: "SciML/.github/.github/workflows/tests.yml@v1"

Project.toml

@@ -93,7 +93,7 @@ StaticArraysCore = "1.4"
 Statistics = "1.10"
 SymbolicIndexingInterface = "0.3.36"
 Tables = "1.11"
-Zygote = "0.6.67, 0.7"
+Zygote = "0.7.10"
 julia = "1.10"
 
 [extras]

test/downstream/Project.toml

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

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, t; parameter_dependencies = [q ~ 3β])
 sys = complete(sys)
 u0 = [x => 1.0,
     y => 0.0,
@@ -70,7 +70,7 @@ 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])
+    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
@@ -338,7 +338,7 @@ end
 @testset "Lazy initialization" begin
     @variables _x(..) [guess = 1.0] y(t) [guess = 1.0]
     @parameters p=1.0 [guess = 1.0]
-    @brownian a
+    @brownians a
     x = _x(t)
 
     initprob = NonlinearProblem(nothing) do args...

test/downstream/observables_autodiff.jl

@@ -33,16 +33,16 @@ sol = solve(prob, Tsit5())
     gs, = gradient(sol) do sol
         sum(sol[sys.w])
     end
-    du_ = [0.0, 1.0, 1.0, 1.0]
-    du = [du_ for _ in sol.u]
+    du_ = [1.0, 1.0, 1.0, 0.0]
+    du = [du_ for _ in sol[[D(x), x, y, z]]]
     @test du == gs.u
 
     # Observable in a vector
     gs, = gradient(sol) do sol
         sum(sum.(sol[[sys.w, sys.x]]))
     end
-    du_ = [0.0, 1.0, 1.0, 2.0]
-    du = [du_ for _ in sol.u]
+    du_ = [1.0, 1.0, 2.0, 0.0]
+    du = [du_ for _ in sol[[D(x), x, y, z]]]
     @test du == gs.u
 end
 
@@ -118,14 +118,19 @@ end
 end
 
 @testset "Adjoints with DAE" begin
-    gs_mtkp, gs_p_new = gradient(prob.p, prob.p.tunable) do p, new_tunables
+    model = create_model()
+    sys = mtkcompile(model)
+    prob = ODEProblem(sys, [], (0.0, 1.0))
+    tunables, _, _ = SS.canonicalize(SS.Tunable(), prob.p)
+
+    gs_mtkp, gs_p_new = gradient(prob.p, tunables) do p, new_tunables
         new_p = SS.replace(SS.Tunable(), p, new_tunables)
         new_prob = remake(prob, p = new_p)
         sol = solve(new_prob, Rodas4())
-        mean(abs.(sol[sys.ampermeter.i] .- gt))
         sum(sol[sys.ampermeter.i])
     end
 
     @test isnothing(gs_mtkp)
-    @test length(gs_p_new) == length(p_new)
+    @test !isnothing(gs_p_new)
+    @test length(gs_p_new) == length(tunables)
 end

test/downstream/solution_interface.jl

@@ -6,29 +6,29 @@ using Plots: Plots, plot
 
 ### Tests on non-layered model (everything should work). ###
 
-@parameters a b c d
-@variables s1(t) s2(t)
+@testset "Basic indexing" begin
+    @parameters a b c d
+    @variables s1(t) s2(t)
+    
+    eqs = [D(s1) ~ a * s1 / (1 + s1 + s2) - b * s1,
+        D(s2) ~ +c * s2 / (1 + s1 + s2) - d * s2]
+    
+    @mtkcompile population_model = System(eqs, t)
+    
+    # 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; p], tspan)
+    sol = solve(oprob, Rodas4())
+    
+    @test sol[s1] == sol[population_model.s1] == sol[:s1]
+    @test sol[s2] == sol[population_model.s2] == sol[:s2]
+    @test sol[s1][end] ≈ 1.0
+    @test_throws Exception sol[a]
+    @test_throws Exception sol[population_model.a]
+    @test_throws Exception sol[:a]
 
-eqs = [D(s1) ~ a * s1 / (1 + s1 + s2) - b * s1,
-    D(s2) ~ +c * s2 / (1 + s1 + s2) - d * s2]
-
-@mtkcompile population_model = System(eqs, t)
-
-# 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; p], tspan)
-sol = solve(oprob, Rodas4())
-
-@test sol[s1] == sol[population_model.s1] == sol[:s1]
-@test sol[s2] == sol[population_model.s2] == sol[:s2]
-@test sol[s1][end] ≈ 1.0
-@test_throws Exception sol[a]
-@test_throws Exception sol[population_model.a]
-@test_throws Exception sol[:a]
-
-@testset "plot ODE solution" begin
     Plots.unicodeplots()
     f = ODEFunction((u, p, t) -> -u, analytic = (u0, p, t) -> u0 * exp(-t))
 
@@ -49,104 +49,104 @@ sol = solve(oprob, Rodas4())
     sol = solve(ode, Tsit5())
     @test_nowarn plot(sol)
     @test_nowarn plot(sol; plot_analytic = true)
-end
-
-# Tests on SDEProblem
-noiseeqs = [0.1 * s1,
-    0.1 * s2]
-@named noisy_population_model = SDESystem(population_model, noiseeqs)
-noisy_population_model = complete(noisy_population_model)
-sprob = SDEProblem(noisy_population_model, [u0; p], (0.0, 100.0))
-sol = solve(sprob, ImplicitEM())
-
-@test sol[s1] == sol[noisy_population_model.s1] == sol[:s1]
-@test sol[s2] == sol[noisy_population_model.s2] == sol[:s2]
-@test_throws Exception sol[a]
-@test_throws Exception sol[noisy_population_model.a]
-@test_throws Exception sol[:a]
-@test_nowarn sol(0.5, idxs = noisy_population_model.s1)
-### Tests on layered model (some things should not work). ###
-
-@parameters σ ρ β
-@variables x(t) y(t) z(t)
-
-eqs = [D(x) ~ σ * (y - x),
-    D(y) ~ x * (ρ - z) - y,
-    D(z) ~ x * y - β * z]
-
-@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 * γ]
-@mtkcompile sys = System(connections, t, [a, α], [γ], systems = [lorenz1, lorenz2])
-
-u0 = [lorenz1.x => 1.0,
-    lorenz1.y => 0.0,
-    lorenz1.z => 0.0,
-    lorenz2.x => 0.0,
-    lorenz2.y => 1.0,
-    lorenz2.z => 0.0]
 
-p = [lorenz1.σ => 10.0,
-    lorenz1.ρ => 28.0,
-    lorenz1.β => 8 / 3,
-    lorenz2.σ => 10.0,
-    lorenz2.ρ => 28.0,
-    lorenz2.β => 8 / 3,
-    γ => 2.0]
-
-tspan = (0.0, 100.0)
-prob = ODEProblem(sys, [u0; p], tspan)
-sol = solve(prob, Rodas4())
-
-@test_throws ArgumentError sol[x]
-@test in(sol[lorenz1.x], [getindex.(sol.u, i) for i in 1:length(unknowns(sol.prob.f.sys))])
-@test_throws KeyError sol[:x]
-
-### Non-symbolic indexing tests
-@test sol[:, 1] isa AbstractVector
-@test sol[:, 1:2] isa AbstractDiffEqArray
-@test sol[:, [1, 2]] isa AbstractDiffEqArray
-
-sol1 = sol(0.0:1.0:10.0)
-@test sol1.u isa Vector
-@test first(sol1.u) isa Vector
-@test length(sol1.u) == 11
-@test length(sol1.t) == 11
-
-sol2 = sol(0.1)
-@test sol2 isa Vector
-@test length(sol2) == length(unknowns(sys))
-@test first(sol2) isa Real
-
-sol3 = sol(0.0:1.0:10.0, idxs = [lorenz1.x, lorenz2.x])
-
-sol7 = sol(0.0:1.0:10.0, idxs = [2, 1])
-@test sol7.u isa Vector
-@test first(sol7.u) isa Vector
-@test length(sol7.u) == 11
-@test length(sol7.t) == 11
-@test collect(sol7[t]) ≈ sol3.t
-@test collect(sol7[t, 1:5]) ≈ sol3.t[1:5]
-
-sol8 = sol(0.1, idxs = [2, 1])
-@test sol8 isa Vector
-@test length(sol8) == 2
-@test first(sol8) isa Real
-
-sol9 = sol(0.0:1.0:10.0, idxs = 2)
-@test sol9.u isa Vector
-@test first(sol9.u) isa Real
-@test length(sol9.u) == 11
-@test length(sol9.t) == 11
-@test collect(sol9[t]) ≈ sol3.t
-@test collect(sol9[t, 1:5]) ≈ sol3.t[1:5]
-
-sol10 = sol(0.1, idxs = 2)
-@test sol10 isa Real
+    # Tests on SDEProblem
+    noiseeqs = [0.1 * s1,
+        0.1 * s2]
+    @named noisy_population_model = SDESystem(population_model, noiseeqs)
+    noisy_population_model = complete(noisy_population_model)
+    sprob = SDEProblem(noisy_population_model, [u0; p], (0.0, 100.0))
+    sol = solve(sprob, ImplicitEM())
+    
+    @test sol[s1] == sol[noisy_population_model.s1] == sol[:s1]
+    @test sol[s2] == sol[noisy_population_model.s2] == sol[:s2]
+    @test_throws Exception sol[a]
+    @test_throws Exception sol[noisy_population_model.a]
+    @test_throws Exception sol[:a]
+    @test_nowarn sol(0.5, idxs = noisy_population_model.s1)
+    ### Tests on layered model (some things should not work). ###
+    
+    @parameters σ ρ β
+    @variables x(t) y(t) z(t)
+    
+    eqs = [D(x) ~ σ * (y - x),
+        D(y) ~ x * (ρ - z) - y,
+        D(z) ~ x * y - β * z]
+    
+    @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 * γ]
+    @mtkcompile sys = System(connections, t, [a, α], [γ], systems = [lorenz1, lorenz2])
+    
+    u0 = [lorenz1.x => 1.0,
+        lorenz1.y => 0.0,
+        lorenz1.z => 0.0,
+        lorenz2.x => 0.0,
+        lorenz2.y => 1.0,
+        lorenz2.z => 0.0]
+    
+    p = [lorenz1.σ => 10.0,
+        lorenz1.ρ => 28.0,
+        lorenz1.β => 8 / 3,
+        lorenz2.σ => 10.0,
+        lorenz2.ρ => 28.0,
+        lorenz2.β => 8 / 3,
+        γ => 2.0]
+    
+    tspan = (0.0, 100.0)
+    prob = ODEProblem(sys, [u0; p], tspan)
+    sol = solve(prob, Rodas4())
+    
+    @test_throws ArgumentError sol[x]
+    @test in(sol[lorenz1.x], [getindex.(sol.u, i) for i in 1:length(unknowns(sol.prob.f.sys))])
+    @test_throws KeyError sol[:x]
+    
+    ### Non-symbolic indexing tests
+    @test sol[:, 1] isa AbstractVector
+    @test sol[:, 1:2] isa AbstractDiffEqArray
+    @test sol[:, [1, 2]] isa AbstractDiffEqArray
+    
+    sol1 = sol(0.0:1.0:10.0)
+    @test sol1.u isa Vector
+    @test first(sol1.u) isa Vector
+    @test length(sol1.u) == 11
+    @test length(sol1.t) == 11
+    
+    sol2 = sol(0.1)
+    @test sol2 isa Vector
+    @test length(sol2) == length(unknowns(sys))
+    @test first(sol2) isa Real
+    
+    sol3 = sol(0.0:1.0:10.0, idxs = [lorenz1.x, lorenz2.x])
+    
+    sol7 = sol(0.0:1.0:10.0, idxs = [2, 1])
+    @test sol7.u isa Vector
+    @test first(sol7.u) isa Vector
+    @test length(sol7.u) == 11
+    @test length(sol7.t) == 11
+    @test collect(sol7[t]) ≈ sol3.t
+    @test collect(sol7[t, 1:5]) ≈ sol3.t[1:5]
+    
+    sol8 = sol(0.1, idxs = [2, 1])
+    @test sol8 isa Vector
+    @test length(sol8) == 2
+    @test first(sol8) isa Real
+    
+    sol9 = sol(0.0:1.0:10.0, idxs = 2)
+    @test sol9.u isa Vector
+    @test first(sol9.u) isa Real
+    @test length(sol9.u) == 11
+    @test length(sol9.t) == 11
+    @test collect(sol9[t]) ≈ sol3.t
+    @test collect(sol9[t, 1:5]) ≈ sol3.t[1:5]
+    
+    sol10 = sol(0.1, idxs = 2)
+    @test sol10 isa Real
+end
 
 @testset "Plot idxs" begin
     @variables x(t) y(t)
@@ -198,7 +198,7 @@ end
             (0.0, 1.0); build_initializeprob = false)
         dae_sol = solve(prob, DFBDF(); save_idxs = [x])
 
-        @brownian a b
+        @brownians a b
         @mtkcompile sys = System([D(x) ~ x + p * y + x * a, D(y) ~ 2p + x^2 + y * b], t)
         xidx = variable_index(sys, x)
         prob = SDEProblem(sys, [x => 1.0, y => 2.0, p => 2.0], (0.0, 1.0))

test/runtests.jl

@@ -123,27 +123,25 @@ end
         end
     end
 
-    if !is_APPVEYOR && (GROUP == "Downstream" || GROUP == "SymbolicIndexingInterface")
+    if !is_APPVEYOR && GROUP == "SymbolicIndexingInterface"
         if GROUP != "Downstream"
             activate_downstream_env()
         end
-        @time @safetestset "ModelingToolkit Remake" begin
-            include("downstream/modelingtoolkit_remake.jl")
-        end
         @time @safetestset "Symbol and integer based indexing of interpolated solutions" begin
             include("downstream/comprehensive_indexing.jl")
         end
-        if VERSION >= v"1.8"
-            @time @safetestset "Symbol and integer based indexing of integrators" begin
-                include("downstream/integrator_indexing.jl")
-            end
-            @time @safetestset "Problem Indexing" begin
-                include("downstream/problem_interface.jl")
-            end
+        @time @safetestset "Symbol and integer based indexing of integrators" begin
+            include("downstream/integrator_indexing.jl")
+        end
+        @time @safetestset "Problem Indexing" begin
+            include("downstream/problem_interface.jl")
         end
         @time @safetestset "Adjoints" begin
             include("downstream/adjoints.jl")
         end
+        @time @safetestset "ModelingToolkit Remake" begin
+            include("downstream/modelingtoolkit_remake.jl")
+        end
     end
 
     if !is_APPVEYOR && GROUP == "Python"