feat: store InitializationMetadata in OverrideInitData

Author: AayushSabharwal

src/systems/nonlinear/initializesystem.jl

@@ -481,22 +481,19 @@ function SciMLBase.remake_initialization_data(
     if u0 === missing && p === missing
         return odefn.initialization_data
     end
+
+    oldinitdata = odefn.initialization_data
+
     if !(eltype(u0) <: Pair) && !(eltype(p) <: Pair)
-        oldinitdata = odefn.initialization_data
         oldinitdata === nothing && return nothing
 
         oldinitprob = oldinitdata.initializeprob
         oldinitprob === nothing && return nothing
-        if !SciMLBase.has_sys(oldinitprob.f) || !(oldinitprob.f.sys isa NonlinearSystem)
-            return oldinitdata
-        end
-        oldinitsys = oldinitprob.f.sys
-        meta = get_metadata(oldinitsys)
-        if meta isa InitializationSystemMetadata && meta.oop_reconstruct_u0_p !== nothing
-            reconstruct_fn = meta.oop_reconstruct_u0_p
-        else
-            reconstruct_fn = ReconstructInitializeprob(sys, oldinitsys)
-        end
+
+        meta = oldinitdata.metadata
+        meta isa InitializationMetadata || return oldinitdata
+
+        reconstruct_fn = meta.oop_reconstruct_u0_p
         # the history function doesn't matter because `reconstruct_fn` is only going to
         # update the values of parameters, which aren't time dependent. The reason it
         # is called is because `Initial` parameters are calculated from the corresponding
@@ -507,16 +504,15 @@ function SciMLBase.remake_initialization_data(
         if oldinitprob.f.resid_prototype === nothing
             newf = oldinitprob.f
         else
-            newf = NonlinearFunction{
-                SciMLBase.isinplace(oldinitprob.f), SciMLBase.specialization(oldinitprob.f)}(
-                oldinitprob.f;
+            newf = remake(oldinitprob.f;
                 resid_prototype = calculate_resid_prototype(
                     length(oldinitprob.f.resid_prototype), new_initu0, new_initp))
         end
         initprob = remake(oldinitprob; f = newf, u0 = new_initu0, p = new_initp)
         return SciMLBase.OverrideInitData(initprob, oldinitdata.update_initializeprob!,
-            oldinitdata.initializeprobmap, oldinitdata.initializeprobpmap)
+            oldinitdata.initializeprobmap, oldinitdata.initializeprobpmap; metadata = oldinitdata.metadata)
     end
+
     dvs = unknowns(sys)
     ps = parameters(sys)
     u0map = to_varmap(u0, dvs)
@@ -530,16 +526,13 @@ function SciMLBase.remake_initialization_data(
     use_scc = true
     initialization_eqs = Equation[]
 
-    if SciMLBase.has_initializeprob(odefn)
-        oldsys = odefn.initialization_data.initializeprob.f.sys
-        meta = get_metadata(oldsys)
-        if meta isa InitializationSystemMetadata
-            u0map = merge(meta.u0map, u0map)
-            pmap = merge(meta.pmap, pmap)
-            merge!(guesses, meta.additional_guesses)
-            use_scc = get(meta.extra_metadata, :use_scc, true)
-            initialization_eqs = meta.additional_initialization_eqs
-        end
+    if oldinitdata !== nothing && oldinitdata.metadata isa InitializationMetadata
+        meta = oldinitdata.metadata
+        u0map = merge(meta.u0map, u0map)
+        pmap = merge(meta.pmap, pmap)
+        merge!(guesses, meta.guesses)
+        use_scc = meta.use_scc
+        initialization_eqs = meta.additional_initialization_eqs
     else
         # there is no initializeprob, so the original problem construction
         # had no solvable parameters and had the differential variables
@@ -600,19 +593,22 @@ function SciMLBase.late_binding_update_u0_p(
     if !(eltype(u0) <: Pair)
         # if `p` is not provided or is symbolic
         p === missing || eltype(p) <: Pair || return newu0, newp
-        newu0 === nothing && return newu0, newp
-        all(is_parameter(sys, Initial(x)) for x in unknowns(sys)) || return newu0, newp
+        (newu0 === nothing || isempty(newu0)) && return newu0, newp
+        initdata = prob.f.initialization_data
+        initdata === nothing && return newu0, newp
+        meta = initdata.metadata
+        meta isa InitializationMetadata || return newu0, newp
         newp = p === missing ? copy(newp) : newp
         initials, repack, alias = SciMLStructures.canonicalize(
             SciMLStructures.Initials(), newp)
         if eltype(initials) != eltype(newu0)
             initials = DiffEqBase.promote_u0(initials, newu0, t0)
             newp = repack(initials)
         end
-        if length(newu0) != length(unknowns(sys))
-            throw(ArgumentError("Expected `newu0` to be of same length as unknowns ($(length(unknowns(sys)))). Got $(typeof(newu0)) of length $(length(newu0))"))
+        if length(newu0) != length(prob.u0)
+            throw(ArgumentError("Expected `newu0` to be of same length as unknowns ($(length(prob.u0))). Got $(typeof(newu0)) of length $(length(newu0))"))
         end
-        setp(sys, Initial.(unknowns(sys)))(newp, newu0)
+        meta.set_initial_unknowns!(newp, newu0)
         return newu0, newp
     end
 

src/systems/problem_utils.jl

@@ -620,6 +620,49 @@ function build_operating_point!(sys::AbstractSystem,
     return op, missing_unknowns, missing_pars
 end
 
+"""
+    $(TYPEDEF)
+
+Metadata attached to `OverrideInitData` used in `remake` hooks for handling initialization
+properly.
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
+struct InitializationMetadata{R <: ReconstructInitializeprob, SIU}
+    """
+    The `u0map` used to construct the initialization.
+    """
+    u0map::Dict{Any, Any}
+    """
+    The `pmap` used to construct the initialization.
+    """
+    pmap::Dict{Any, Any}
+    """
+    The `guesses` used to construct the initialization.
+    """
+    guesses::Dict{Any, Any}
+    """
+    The `initialization_eqs` in addition to those of the system that were used to construct
+    the initialization.
+    """
+    additional_initialization_eqs::Vector{Equation}
+    """
+    Whether to use `SCCNonlinearProblem` if possible.
+    """
+    use_scc::Bool
+    """
+    `ReconstructInitializeprob` for this initialization problem.
+    """
+    oop_reconstruct_u0_p::R
+    """
+    A function which takes the `u0` of the problem and sets
+    `Initial.(unknowns(sys))`.
+    """
+    set_initial_unknowns!::SIU
+end
+
 """
     $(TYPEDSIGNATURES)
 
@@ -632,16 +675,16 @@ All other keyword arguments are forwarded to `InitializationProblem`.
 """
 function maybe_build_initialization_problem(
         sys::AbstractSystem, op::AbstractDict, u0map, pmap, t, defs,
-        guesses, missing_unknowns; implicit_dae = false,
-        u0_constructor = identity, floatT = Float64, kwargs...)
+        guesses, missing_unknowns; implicit_dae = false, u0_constructor = identity,
+        floatT = Float64, initialization_eqs = [], use_scc = true, kwargs...)
     guesses = merge(ModelingToolkit.guesses(sys), todict(guesses))
 
     if t === nothing && is_time_dependent(sys)
         t = zero(floatT)
     end
 
     initializeprob = ModelingToolkit.InitializationProblem{true, SciMLBase.FullSpecialize}(
-        sys, t, u0map, pmap; guesses, kwargs...)
+        sys, t, u0map, pmap; guesses, initialization_eqs, use_scc, kwargs...)
     if state_values(initializeprob) !== nothing
         initializeprob = remake(initializeprob; u0 = floatT.(state_values(initializeprob)))
     end
@@ -658,7 +701,10 @@ function maybe_build_initialization_problem(
     end
     initializeprob = remake(initializeprob; p = initp)
 
-    meta = get_metadata(initializeprob.f.sys)
+    meta = InitializationMetadata(
+        u0map, pmap, guesses, Equation[get_initialization_eqs(sys); initialization_eqs],
+        use_scc, ReconstructInitializeprob(sys, initializeprob.f.sys),
+        setp(sys, Initial.(unknowns(sys))))
 
     if is_time_dependent(sys)
         all_init_syms = Set(all_symbols(initializeprob))
@@ -710,7 +756,7 @@ function maybe_build_initialization_problem(
     return (;
         initialization_data = SciMLBase.OverrideInitData(
             initializeprob, update_initializeprob!, initializeprobmap,
-            initializeprobpmap))
+            initializeprobpmap; metadata = meta))
 end
 
 """