late binding initialization_eqs

Author: ChrisRackauckas

src/systems/abstractsystem.jl

@@ -576,6 +576,7 @@ for prop in [:eqs
              :preface
              :torn_matching
              :initializesystem
+             :initialization_eqs
              :schedule
              :tearing_state
              :substitutions

src/systems/diffeqs/abstractodesystem.jl

@@ -864,7 +864,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
 
     # TODO: Pass already computed information to varmap_to_vars call
     # in process_u0? That would just be a small optimization
-    varmap = isempty(u0map) ? defaults : merge(defaults, todict(u0map))
+    varmap = isempty(u0map) ? defaults(sys) : merge(defaults(sys), todict(u0map))
     varlist = collect(map(unwrap, dvs))
     missingvars = setdiff(varlist, collect(keys(varmap)))
 

src/systems/diffeqs/odesystem.jl

@@ -101,6 +101,10 @@ struct ODESystem <: AbstractODESystem
     """
     initializesystem::Union{Nothing, NonlinearSystem}
     """
+    Extra equations to be enforced during the initialization sequence.
+    """
+    initialization_eqs::Vector{Equation}
+    """
     The schedule for the code generation process.
     """
     schedule::Any
@@ -171,7 +175,8 @@ struct ODESystem <: AbstractODESystem
 
     function ODESystem(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad,
             jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, guesses,
-            torn_matching, initializesystem, schedule, connector_type, preface, cevents,
+            torn_matching, initializesystem, initialization_eqs, schedule,
+            connector_type, preface, cevents,
             devents, parameter_dependencies,
             metadata = nothing, gui_metadata = nothing,
             tearing_state = nothing,
@@ -190,8 +195,8 @@ struct ODESystem <: AbstractODESystem
         end
         new(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, guesses, torn_matching,
-            initializesystem, schedule, connector_type, preface, cevents, devents, parameter_dependencies,
-            metadata,
+            initializesystem, initialization_eqs, schedule, connector_type, preface,
+            cevents, devents, parameter_dependencies, metadata,
             gui_metadata, tearing_state, substitutions, complete, index_cache,
             discrete_subsystems, solved_unknowns, split_idxs, parent)
     end
@@ -208,6 +213,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         defaults = _merge(Dict(default_u0), Dict(default_p)),
         guesses = Dict(),
         initializesystem = nothing,
+        initialization_eqs = Equation[],
         schedule = nothing,
         connector_type = nothing,
         preface = nothing,
@@ -260,7 +266,8 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     ODESystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         deqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, tgrad, jac,
         ctrl_jac, Wfact, Wfact_t, name, systems, defaults, guesses, nothing, initializesystem,
-        schedule, connector_type, preface, cont_callbacks, disc_callbacks, parameter_dependencies,
+        initialization_eqs, schedule, connector_type, preface, cont_callbacks,
+        disc_callbacks, parameter_dependencies,
         metadata, gui_metadata, checks = checks)
 end
 

src/systems/nonlinear/initializesystem.jl

@@ -55,7 +55,7 @@ function generate_initializesystem(sys::ODESystem;
     end
 
     pars = [parameters(sys); get_iv(sys)]
-    nleqs = [eqs_ics; observed(sys)]
+    nleqs = [eqs_ics; get_initialization_eqs(sys); observed(sys)]
 
     sys_nl = NonlinearSystem(nleqs,
         full_states,

test/initializationsystem.jl

@@ -355,3 +355,24 @@ prob = ODEProblem(sys, [sys.ddx => -2], (0, 1), guesses = [sys.dx => 1])
 sol = solve(prob, Tsit5())
 @test SciMLBase.successful_retcode(sol)
 @test sol[1] == [1.0]
+
+## Late binding initialization_eqs
+
+function System(; name)
+    vars = @variables begin
+        dx(t), [guess = 0]
+        ddx(t), [guess = 0]
+    end
+    eqs = [D(dx) ~ ddx
+           0 ~ ddx + dx + 1]
+    initialization_eqs = [
+        ddx ~ -2
+    ]
+    return ODESystem(eqs, t, vars, []; name, initialization_eqs)
+end
+
+@mtkbuild sys = System()
+prob = ODEProblem(sys, [], (0, 1), guesses = [sys.dx => 1])
+sol = solve(prob, Tsit5())
+@test SciMLBase.successful_retcode(sol)
+@test sol[1] == [1.0]