feat: always build initialization problem

Project.toml

@@ -132,7 +132,7 @@ RecursiveArrayTools = "3.26"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
 SCCNonlinearSolve = "1.0.0"
-SciMLBase = "2.71.1"
+SciMLBase = "2.72"
 SciMLStructures = "1.0"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"

src/ModelingToolkit.jl

@@ -45,7 +45,7 @@ using Compat
 using AbstractTrees
 using DiffEqBase, SciMLBase, ForwardDiff
 using SciMLBase: StandardODEProblem, StandardNonlinearProblem, handle_varmap, TimeDomain,
-                 PeriodicClock, Clock, SolverStepClock, Continuous
+                 PeriodicClock, Clock, SolverStepClock, Continuous, OverrideInit, NoInit
 using Distributed
 import JuliaFormatter
 using MLStyle

src/systems/abstractsystem.jl

@@ -2377,6 +2377,9 @@ See also [`linearize`](@ref) which provides a higher-level interface.
 function linearization_function(sys::AbstractSystem, inputs,
         outputs; simplify = false,
         initialize = true,
+        initializealg = nothing,
+        initialization_abstol = 1e-6,
+        initialization_reltol = 1e-3,
         op = Dict(),
         p = DiffEqBase.NullParameters(),
         zero_dummy_der = false,
@@ -2403,88 +2406,29 @@ function linearization_function(sys::AbstractSystem, inputs,
         op = merge(defs, op)
     end
     sys = ssys
-    u0map = Dict(k => v for (k, v) in op if is_variable(ssys, k))
-    initsys = structural_simplify(
-        generate_initializesystem(
-            sys, u0map = u0map, guesses = guesses(sys), algebraic_only = true),
-        fully_determined = false)
-
-    # HACK: some unknowns may not be involved in any initialization equations, and are
-    # thus removed from the system during `structural_simplify`.
-    # This causes `getu(initsys, unknowns(sys))` to fail, so we add them back as parameters
-    # for now.
-    missing_unknowns = setdiff(unknowns(sys), all_symbols(initsys))
-    if !isempty(missing_unknowns)
-        if warn_initialize_determined
-            @warn "Initialization system is underdetermined. No equations for $(missing_unknowns). Initialization will default to using least squares. To suppress this warning pass warn_initialize_determined = false."
-        end
-        new_parameters = [parameters(initsys); missing_unknowns]
-        @set! initsys.ps = new_parameters
-        initsys = complete(initsys)
-    end
-
-    if p isa SciMLBase.NullParameters
-        p = Dict()
-    else
-        p = todict(p)
-    end
-    x0 = merge(defaults_and_guesses(sys), op)
-    if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        sys_ps = MTKParameters(sys, p, x0)
-    else
-        sys_ps = varmap_to_vars(p, parameters(sys); defaults = x0)
-    end
-    p[get_iv(sys)] = NaN
-    if has_index_cache(initsys) && get_index_cache(initsys) !== nothing
-        oldps = MTKParameters(initsys, p, merge(guesses(sys), defaults(sys), op))
-        initsys_ps = parameters(initsys)
-        p_getter = build_explicit_observed_function(
-            sys, initsys_ps; eval_expression, eval_module)
-
-        u_getter = isempty(unknowns(initsys)) ? (_...) -> nothing :
-                   build_explicit_observed_function(
-            sys, unknowns(initsys); eval_expression, eval_module)
-        get_initprob_u_p = let p_getter = p_getter,
-            p_setter! = setp(initsys, initsys_ps),
-            u_getter = u_getter
-
-            function (u, p, t)
-                p_setter!(oldps, p_getter(u, p, t))
-                newu = u_getter(u, p, t)
-                return newu, oldps
-            end
-        end
-    else
-        get_initprob_u_p = let p_getter = getu(sys, parameters(initsys)),
-            u_getter = build_explicit_observed_function(
-                sys, unknowns(initsys); eval_expression, eval_module)
-
-            function (u, p, t)
-                state = ProblemState(; u, p, t)
-                return u_getter(
-                    state_values(state), parameter_values(state), current_time(state)),
-                p_getter(state)
-            end
-        end
+
+    if initializealg === nothing
+        initializealg = initialize ? OverrideInit() : NoInit()
     end
-    initfn = NonlinearFunction(initsys; eval_expression, eval_module)
-    initprobmap = build_explicit_observed_function(
-        initsys, unknowns(sys); eval_expression, eval_module)
+
+    fun, u0, p = process_SciMLProblem(ODEFunction{true, SciMLBase.FullSpecialize}, sys, op, p; t = 0.0, build_initializeprob = initializealg isa OverrideInit, allow_incomplete = true, algebraic_only = true)
+    prob = ODEProblem(fun, u0, (nothing, nothing), p)
+
     ps = parameters(sys)
     h = build_explicit_observed_function(sys, outputs; eval_expression, eval_module)
     lin_fun = let diff_idxs = diff_idxs,
         alge_idxs = alge_idxs,
         input_idxs = input_idxs,
         sts = unknowns(sys),
-        get_initprob_u_p = get_initprob_u_p,
-        fun = ODEFunction{true, SciMLBase.FullSpecialize}(
-            sys, unknowns(sys), ps; eval_expression, eval_module),
-        initfn = initfn,
-        initprobmap = initprobmap,
+        fun = fun,
+        prob = prob,
+        sys_ps = p,
         h = h,
+        integ_cache = (similar(u0)),
         chunk = ForwardDiff.Chunk(input_idxs),
-        sys_ps = sys_ps,
-        initialize = initialize,
+        initializealg = initializealg,
+        initialization_abstol = initialization_abstol,
+        initialization_reltol = initialization_reltol,
         initialization_solver_alg = initialization_solver_alg,
         sys = sys
 
@@ -2504,14 +2448,11 @@ function linearization_function(sys::AbstractSystem, inputs,
             if u !== nothing # Handle systems without unknowns
                 length(sts) == length(u) ||
                     error("Number of unknown variables ($(length(sts))) does not match the number of input unknowns ($(length(u)))")
-                if initialize && !isempty(alge_idxs) # This is expensive and can be omitted if the user knows that the system is already initialized
-                    residual = fun(u, p, t)
-                    if norm(residual[alge_idxs]) > √(eps(eltype(residual)))
-                        initu0, initp = get_initprob_u_p(u, p, t)
-                        initprob = NonlinearLeastSquaresProblem(initfn, initu0, initp)
-                        nlsol = solve(initprob, initialization_solver_alg)
-                        u = initprobmap(state_values(nlsol), parameter_values(nlsol))
-                    end
+
+                integ = MockIntegrator{true}(u, p, t, integ_cache)
+                u, p, success = SciMLBase.get_initial_values(prob, integ, fun, initializealg, Val(true); abstol = initialization_abstol, reltol = initialization_reltol, nlsolve_alg = initialization_solver_alg)
+                if !success
+                    error("Initialization algorithm $(initializealg) failed with `u = $u` and `p = $p`.")
                 end
                 uf = SciMLBase.UJacobianWrapper(fun, t, p)
                 fg_xz = ForwardDiff.jacobian(uf, u)
@@ -2546,6 +2487,44 @@ function linearization_function(sys::AbstractSystem, inputs,
     return lin_fun, sys
 end
 
+"""
+    $(TYPEDEF)
+
+Mock `DEIntegrator` to allow using `CheckInit` without having to create a new integrator
+(and consequently depend on `OrdinaryDiffEq`).
+
+# Fields
+
+$(TYPEDFIELDS)
+"""
+struct MockIntegrator{iip, U, P, T, C} <: SciMLBase.DEIntegrator{Nothing, iip, U, T}
+    """
+    The state vector.
+    """
+    u::U
+    """
+    The parameter object.
+    """
+    p::P
+    """
+    The current time.
+    """
+    t::T
+    """
+    The integrator cache.
+    """
+    cache::C
+end
+
+function MockIntegrator{iip}(u::U, p::P, t::T, cache::C) where {iip, U, P, T, C}
+    return MockIntegrator{iip, U, P, T, C}(u, p, t, cache)
+end
+
+SymbolicIndexingInterface.state_values(integ::MockIntegrator) = integ.u
+SymbolicIndexingInterface.parameter_values(integ::MockIntegrator) = integ.p
+SymbolicIndexingInterface.current_time(integ::MockIntegrator) = integ.t
+SciMLBase.get_tmp_cache(integ::MockIntegrator) = integ.cache
+
 """
     (; A, B, C, D), simplified_sys = linearize_symbolic(sys::AbstractSystem, inputs, outputs; simplify = false, allow_input_derivatives = false, kwargs...)
 

src/systems/analysis_points.jl

@@ -510,7 +510,7 @@
         ap_idx = analysis_point_index(ap_sys, tf.ap)
         ap_idx === nothing &&
             error("Analysis point $(nameof(tf.ap)) not found in system $(nameof(sys)).")
         # get the anlysis point
         ap_sys_eqs = copy(get_eqs(ap_sys))
         ap = ap_sys_eqs[ap_idx].rhs
 
@@ -564,7 +564,7 @@
         ap_idx = analysis_point_index(ap_sys, tf.ap)
         ap_idx === nothing &&
             error("Analysis point $(nameof(tf.ap)) not found in system $(nameof(sys)).")
         # modified quations
         ap_sys_eqs = copy(get_eqs(ap_sys))
         @set! ap_sys.eqs = ap_sys_eqs
         ap = ap_sys_eqs[ap_idx].rhs
@@ -863,7 +863,7 @@
             sys, ap, args...; loop_openings = [], system_modifier = identity, kwargs...)
         lin_fun, ssys = $(utility_fun)(
             sys, ap, args...; loop_openings, system_modifier, kwargs...)
-        ModelingToolkit.linearize(ssys, lin_fun; kwargs...), ssys
+        ModelingToolkit.linearize(ssys, lin_fun), ssys
     end
 end
 
@@ -876,7 +876,7 @@
 # Keyword Arguments
 
 - `system_modifier`: a function which takes the modified system and returns a new system
   with any required further modifications peformed.
 """
 function open_loop(sys, ap::Union{Symbol, AnalysisPoint}; system_modifier = identity)
     ap = only(canonicalize_ap(sys, ap))

src/systems/diffeqs/abstractodesystem.jl

@@ -1120,7 +1120,7 @@ function DiffEqBase.SteadyStateProblem{iip}(sys::AbstractODESystem, u0map,
     end
     f, u0, p = process_SciMLProblem(ODEFunction{iip}, sys, u0map, parammap;
         steady_state = true,
-        check_length, kwargs...)
+        check_length, force_initialization_time_independent = true, kwargs...)
     kwargs = filter_kwargs(kwargs)
     SteadyStateProblem{iip}(f, u0, p; kwargs...)
 end
@@ -1295,27 +1295,41 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
         check_units = true,
         use_scc = true,
         allow_incomplete = false,
+        force_time_independent = false,
+        algebraic_only = 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`")
     end
     if isempty(u0map) && get_initializesystem(sys) !== nothing
         isys = get_initializesystem(sys; initialization_eqs, check_units)
+        simplify_system = false
     elseif isempty(u0map) && get_initializesystem(sys) === nothing
-        isys = structural_simplify(
-            generate_initializesystem(
+        isys = generate_initializesystem(
                 sys; initialization_eqs, check_units, pmap = parammap,
-                guesses, extra_metadata = (; use_scc)); fully_determined)
+                guesses, extra_metadata = (; use_scc), algebraic_only)
+        simplify_system = true
     else
-        isys = structural_simplify(
-            generate_initializesystem(
+        isys = generate_initializesystem(
                 sys; u0map, initialization_eqs, check_units,
-                pmap = parammap, guesses, extra_metadata = (; use_scc)); fully_determined)
+                pmap = parammap, guesses, extra_metadata = (; use_scc), algebraic_only)
+        simplify_system = true
+    end
+
+    # useful for `SteadyStateProblem` since `f` has to be autonomous and the
+    # initialization should be too
+    if force_time_independent
+        idx = findfirst(isequal(get_iv(sys)), get_ps(isys))
+        idx === nothing || deleteat!(get_ps(isys), idx)
+    end
+
+    if simplify_system
+        isys = structural_simplify(isys; fully_determined)
     end
 
     meta = get_metadata(isys)
     if meta isa InitializationSystemMetadata
-        @set! isys.metadata.oop_reconstruct_u0_p = ReconstructInitializeprob(sys, isys)
+        @set! isys.metadata.oop_reconstruct_u0_p = ReconstructInitializeprob(sys, isys; remap = meta.new_params)
     end
 
     ts = get_tearing_state(isys)
@@ -1374,20 +1388,6 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
 
     u0map = merge(ModelingToolkit.guesses(sys), todict(guesses), todict(u0map))
 
-    # Replace dummy derivatives in u0map: D(x) -> x_t etc.
-    if has_schedule(sys)
-        schedule = get_schedule(sys)
-        if !isnothing(schedule)
-            for (var, val) in u0map
-                dvar = get(schedule.dummy_sub, var, var) # with dummy derivatives
-                if dvar !== var # then replace it
-                    delete!(u0map, var)
-                    push!(u0map, dvar => val)
-                end
-            end
-        end
-    end
-
     fullmap = merge(u0map, parammap)
     u0T = Union{}
     for sym in unknowns(isys)

src/systems/diffeqs/odesystem.jl

@@ -687,13 +687,17 @@ function populate_delays(delays::Set, obsexprs, histfn, sys, sym)
 end
 
 function _eq_unordered(a, b)
+    a = vec(a)
+    b = vec(b)
     length(a) === length(b) || return false
     n = length(a)
     idxs = Set(1:n)
     for x in a
         idx = findfirst(isequal(x), b)
+        while idx !== nothing && !(idx in idxs)
+            idx = findnext(isequal(x), b, idx + 1)
+        end
         idx === nothing && return false
-        idx ∈ idxs || return false
         delete!(idxs, idx)
     end
     return true

src/systems/discrete_system/discrete_system.jl

@@ -309,7 +309,7 @@ function SciMLBase.DiscreteProblem(
     u0map = to_varmap(u0map, dvs)
     u0map = shift_u0map_forward(sys, u0map, defaults(sys))
     f, u0, p = process_SciMLProblem(
-        DiscreteFunction, sys, u0map, parammap; eval_expression, eval_module)
+        DiscreteFunction, sys, u0map, parammap; eval_expression, eval_module, build_initializeprob = false)
     u0 = f(u0, p, tspan[1])
     DiscreteProblem(f, u0, tspan, p; kwargs...)
 end
@@ -336,6 +336,7 @@ function SciMLBase.DiscreteFunction{iip, specialize}(
         eval_expression = false,
         eval_module = @__MODULE__,
         analytic = nothing,
+        initialization_data = nothing,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed `DiscreteSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `DiscreteProblem`")
@@ -359,7 +360,8 @@ function SciMLBase.DiscreteFunction{iip, specialize}(
     DiscreteFunction{iip, specialize}(f;
         sys = sys,
         observed = observedfun,
-        analytic = analytic)
+        analytic = analytic,
+        initialization_data = initialization_data)
 end
 
 """

src/systems/jumps/jumpsystem.jl

@@ -426,7 +426,7 @@ function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple,
     end
 
     _f, u0, p = process_SciMLProblem(EmptySciMLFunction, sys, u0map, parammap;
-        t = tspan === nothing ? nothing : tspan[1], use_union, tofloat = false, check_length = false)
+        t = tspan === nothing ? nothing : tspan[1], use_union, tofloat = false, check_length = false, build_initializeprob = false)
     f = DiffEqBase.DISCRETE_INPLACE_DEFAULT
 
     observedfun = ObservedFunctionCache(
@@ -523,7 +523,7 @@ function DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothi
             parameter_dependencies = parameter_dependencies(sys),
             metadata = get_metadata(sys), gui_metadata = get_gui_metadata(sys))
         osys = complete(osys)
-        return ODEProblem(osys, u0map, tspan, parammap; check_length = false, kwargs...)
+        return ODEProblem(osys, u0map, tspan, parammap; check_length = false, build_initializeprob = false, kwargs...)
     else
         _, u0, p = process_SciMLProblem(EmptySciMLFunction, sys, u0map, parammap;
             t = tspan === nothing ? nothing : tspan[1], use_union, tofloat = false,