Merge pull request #2789 from SciML/late_initial

allow for late binding of initialization equations

Author: ChrisRackauckas

src/systems/diffeqs/abstractodesystem.jl

@@ -853,6 +853,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
         t = nothing,
         warn_initialize_determined = true,
         build_initializeprob = true,
+        initialization_eqs = [],
         kwargs...)
     eqs = equations(sys)
     dvs = unknowns(sys)
@@ -925,7 +926,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
             u0map = Dict()
         end
         initializeprob = ModelingToolkit.InitializationProblem(
-            sys, t, u0map, parammap; guesses, warn_initialize_determined)
+            sys, t, u0map, parammap; guesses, warn_initialize_determined, initialization_eqs)
         initializeprobmap = getu(initializeprob, unknowns(sys))
 
         zerovars = Dict(setdiff(unknowns(sys), keys(defaults(sys))) .=> 0.0)
@@ -1555,6 +1556,7 @@ InitializationProblem{iip}(sys::AbstractODESystem, u0map, tspan,
                            checkbounds = false, sparse = false,
                            simplify = false,
                            linenumbers = true, parallel = SerialForm(),
+                           initialization_eqs = [],
                            kwargs...) where {iip}
 ```
 
@@ -1603,17 +1605,19 @@ function InitializationProblem{iip, specialize}(sys::AbstractODESystem,
         guesses = [],
         check_length = true,
         warn_initialize_determined = true,
+        initialization_eqs = [],
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEProblem`")
     end
     if isempty(u0map) && get_initializesystem(sys) !== nothing
-        isys = get_initializesystem(sys)
+        isys = get_initializesystem(sys; initialization_eqs)
     elseif isempty(u0map) && get_initializesystem(sys) === nothing
-        isys = structural_simplify(generate_initializesystem(sys); fully_determined = false)
+        isys = structural_simplify(
+            generate_initializesystem(sys; initialization_eqs); fully_determined = false)
     else
         isys = structural_simplify(
-            generate_initializesystem(sys; u0map); fully_determined = false)
+            generate_initializesystem(sys; u0map, initialization_eqs); fully_determined = false)
     end
 
     uninit = setdiff(unknowns(sys), [unknowns(isys); getfield.(observed(isys), :lhs)])

src/systems/nonlinear/initializesystem.jl

@@ -9,6 +9,7 @@ function generate_initializesystem(sys::ODESystem;
         guesses = Dict(), check_defguess = false,
         default_dd_value = 0.0,
         algebraic_only = false,
+        initialization_eqs = [],
         kwargs...)
     sts, eqs = unknowns(sys), equations(sys)
     idxs_diff = isdiffeq.(eqs)
@@ -95,7 +96,7 @@ function generate_initializesystem(sys::ODESystem;
     nleqs = if algebraic_only
         [eqs_ics; observed(sys)]
     else
-        [eqs_ics; get_initialization_eqs(sys); observed(sys)]
+        [eqs_ics; get_initialization_eqs(sys); initialization_eqs; observed(sys)]
     end
 
     sys_nl = NonlinearSystem(nleqs,

test/initializationsystem.jl

@@ -430,3 +430,16 @@ sol = solve(prob, Tsit5())
 
 # This should warn, but logging tests can't be marked as broken
 @test_logs prob = ODEProblem(simpsys, [], tspan, guesses = [x => 2.0])
+
+# Late Binding initialization_eqs
+# https://github.com/SciML/ModelingToolkit.jl/issues/2787
+
+@parameters g
+@variables x(t) y(t) [state_priority = 10] λ(t)
+eqs = [D(D(x)) ~ λ * x
+       D(D(y)) ~ λ * y - g
+       x^2 + y^2 ~ 1]
+@mtkbuild pend = ODESystem(eqs, t)
+
+prob = ODEProblem(pend, [x => 1], (0.0, 1.5), [g => 1],
+    guesses = [λ => 0, y => 1], initialization_eqs = [y ~ 1])