Automate tagging of the initialization system to the ODEProblem

Author: ChrisRackauckas

Project.toml

@@ -94,7 +94,7 @@ PrecompileTools = "1"
 RecursiveArrayTools = "2.3, 3"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
-SciMLBase = "2.0.1"
+SciMLBase = "2.27"
 SciMLStructures = "1.0"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"

src/systems/abstractsystem.jl

@@ -537,8 +537,8 @@ function complete(sys::AbstractSystem; split = true)
     if split && has_index_cache(sys)
         @set! sys.index_cache = IndexCache(sys)
     end
-    if isdefined(sys, :initializationsystem)
-        @set! sys.initializationsystem = complete(get_initializationsystem(sys); split)
+    if isdefined(sys, :initializesystem) && get_initializesystem(sys) !== nothing
+        @set! sys.initializesystem = complete(get_initializesystem(sys); split)
     end
     isdefined(sys, :complete) ? (@set! sys.complete = true) : sys
 end

src/systems/diffeqs/abstractodesystem.jl

@@ -315,7 +315,9 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         checkbounds = false,
         sparsity = false,
         analytic = nothing,
-        split_idxs = nothing,
+        split_idxs = nothing,  
+        initializeprob =  nothing,
+        initializeprobmap = nothing,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEFunction`")
@@ -487,6 +489,7 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
     end
 
     @set! sys.split_idxs = split_idxs
+
     ODEFunction{iip, specialize}(f;
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
@@ -495,7 +498,9 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem,
         jac_prototype = jac_prototype,
         observed = observedfun,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
-        analytic = analytic)
+        analytic = analytic,
+        initializeprob = initializeprob,
+        initializeprobmap = initializeprobmap)
 end
 
 """
@@ -525,6 +530,8 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         sparse = false, simplify = false,
         eval_module = @__MODULE__,
         checkbounds = false,
+        initializeprob =  nothing,
+        initializeprobmap = nothing,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `DAEFunction`")
@@ -596,7 +603,9 @@ function DiffEqBase.DAEFunction{iip}(sys::AbstractODESystem, dvs = unknowns(sys)
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
         jac_prototype = jac_prototype,
-        observed = observedfun)
+        observed = observedfun,
+        initializeprob = initializeprob,
+        initializeprobmap = initializeprobmap)
 end
 
 function DiffEqBase.DDEFunction(sys::AbstractODESystem, args...; kwargs...)
@@ -877,10 +886,15 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
 
     check_eqs_u0(eqs, dvs, u0; kwargs...)
 
+    initializeprob =  ModelingToolkit.InitializationProblem(sys, u0map, parammap)
+    initializeprobmap = getu(initializeprob, unknowns(sys))
+
     f = constructor(sys, dvs, ps, u0; ddvs = ddvs, tgrad = tgrad, jac = jac,
         checkbounds = checkbounds, p = p,
         linenumbers = linenumbers, parallel = parallel, simplify = simplify,
         sparse = sparse, eval_expression = eval_expression,
+        initializeprob = initializeprob, 
+        initializeprobmap = initializeprobmap,
         kwargs...)
     implicit_dae ? (f, du0, u0, p) : (f, u0, p)
 end

test/initializationsystem.jl

@@ -181,6 +181,24 @@ prob = ODEProblem(sys, allinit, (0, 0.1))
 sol = solve(prob, Rodas5P())
 # If initialized incorrectly, then it would be InitialFailure
 @test sol.retcode == SciMLBase.ReturnCode.Unstable
+SciMLBase.has_initializeprob(prob.f)
+
+isys = ModelingToolkit.get_initializesystem(sys)
+unknowns(isys)
+
+initprob = ModelingToolkit.InitializationProblem(sys)
+sol = solve(initprob)
+
+unknowns(sys)
+
+[sys.act.vol₁.dr]
+
+getter = ModelingToolkit.getu(initprob, unknowns(sys)[end-1:end])
+getter(sol)
+
+prob.f.initializeprobmap(initsol)
+
+initsol[unknowns(isys)]
 
 @connector Flange begin
     dx(t), [guess = 0]
@@ -250,3 +268,4 @@ prob = ODEProblem(sys, allinit, (0, 0.1))
 sol = solve(prob, Rodas5P())
 # If initialized incorrectly, then it would be InitialFailure
 @test sol.retcode == SciMLBase.ReturnCode.Success
+