Fix up initializesystem for hierarchical models

Project.toml

@@ -89,12 +89,12 @@ Libdl = "1"
 LinearAlgebra = "1"
 MLStyle = "0.4.17"
 NaNMath = "0.3, 1"
-OrdinaryDiffEq = "6"
+OrdinaryDiffEq = "6.72.0"
 PrecompileTools = "1"
 RecursiveArrayTools = "2.3, 3"
 Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
-SciMLBase = "2.0.1"
+SciMLBase = "2.28.0"
 SciMLStructures = "1.0"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"

src/ModelingToolkit.jl

@@ -137,19 +137,18 @@ include("systems/model_parsing.jl")
 include("systems/connectors.jl")
 include("systems/callbacks.jl")
 
+include("systems/nonlinear/nonlinearsystem.jl")
 include("systems/diffeqs/odesystem.jl")
 include("systems/diffeqs/sdesystem.jl")
 include("systems/diffeqs/abstractodesystem.jl")
+include("systems/nonlinear/modelingtoolkitize.jl")
+include("systems/nonlinear/initializesystem.jl")
 include("systems/diffeqs/first_order_transform.jl")
 include("systems/diffeqs/modelingtoolkitize.jl")
 include("systems/diffeqs/basic_transformations.jl")
 
 include("systems/jumps/jumpsystem.jl")
 
-include("systems/nonlinear/nonlinearsystem.jl")
-include("systems/nonlinear/modelingtoolkitize.jl")
-include("systems/nonlinear/initializesystem.jl")
-
 include("systems/optimization/constraints_system.jl")
 include("systems/optimization/optimizationsystem.jl")
 include("systems/optimization/modelingtoolkitize.jl")
@@ -253,7 +252,7 @@ export toexpr, get_variables
 export simplify, substitute
 export build_function
 export modelingtoolkitize
-export initializesystem
+export initializesystem, generate_initializesystem
 
 export @variables, @parameters, @constants, @brownian
 export @named, @nonamespace, @namespace, extend, compose, complete

src/bipartite_graph.jl

@@ -423,7 +423,7 @@ return `false` may not be matched.
 """
 function maximal_matching(g::BipartiteGraph, srcfilter = vsrc -> true,
         dstfilter = vdst -> true, ::Type{U} = Unassigned) where {U}
-    matching = Matching{U}(ndsts(g))
+    matching = Matching{U}(max(nsrcs(g), ndsts(g)))
     foreach(Iterators.filter(srcfilter, 𝑠vertices(g))) do vsrc
         construct_augmenting_path!(matching, g, vsrc, dstfilter)
     end

src/systems/abstractsystem.jl

@@ -537,6 +537,9 @@
     if split && has_index_cache(sys)
         @set! sys.index_cache = IndexCache(sys)
     end
+    if isdefined(sys, :initializesystem) && get_initializesystem(sys) !== nothing
+        @set! sys.initializesystem = complete(get_initializesystem(sys); split)
+    end
     isdefined(sys, :complete) ? (@set! sys.complete = true) : sys
 end
 
@@ -551,6 +554,7 @@
              :var_to_name
              :ctrls
              :defaults
+             :guesses
              :observed
              :tgrad
              :jac
@@ -571,6 +575,7 @@
              :connections
              :preface
              :torn_matching
+             :initializesystem
              :tearing_state
              :substitutions
              :metadata
@@ -933,6 +938,10 @@
     vcat(parameters(sys), dependent_parameters(sys))
 end
 
+function guesses(sys::AbstractSystem)
+    get_guesses(sys)
+end
+
 # required in `src/connectors.jl:437`
 parameters(_) = []
 
@@ -2259,14 +2268,15 @@
 end
 
 """
