refactor: copy initials to u0 in solve/init instead of remake

Author: AayushSabharwal

src/systems/nonlinear/initializesystem.jl

@@ -595,27 +595,6 @@ function SciMLBase.late_binding_update_u0_p(
 
     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) &&
-           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, getter(newp))
-            else
-                T = StaticArrays.similar_type(newu0)
-                newu0 = T(getter(newp))
-            end
-        end
-        return newu0, newp
-    end
 
     # non-symbolic u0 updates initials...
     if !(eltype(u0) <: Pair)
@@ -669,6 +648,40 @@ function SciMLBase.late_binding_update_u0_p(
     return newu0, newp
 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
+
+    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
+    else
+        buffer = p
+    end
+
+    u0 = DiffEqBase.promote_u0(u0, buffer, t0)
+
+    if ArrayInterface.ismutable(u0)
+        T = typeof(u0)
+    else
+        T = StaticArrays.similar_type(u0)
+    end
+    return remake(prob; u0 = T(getter(p)))
+end
+
 """
     $(TYPEDSIGNATURES)
 

test/initializationsystem.jl

@@ -1513,7 +1513,7 @@ end
     @inferred solve(prob)
 end
 
-@testset "Issue#3570: `Initial`s are copied to `u0` if `u0` not provided to `remake`" begin
+@testset "Issue#3570, #3552: `Initial`s are copied to `u0` during `solve`/`init`" begin
     @parameters g
     @variables x(t) [state_priority = 10] y(t) λ(t)
     eqs = [D(D(x)) ~ λ * x
@@ -1525,11 +1525,22 @@ end
         pend, [x => (√2 / 2)], (0.0, 1.5), [g => 1], guesses = [λ => 1, y => √2 / 2])
     sol = solve(prob)
 
-    setter = setsym_oop(prob, [Initial(x)])
-    (u0, p) = setter(prob, [0.8])
+    @testset "`setsym_oop`" begin
+        setter = setsym_oop(prob, [Initial(x)])
+        (u0, p) = setter(prob, [0.8])
+        new_prob = remake(prob; u0, p, initializealg = BrownFullBasicInit())
+        new_sol = solve(new_prob)
+        @test new_sol[x, 1] ≈ 0.8
+        integ = init(new_prob)
+        @test integ[x] ≈ 0.8
+    end
 
-    new_prob = remake(prob; p, initializealg = BrownFullBasicInit())
-    @test new_prob[x] ≈ 0.8
-    new_sol = solve(new_prob)
-    @test new_sol[x, 1] ≈ 0.8
+    @testset "`setsym`" begin
+        @test prob.ps[Initial(x)] ≈ √2 / 2
+        prob.ps[Initial(x)] = 0.8
+        sol = solve(prob; initializealg = BrownFullBasicInit())
+        @test sol[x, 1] ≈ 0.8
+        integ = init(prob; initializealg = BrownFullBasicInit())
+        @test integ[x] ≈ 0.8
+    end
 end