feat: allow manually choosing time-independent initialization

Author: AayushSabharwal

src/systems/diffeqs/abstractodesystem.jl

@@ -1378,6 +1378,7 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
         allow_incomplete = false,
         force_time_independent = false,
         algebraic_only = false,
+        time_dependent_init = is_time_dependent(sys),
         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`")
@@ -1392,7 +1393,7 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
         simplify_system = true
     else
         isys = generate_initializesystem(
-            sys; u0map, initialization_eqs, check_units,
+            sys; u0map, initialization_eqs, check_units, time_dependent_init,
             pmap = parammap, guesses, extra_metadata = (; use_scc), algebraic_only)
         simplify_system = true
     end
@@ -1431,7 +1432,7 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
 
     # TODO: throw on uninitialized arrays
     filter!(x -> !(x isa Symbolics.Arr), uninit)
-    if is_time_dependent(sys) && !isempty(uninit)
+    if time_dependent_init && !isempty(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

src/systems/nonlinear/initializesystem.jl

@@ -1,5 +1,6 @@
-function generate_initializesystem(sys::AbstractSystem; kwargs...)
-    if is_time_dependent(sys)
+function generate_initializesystem(
+        sys::AbstractSystem; time_dependent_init = is_time_dependent(sys), kwargs...)
+    if time_dependent_init
         generate_initializesystem_timevarying(sys; kwargs...)
     else
         generate_initializesystem_timeindependent(sys; kwargs...)
@@ -153,7 +154,7 @@ function generate_initializesystem_timevarying(sys::AbstractSystem;
     end
     meta = InitializationSystemMetadata(
         anydict(u0map), anydict(pmap), additional_guesses,
-        additional_initialization_eqs, extra_metadata, nothing)
+        additional_initialization_eqs, extra_metadata, nothing, true)
     return NonlinearSystem(eqs_ics,
         vars,
         pars;
@@ -254,7 +255,7 @@ function generate_initializesystem_timeindependent(sys::AbstractSystem;
     end
     meta = InitializationSystemMetadata(
         anydict(u0map), anydict(pmap), additional_guesses,
-        additional_initialization_eqs, extra_metadata, nothing)
+        additional_initialization_eqs, extra_metadata, nothing, false)
     return NonlinearSystem(eqs_ics,
         vars,
         pars;
@@ -500,6 +501,7 @@ struct InitializationSystemMetadata
     additional_initialization_eqs::Vector{Equation}
     extra_metadata::NamedTuple
     oop_reconstruct_u0_p::Union{Nothing, ReconstructInitializeprob}
+    time_dependent_init::Bool
 end
 
 function get_possibly_array_fallback_singletons(varmap, p)
@@ -609,6 +611,7 @@ function SciMLBase.remake_initialization_data(
             merge!(guesses, meta.additional_guesses)
             use_scc = get(meta.extra_metadata, :use_scc, true)
             initialization_eqs = meta.additional_initialization_eqs
+            time_dependent_init = meta.time_dependent_init
         end
     else
         # there is no initializeprob, so the original problem construction
@@ -656,7 +659,7 @@ function SciMLBase.remake_initialization_data(
         u0map, pmap, defs, cmap, dvs, ps)
     floatT = float_type_from_varmap(op)
     kws = maybe_build_initialization_problem(
-        sys, op, u0map, pmap, t0, defs, guesses, missing_unknowns;
+        sys, op, u0map, pmap, t0, defs, guesses, missing_unknowns; time_dependent_init,
         use_scc, initialization_eqs, floatT, allow_incomplete = true)
 
     return SciMLBase.remake_initialization_data(sys, kws, newu0, t0, newp, newu0, newp)

src/systems/problem_utils.jl

@@ -633,15 +633,16 @@ All other keyword arguments are forwarded to `InitializationProblem`.
 function maybe_build_initialization_problem(
         sys::AbstractSystem, op::AbstractDict, u0map, pmap, t, defs,
         guesses, missing_unknowns; implicit_dae = false,
-        u0_constructor = identity, floatT = Float64, kwargs...)
+        time_dependent_init = is_time_dependent(sys), u0_constructor = identity,
+        floatT = Float64, kwargs...)
     guesses = merge(ModelingToolkit.guesses(sys), todict(guesses))
 
     if t === nothing && is_time_dependent(sys)
         t = zero(floatT)
     end
 
     initializeprob = ModelingToolkit.InitializationProblem{true, SciMLBase.FullSpecialize}(
-        sys, t, u0map, pmap; guesses, kwargs...)
+        sys, t, u0map, pmap; guesses, time_dependent_init, kwargs...)
     if state_values(initializeprob) !== nothing
         initializeprob = remake(initializeprob; u0 = floatT.(state_values(initializeprob)))
     end
@@ -660,7 +661,10 @@ function maybe_build_initialization_problem(
 
     meta = get_metadata(initializeprob.f.sys)
 
-    if is_time_dependent(sys)
+    if time_dependent_init === nothing
+        time_dependent_init = is_time_dependent(sys)
+    end
+    if time_dependent_init
         all_init_syms = Set(all_symbols(initializeprob))
         solved_unknowns = filter(var -> var in all_init_syms, unknowns(sys))
         initializeprobmap = u0_constructor ∘ getu(initializeprob, solved_unknowns)
@@ -700,7 +704,7 @@ function maybe_build_initialization_problem(
         end
     end
 
-    if is_time_dependent(sys)
+    if time_dependent_init
         for v in missing_unknowns
             op[v] = get_temporary_value(v, floatT)
         end
@@ -803,7 +807,7 @@ function process_SciMLProblem(
         symbolic_u0 = false, warn_cyclic_dependency = false,
         circular_dependency_max_cycle_length = length(all_symbols(sys)),
         circular_dependency_max_cycles = 10,
-        substitution_limit = 100, use_scc = true,
+        substitution_limit = 100, use_scc = true, time_dependent_init = is_time_dependent(sys),
         force_initialization_time_independent = false, algebraic_only = false,
         allow_incomplete = false, is_initializeprob = false, kwargs...)
     dvs = unknowns(sys)
@@ -858,7 +862,7 @@ function process_SciMLProblem(
             warn_cyclic_dependency, check_units = check_initialization_units,
             circular_dependency_max_cycle_length, circular_dependency_max_cycles, use_scc,
             force_time_independent = force_initialization_time_independent, algebraic_only, allow_incomplete,
-            u0_constructor, floatT)
+            u0_constructor, floatT, time_dependent_init)
 
         kwargs = merge(kwargs, kws)
     end