-    missing_variable_defaults(sys::AbstractSystem, default = 0.0)
+    missing_variable_defaults(sys::AbstractSystem, default = 0.0; subset = unknowns(sys))
 
 returns a `Vector{Pair}` of variables set to `default` which are missing from `get_defaults(sys)`.  The `default` argument can be a single value or vector to set the missing defaults respectively.
 """
-function missing_variable_defaults(sys::AbstractSystem, default = 0.0)
+function missing_variable_defaults(
+        sys::AbstractSystem, default = 0.0; subset = unknowns(sys))
     varmap = get_defaults(sys)
     varmap = Dict(Symbolics.diff2term(value(k)) => value(varmap[k]) for k in keys(varmap))
-    missingvars = setdiff(unknowns(sys), keys(varmap))
+    missingvars = setdiff(subset, keys(varmap))
     ds = Pair[]
 
     n = length(missingvars)

src/systems/diffeqs/abstractodesystem.jl

@@ -316,6 +316,8 @@
         sparsity = false,
         analytic = nothing,
         split_idxs = nothing,
+        initializeprob = nothing,
+        initializeprobmap = nothing,
         kwargs...) where {iip, specialize}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating an `ODEFunction`")
@@ -487,6 +489,7 @@
     end
 
     @set! sys.split_idxs = split_idxs
+
     ODEFunction{iip, specialize}(f;
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
@@ -495,7 +498,9 @@
         jac_prototype = jac_prototype,
         observed = observedfun,
         sparsity = sparsity ? jacobian_sparsity(sys) : nothing,
-        analytic = analytic)
+        analytic = analytic,
+        initializeprob = initializeprob,
+        initializeprobmap = initializeprobmap)
 end
 
 """
@@ -525,6 +530,8 @@
         sparse = false, simplify = false,
         eval_module = @__MODULE__,
         checkbounds = false,
+        initializeprob = nothing,
+        initializeprobmap = nothing,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `DAEFunction`")
@@ -596,7 +603,9 @@
         sys = sys,
         jac = _jac === nothing ? nothing : _jac,
         jac_prototype = jac_prototype,
-        observed = observedfun)
+        observed = observedfun,
+        initializeprob = initializeprob,
+        initializeprobmap = initializeprobmap)
 end
 
 function DiffEqBase.DDEFunction(sys::AbstractODESystem, args...; kwargs...)
@@ -839,18 +848,45 @@
         tofloat = true,
         symbolic_u0 = false,
         u0_constructor = identity,
+        guesses = Dict(),
+        warn_initialize_determined = true,
         kwargs...)
     eqs = equations(sys)
     dvs = unknowns(sys)
     ps = full_parameters(sys)
     iv = get_iv(sys)
 
+    # Append zeros to the variables which are determined by the initialization system
+    # This essentially bypasses the check for if initial conditions are defined for DAEs
+    # since they will be checked in the initialization problem's construction
+    # TODO: make check for if a DAE cheaper than calculating the mass matrix a second time!
+    ci = infer_clocks!(ClockInference(TearingState(sys)))
+    # TODO: make it work with clocks
+    # ModelingToolkit.get_tearing_state(sys) !== nothing => Requires structural_simplify first
+    if (implicit_dae || calculate_massmatrix(sys) !== I) &&
+       all(isequal(Continuous()), ci.var_domain) &&
+       ModelingToolkit.get_tearing_state(sys) !== nothing
+        if eltype(u0map) <: Number
+            u0map = unknowns(sys) .=> u0map
+        end
+        initializeprob = ModelingToolkit.InitializationProblem(
+            sys, u0map, parammap; guesses, warn_initialize_determined)
+        initializeprobmap = getu(initializeprob, unknowns(sys))
+
+        zerovars = setdiff(unknowns(sys), keys(defaults(sys))) .=> 0.0
+        trueinit = identity.([zerovars; u0map])
+    else
+        initializeprob = nothing
+        initializeprobmap = nothing
+        trueinit = u0map
+    end
+
     if has_index_cache(sys) && get_index_cache(sys) !== nothing
