fix: propagate guesses for algebraic variables instead of Initial(x)

Author: AayushSabharwal

src/systems/nonlinear/initializesystem.jl

@@ -653,20 +653,16 @@ end
 function DiffEqBase.get_updated_symbolic_problem(sys::AbstractSystem, prob; kw...)
     supports_initialization(sys) || return prob
     initdata = prob.f.initialization_data
-    initdata === nothing && return prob
+    initdata isa SciMLBase.OverrideInitData || return prob
+    meta = initdata.metadata
+    meta isa InitializationMetadata || return prob
+    meta.get_updated_u0 === nothing && return prob
 
     u0 = state_values(prob)
     u0 === nothing && return prob
 
     p = parameter_values(prob)
     t0 = is_time_dependent(prob) ? current_time(prob) : nothing
-    meta = initdata.metadata
-
-    getter = if meta isa InitializationMetadata
-        meta.get_initial_unknowns
-    else
-        getu(sys, Initial.(unknowns(sys)))
-    end
 
     if p isa MTKParameters
         buffer = p.initials
@@ -681,7 +677,8 @@ function DiffEqBase.get_updated_symbolic_problem(sys::AbstractSystem, prob; kw..
     else
         T = StaticArrays.similar_type(u0)
     end
-    return remake(prob; u0 = T(getter(p)))
+
+    return remake(prob; u0 = T(meta.get_updated_u0(prob, initdata.initializeprob)))
 end
 
 """

src/systems/problem_utils.jl

@@ -769,7 +769,7 @@ properly.
 
 $(TYPEDFIELDS)
 """
-struct InitializationMetadata{R <: ReconstructInitializeprob, GIU, SIU}
+struct InitializationMetadata{R <: ReconstructInitializeprob, GUU, SIU}
     """
     The `u0map` used to construct the initialization.
     """
@@ -796,17 +796,58 @@ struct InitializationMetadata{R <: ReconstructInitializeprob, GIU, SIU}
     """
     oop_reconstruct_u0_p::R
     """
-    A function which takes the parameter object of the problem and returns
-    `Initial.(unknowns(sys))`.
+    A function which takes `(prob, initializeprob)` and return the `u0` to use for the problem.
     """
-    get_initial_unknowns::GIU
+    get_updated_u0::GUU
     """
     A function which takes the `u0` of the problem and sets
     `Initial.(unknowns(sys))`.
     """
     set_initial_unknowns!::SIU
 end
 
+"""
+    $(TYPEDEF)
+
+A callable struct to use as the `get_updated_u0` field of `InitializationMetadata`.
+Returns the value of `Initial.(unknowns(sys))`, except with algebraic variables replaced
+by their guess values in the initialization problem.
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
+struct GetUpdatedU0{GA, GIU}
+    """
+    Mask with length `length(unknowns(sys))` denoting indices of algebraic variables.
+    """
+    algevars::BitVector
+    """
+    Function which returns the values of algebraic variables in `initializeprob`, in the
+    order the algebraic variables occur in `unknowns(sys)`.
+    """
+    get_algevars::GA
+    """
+    Function which returns `Initial.(unknowns(sys))` as a `Vector`.
+    """
+    get_initial_unknowns::GIU
+end
+
+function GetUpdatedU0(sys::AbstractSystem, initsys::AbstractSystem)
+    algevaridxs = BitVector(is_alg_equation.(equations(sys)))
+    algevars = unknowns(sys)[algevaridxs]
+    get_algevars = getu(initsys, algevars)
+    get_initial_unknowns = getu(sys, Initial.(unknowns(sys)))
+    return GetUpdatedU0(algevaridxs, get_algevars, get_initial_unknowns)
+end
+
+function (guu::GetUpdatedU0)(prob, initprob)
+    buffer = guu.get_initial_unknowns(prob)
+    algebuf = view(buffer, guu.algevars)
+    copyto!(algebuf, guu.get_algevars(initprob))
+    return buffer
+end
+
 """
     $(TYPEDSIGNATURES)
 
@@ -845,10 +886,15 @@ function maybe_build_initialization_problem(
     end
     initializeprob = remake(initializeprob; p = initp)
 
+    get_initial_unknowns = if is_time_dependent(sys)
+        GetUpdatedU0(sys, initializeprob.f.sys)
+    else
+        nothing
+    end
     meta = InitializationMetadata(
         u0map, pmap, guesses, Vector{Equation}(initialization_eqs),
         use_scc, ReconstructInitializeprob(sys, initializeprob.f.sys),
-        getp(sys, Initial.(unknowns(sys))), setp(sys, Initial.(unknowns(sys))))
+        get_initial_unknowns, setp(sys, Initial.(unknowns(sys))))
 
     if is_time_dependent(sys)
         all_init_syms = Set(all_symbols(initializeprob))