refactor: store getu function in InitializationMetadata

Author: AayushSabharwal

src/systems/nonlinear/initializesystem.jl

@@ -588,17 +588,26 @@ end
 function SciMLBase.late_binding_update_u0_p(
         prob, sys::AbstractSystem, u0, p, t0, newu0, newp)
     supports_initialization(sys) || return newu0, newp
+
+    initdata = prob.f.initialization_data
+    meta = initdata === nothing ? nothing : initdata.metadata
     # If the user passes `p` to `remake` but not `u0` and `u0` isn't empty,
     # and if the system supports initialization (so it has initial parameters),
     # and if the initialization solves for `u0`,
     # THEN copy the values of `Initial`s to `newu0`.
     if u0 === missing
-        if newu0 !== nothing && p !== missing && supports_initialization(sys) && prob.f.initialization_data !== nothing && prob.f.initialization_data.initializeprobmap !== nothing
+        if newu0 !== nothing && p !== missing && supports_initialization(sys) &&
+           initdata !== nothing && initdata.initializeprobmap !== nothing
+            getter = if meta isa InitializationMetadata
+                meta.get_initial_unknowns
+            else
+                getu(sys, Initial.(unknowns(sys)))
+            end
             if ArrayInterface.ismutable(newu0)
-                copyto!(newu0, getu(sys, Initial.(unknowns(sys)))(newp))
+                copyto!(newu0, getter(newp))
             else
                 T = StaticArrays.similar_type(newu0)
-                newu0 = T(getu(sys, Initial.(unknowns(sys)))(newp))
+                newu0 = T(getter(newp))
             end
         end
         return newu0, newp
@@ -609,7 +618,6 @@ function SciMLBase.late_binding_update_u0_p(
         # if `p` is not provided or is symbolic
         p === missing || eltype(p) <: Pair || 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

src/systems/problem_utils.jl

@@ -769,7 +769,7 @@ properly.
 
 $(TYPEDFIELDS)
 """
-struct InitializationMetadata{R <: ReconstructInitializeprob, SIU}
+struct InitializationMetadata{R <: ReconstructInitializeprob, GIU, SIU}
     """
     The `u0map` used to construct the initialization.
     """
@@ -796,6 +796,11 @@ struct InitializationMetadata{R <: ReconstructInitializeprob, SIU}
     """
     oop_reconstruct_u0_p::R
     """
+    A function which takes the parameter object of the problem and returns
+    `Initial.(unknowns(sys))`.
+    """
+    get_initial_unknowns::GIU
+    """
     A function which takes the `u0` of the problem and sets
     `Initial.(unknowns(sys))`.
     """
@@ -843,7 +848,7 @@ function maybe_build_initialization_problem(
     meta = InitializationMetadata(
         u0map, pmap, guesses, Vector{Equation}(initialization_eqs),
         use_scc, ReconstructInitializeprob(sys, initializeprob.f.sys),
-        setp(sys, Initial.(unknowns(sys))))
+        getp(sys, Initial.(unknowns(sys))), setp(sys, Initial.(unknowns(sys))))
 
     if is_time_dependent(sys)
         all_init_syms = Set(all_symbols(initializeprob))