-        u0, defs = get_u0(sys, u0map, parammap; symbolic_u0)
+        u0, defs = get_u0(sys, trueinit, parammap; symbolic_u0)
         p = MTKParameters(sys, parammap)
     else
         u0, p, defs = get_u0_p(sys,
-            u0map,
+            trueinit,
             parammap;
             tofloat,
             use_union,
@@ -881,6 +917,8 @@
         checkbounds = checkbounds, p = p,
         linenumbers = linenumbers, parallel = parallel, simplify = simplify,
         sparse = sparse, eval_expression = eval_expression,
+        initializeprob = initializeprob,
+        initializeprobmap = initializeprobmap,
         kwargs...)
     implicit_dae ? (f, du0, u0, p) : (f, u0, p)
 end
@@ -984,13 +1022,14 @@
         parammap = DiffEqBase.NullParameters();
         callback = nothing,
         check_length = true,
+        warn_initialize_determined = true,
         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
     f, u0, p = process_DEProblem(ODEFunction{iip, specialize}, sys, u0map, parammap;
         t = tspan !== nothing ? tspan[1] : tspan,
-        check_length, kwargs...)
+        check_length, warn_initialize_determined, kwargs...)
     cbs = process_events(sys; callback, kwargs...)
     inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
@@ -1055,13 +1094,14 @@
 
 function DiffEqBase.DAEProblem{iip}(sys::AbstractODESystem, du0map, u0map, tspan,
         parammap = DiffEqBase.NullParameters();
+        warn_initialize_determined = true,
         check_length = true, kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed system is required. Call `complete` or `structural_simplify` on the system before creating a `DAEProblem`")
     end
     f, du0, u0, p = process_DEProblem(DAEFunction{iip}, sys, u0map, parammap;
         implicit_dae = true, du0map = du0map, check_length,
-        kwargs...)
+        warn_initialize_determined, kwargs...)
     diffvars = collect_differential_variables(sys)
     sts = unknowns(sys)
     differential_vars = map(Base.Fix2(in, diffvars), sts)
@@ -1442,3 +1482,76 @@
         end
     end
 end
