Merge pull request #2984 from AayushSabharwal/as/discover-pdeps

fix: discover parameters from parameter dependencies, fix substitute

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -2940,7 +2940,14 @@ function Symbolics.substitute(sys::AbstractSystem, rules::Union{Vector{<:Pair},
         rules = todict(map(r -> Symbolics.unwrap(r[1]) => Symbolics.unwrap(r[2]),
             collect(rules)))
         eqs = fast_substitute(equations(sys), rules)
-        ODESystem(eqs, get_iv(sys); name = nameof(sys))
+        pdeps = fast_substitute(parameter_dependencies(sys), rules)
+        defs = Dict(fast_substitute(k, rules) => fast_substitute(v, rules)
+        for (k, v) in defaults(sys))
+        guess = Dict(fast_substitute(k, rules) => fast_substitute(v, rules)
+        for (k, v) in guesses(sys))
+        subsys = map(s -> substitute(s, rules), get_systems(sys))
+        ODESystem(eqs, get_iv(sys); name = nameof(sys), defaults = defs,
+            guesses = guess, parameter_dependencies = pdeps, systems = subsys)
     else
         error("substituting symbols is not supported for $(typeof(sys))")
     end

src/systems/diffeqs/odesystem.jl

@@ -311,6 +311,15 @@ function ODESystem(eqs, iv; kwargs...)
             push!(algeeq, eq)
         end
     end
+    for eq in get(kwargs, :parameter_dependencies, Equation[])
+        if eq isa Pair
+            collect_vars!(allunknowns, ps, eq[1], iv)
+            collect_vars!(allunknowns, ps, eq[2], iv)
+        else
+            collect_vars!(allunknowns, ps, eq.lhs, iv)
+            collect_vars!(allunknowns, ps, eq.rhs, iv)
+        end
+    end
     for v in allunknowns
         isdelay(v, iv) || continue
         collect_vars!(allunknowns, ps, arguments(v)[1], iv)

src/systems/discrete_system/discrete_system.jl

@@ -183,6 +183,15 @@ function DiscreteSystem(eqs, iv; kwargs...)
             push!(diffvars, eq.lhs)
         end
     end
+    for eq in get(kwargs, :parameter_dependencies, Equation[])
+        if eq isa Pair
+            collect_vars!(allunknowns, ps, eq[1], iv)
+            collect_vars!(allunknowns, ps, eq[2], iv)
+        else
+            collect_vars!(allunknowns, ps, eq.lhs, iv)
+            collect_vars!(allunknowns, ps, eq.rhs, iv)
+        end
+    end
     new_ps = OrderedSet()
     for p in ps
         if iscall(p) && operation(p) === getindex

src/systems/nonlinear/nonlinearsystem.jl

@@ -170,6 +170,15 @@ function NonlinearSystem(eqs; kwargs...)
         collect_vars!(allunknowns, ps, eq.lhs, nothing)
         collect_vars!(allunknowns, ps, eq.rhs, nothing)
     end
+    for eq in get(kwargs, :parameter_dependencies, Equation[])
+        if eq isa Pair
+            collect_vars!(allunknowns, ps, eq[1], nothing)
+            collect_vars!(allunknowns, ps, eq[2], nothing)
+        else
+            collect_vars!(allunknowns, ps, eq.lhs, nothing)
+            collect_vars!(allunknowns, ps, eq.rhs, nothing)
+        end
+    end
     new_ps = OrderedSet()
     for p in ps
         if iscall(p) && operation(p) === getindex

test/odesystem.jl

@@ -1315,3 +1315,19 @@ end
     @named sys = compose(sys, sys) # nest into a hierarchical system
     @test t === sys.t === sys.sys.t
 end
+
+@testset "Substituting preserves parameter dependencies, defaults, guesses" begin
+    @parameters p1 p2
+    @variables x(t) y(t)
+    @named sys = ODESystem([D(x) ~ y + p2], t; parameter_dependencies = [p2 ~ 2p1],
+        defaults = [p1 => 1.0, p2 => 2.0], guesses = [p1 => 2.0, p2 => 3.0])
+    @parameters p3
+    sys2 = substitute(sys, [p1 => p3])
+    @test length(parameters(sys2)) == 1
+    @test is_parameter(sys2, p3)
+    @test !is_parameter(sys2, p1)
+    @test length(ModelingToolkit.defaults(sys2)) == 2
+    @test ModelingToolkit.defaults(sys2)[p3] == 1.0
+    @test length(ModelingToolkit.guesses(sys2)) == 2
+    @test ModelingToolkit.guesses(sys2)[p3] == 2.0
+end

test/parameter_dependencies.jl

@@ -325,3 +325,16 @@ end
     @test getp(sys, p1)(ps2) == 2.0
     @test getp(sys, p2)(ps2) == 4.0
 end
+
+@testset "Discovery of parameters from dependencies" begin
+    @parameters p1 p2
+    @variables x(t) y(t)
+    @named sys = ODESystem([D(x) ~ y + p2], t; parameter_dependencies = [p2 ~ 2p1])
+    @test is_parameter(sys, p1)
+    @named sys = NonlinearSystem([x * y^2 ~ y + p2]; parameter_dependencies = [p2 ~ 2p1])
+    @test is_parameter(sys, p1)
+    k = ShiftIndex(t)
+    @named sys = DiscreteSystem(
+        [x(k - 1) ~ x(k) + y(k) + p2], t; parameter_dependencies = [p2 ~ 2p1])
+    @test is_parameter(sys, p1)
+end

test/split_parameters.jl

@@ -117,7 +117,7 @@ sol = solve(prob, ImplicitEuler());
 
 # ------------------------ Mixed Type Conserved
 
-prob = ODEProblem(sys, [], tspan, []; tofloat = false)
+prob = ODEProblem(sys, [], tspan, []; tofloat = false, use_union = true)
 
 @test prob.p isa Tuple{Vector{Float64}, Vector{Int64}}
 sol = solve(prob, ImplicitEuler());