From 761f5ea75a39e4264c765d6d00b37b302581cf47 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Thu, 9 Jan 2025 13:01:27 +0530
Subject: [PATCH 1/2] feat: allow remade initialization systems to be
 incomplete

---
 src/systems/diffeqs/abstractodesystem.jl  |  9 ++++++++-
 src/systems/nonlinear/initializesystem.jl |  9 ++++++++-
 src/systems/problem_utils.jl              |  5 ++++-
 test/initializationsystem.jl              | 24 +++++++++++++++++++++++
 4 files changed, 44 insertions(+), 3 deletions(-)

diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
index b012bc73d4..3da5b6a2d0 100644
--- a/src/systems/diffeqs/abstractodesystem.jl
+++ b/src/systems/diffeqs/abstractodesystem.jl
@@ -1294,6 +1294,7 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
         fully_determined = nothing,
         check_units = true,
         use_scc = true,
+        allow_incomplete = false,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEProblem`")
@@ -1335,7 +1336,13 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
     # TODO: throw on uninitialized arrays
     filter!(x -> !(x isa Symbolics.Arr), uninit)
     if !isempty(uninit)
-        throw(IncompleteInitializationError(uninit))
+        allow_incomplete || throw(IncompleteInitializationError(uninit))
+        # for incomplete initialization, we will add the missing variables as parameters.
+        # they will be updated by `update_initializeprob!` and `initializeprobmap` will
+        # use them to construct the new `u0`.
+        newparams = map(toparam, uninit)
+        append!(get_ps(isys), newparams)
+        isys = complete(isys)
     end
 
     neqs = length(equations(isys))
diff --git a/src/systems/nonlinear/initializesystem.jl b/src/systems/nonlinear/initializesystem.jl
index 602fc7cac9..c3bbd92317 100644
--- a/src/systems/nonlinear/initializesystem.jl
+++ b/src/systems/nonlinear/initializesystem.jl
@@ -317,6 +317,12 @@ function SciMLBase.remake_initialization_data(
                 u0map[dvs[i]] = newu0[i]
             end
         end
+        # ensure all unknowns have guesses in case they weren't given one
+        # and become solvable
+        for i in eachindex(dvs)
+            haskey(guesses, dvs[i]) && continue
+            guesses[dvs[i]] = newu0[i]
+        end
         if p === missing
             # the user didn't pass `p` to `remake`, so they want to retain
             # existing values. Fill all parameters in `pmap` so that none of
@@ -341,7 +347,8 @@ function SciMLBase.remake_initialization_data(
     op, missing_unknowns, missing_pars = build_operating_point(
         u0map, pmap, defs, cmap, dvs, ps)
     kws = maybe_build_initialization_problem(
-        sys, op, u0map, pmap, t0, defs, guesses, missing_unknowns; use_scc, initialization_eqs)
+        sys, op, u0map, pmap, t0, defs, guesses, missing_unknowns;
+        use_scc, initialization_eqs, allow_incomplete = true)
     return get(kws, :initialization_data, nothing)
 end
 
diff --git a/src/systems/problem_utils.jl b/src/systems/problem_utils.jl
index 6ce5704daa..4cbf495d87 100644
--- a/src/systems/problem_utils.jl
+++ b/src/systems/problem_utils.jl
@@ -545,7 +545,10 @@ function maybe_build_initialization_problem(
        (!is_time_dependent(sys) || t !== nothing)
         initializeprob = ModelingToolkit.InitializationProblem(
             sys, t, u0map, pmap; guesses, kwargs...)
-        initializeprobmap = getu(initializeprob, unknowns(sys))
+
+        all_init_syms = Set(all_symbols(initializeprob))
+        solved_unknowns = filter(var -> var in all_init_syms, unknowns(sys))
+        initializeprobmap = getu(initializeprob, solved_unknowns)
 
         punknowns = [p
                      for p in all_variable_symbols(initializeprob)
diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl
index 2745b9d81f..ecf6899624 100644
--- a/test/initializationsystem.jl
+++ b/test/initializationsystem.jl
@@ -1282,3 +1282,27 @@ end
 
     @test SciMLBase.successful_retcode(solve(newprob))
 end
+
+@testset "Issue#3295: Incomplete initialization of pure-ODE systems" begin
+    @variables X(t) Y(t)
+    @parameters p d
+    eqs = [
+        D(X) ~ p - d * X,
+        D(Y) ~ p - d * Y
+    ]
+    @mtkbuild osys = ODESystem(eqs, t)
+
+    # Make problem.
+    u0_vals = [X => 4, Y => 5.0]
+    tspan = (0.0, 10.0)
+    p_vals = [p => 1.0, d => 0.1]
+    oprob = ODEProblem(osys, u0_vals, tspan, p_vals)
+    integ = init(oprob)
+    @test integ[X] ≈ 4.0
+    @test integ[Y] ≈ 5.0
+    # Attempt to `remake`.
+    rp = remake(oprob; u0 = [Y => 7])
+    integ = init(rp)
+    @test integ[X] ≈ 4.0
+    @test integ[Y] ≈ 7.0
+end

From fa0df9d7e215376967acc2f00e06feaea7e8d730 Mon Sep 17 00:00:00 2001
From: Aayush Sabharwal <aayush.sabharwal@gmail.com>
Date: Fri, 10 Jan 2025 14:22:38 +0530
Subject: [PATCH 2/2] test: fix initialization test

---
 test/initializationsystem.jl | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/test/initializationsystem.jl b/test/initializationsystem.jl
index ecf6899624..6f8c961b45 100644
--- a/test/initializationsystem.jl
+++ b/test/initializationsystem.jl
@@ -116,7 +116,7 @@ end
         x′
     end
     @variables begin
-        p(t) = p′
+        p(t)
         x(t) = x′
         dm(t) = 0
         f(t) = p′ * A