+
+struct InitializationProblem{iip, specialization} end
+
+"""
+```julia
+InitializationProblem{iip}(sys::AbstractODESystem, u0map, tspan,
+                           parammap = DiffEqBase.NullParameters();
+                           version = nothing, tgrad = false,
+                           jac = false,
+                           checkbounds = false, sparse = false,
+                           simplify = false,
+                           linenumbers = true, parallel = SerialForm(),
+                           kwargs...) where {iip}
+```
+
+Generates a NonlinearProblem or NonlinearLeastSquaresProblem from an ODESystem 
+which represents the initialization, i.e. the calculation of the consistent
+initial conditions for the given DAE.
+"""
+function InitializationProblem(sys::AbstractODESystem, args...; kwargs...)
+    InitializationProblem{true}(sys, args...; kwargs...)
+end
+
+function InitializationProblem(sys::AbstractODESystem,
+        u0map::StaticArray,
+        args...;
+        kwargs...)
+    InitializationProblem{false, SciMLBase.FullSpecialize}(sys, u0map, args...; kwargs...)
+end
+
+function InitializationProblem{true}(sys::AbstractODESystem, args...; kwargs...)
+    InitializationProblem{true, SciMLBase.AutoSpecialize}(sys, args...; kwargs...)
+end
+
+function InitializationProblem{false}(sys::AbstractODESystem, args...; kwargs...)
+    InitializationProblem{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
+end
+
+function InitializationProblem{iip, specialize}(sys::AbstractODESystem, u0map = [],
+        parammap = DiffEqBase.NullParameters();
+        guesses = [],
+        check_length = true,
+        warn_initialize_determined = true,
+        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)
+    elseif isempty(u0map) && get_initializesystem(sys) === nothing
+        isys = structural_simplify(generate_initializesystem(sys); fully_determined = false)
+    else
+        isys = structural_simplify(
+            generate_initializesystem(sys; u0map); fully_determined = false)
+    end
+
+    neqs = length(equations(isys))
+    nunknown = length(unknowns(isys))
+
+    if warn_initialize_determined && neqs > nunknown
+        @warn "Initialization system is overdetermined. $neqs equations for $nunknown unknowns. Initialization will default to using least squares. To suppress this warning pass warn_initialize_determined = false."
+    end
+    if warn_initialize_determined && neqs < nunknown
+        @warn "Initialization system is underdetermined. $neqs equations for $nunknown unknowns. Initialization will default to using least squares. To suppress this warning pass warn_initialize_determined = false."
+    end
+
+    if neqs == nunknown
+        NonlinearProblem(isys, guesses, parammap)
+    else
+        NonlinearLeastSquaresProblem(isys, guesses, parammap)
+    end
+end

src/systems/diffeqs/odesystem.jl

@@ -88,10 +88,19 @@ struct ODESystem <: AbstractODESystem
     """
     defaults::Dict
     """
+    The guesses to use as the initial conditions for the
+    initialization system.
+    """
+    guesses::Dict
+    """
     Tearing result specifying how to solve the system.
     """
     torn_matching::Union{Matching, Nothing}
     """
+    The system for performing the initialization.
+    """
+    initializesystem::Union{Nothing, NonlinearSystem}
+    """
     Type of the system.
     """
     connector_type::Any
@@ -157,8 +166,8 @@ struct ODESystem <: AbstractODESystem
     parent::Any
 
     function ODESystem(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad,
-            jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults,
-            torn_matching, connector_type, preface, cevents,
+            jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults, guesses,
+            torn_matching, initializesystem, connector_type, preface, cevents,
             devents, parameter_dependencies, metadata = nothing, gui_metadata = nothing,
             tearing_state = nothing,
             substitutions = nothing, complete = false, index_cache = nothing,
@@ -175,8 +184,8 @@ struct ODESystem <: AbstractODESystem
             check_units(u, deqs)
         end
         new(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
-            ctrl_jac, Wfact, Wfact_t, name, systems, defaults, torn_matching,
-            connector_type, preface, cevents, devents, parameter_dependencies, metadata,
+            ctrl_jac, Wfact, Wfact_t, name, systems, defaults, guesses, torn_matching,
+            initializesystem, connector_type, preface, cevents, devents, parameter_dependencies, metadata,
             gui_metadata, tearing_state, substitutions, complete, index_cache,
             discrete_subsystems, solved_unknowns, split_idxs, parent)
     end
@@ -191,6 +200,8 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         default_u0 = Dict(),
         default_p = Dict(),
         defaults = _merge(Dict(default_u0), Dict(default_p)),
+        guesses = Dict(),
+        initializesystem = nothing,
         connector_type = nothing,
         preface = nothing,
         continuous_events = nothing,
@@ -217,6 +228,13 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
     var_to_name = Dict()
     process_variables!(var_to_name, defaults, dvs′)
     process_variables!(var_to_name, defaults, ps′)
+
+    sysguesses = [ModelingToolkit.getguess(st) for st in dvs′]
+    hasaguess = findall(!isnothing, sysguesses)
+    var_guesses = dvs′[hasaguess] .=> sysguesses[hasaguess]
+    sysguesses = isempty(var_guesses) ? Dict() : todict(var_guesses)
+    guesses = merge(sysguesses, todict(guesses))
+
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
 
     tgrad = RefValue(EMPTY_TGRAD)
@@ -234,7 +252,7 @@ function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps;
         parameter_dependencies, ps′)
     ODESystem(Threads.atomic_add!(SYSTEM_COUNT, UInt(1)),
         deqs, iv′, dvs′, ps′, tspan, var_to_name, ctrl′, observed, tgrad, jac,
-        ctrl_jac, Wfact, Wfact_t, name, systems, defaults, nothing,
+        ctrl_jac, Wfact, Wfact_t, name, systems, defaults, guesses, nothing, initializesystem,
         connector_type, preface, cont_callbacks, disc_callbacks, parameter_dependencies,
         metadata, gui_metadata, checks = checks)
 end