Merge pull request #3301 from AayushSabharwal/as/incomplete-init

feat: allow remade initialization systems to be incomplete

Author: ChrisRackauckas

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))

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
 

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)

test/initializationsystem.jl

@@ -32,7 +32,7 @@ initprob = ModelingToolkit.InitializationProblem(pend, 0.0, [x => 1, y => 0], [g
 @test initprob isa NonlinearLeastSquaresProblem
 sol = solve(initprob)
 @test SciMLBase.successful_retcode(sol)
-@test sol.u == [0.0, 0.0, 0.0, 0.0]
+@test all(iszero, sol.u)
 @test maximum(abs.(sol[conditions])) < 1e-14
 
 initprob = ModelingToolkit.InitializationProblem(
@@ -116,7 +116,7 @@ end
         x′
     end
     @variables begin
-        p(t) = p′
+        p(t)
         x(t) = x′
         dm(t) = 0
         f(t) = p′ * A
@@ -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