Merge remote-tracking branch 'origin' into implicit_discrete_system

Update to master

Author: vyudu

Project.toml

@@ -1,7 +1,7 @@
 name = "ModelingToolkit"
 uuid = "961ee093-0014-501f-94e3-6117800e7a78"
 authors = ["Yingbo Ma <mayingbo5@gmail.com>", "Chris Rackauckas <accounts@chrisrackauckas.com> and contributors"]
-version = "9.64.3"
+version = "9.65.0"
 
 [deps]
 AbstractTrees = "1520ce14-60c1-5f80-bbc7-55ef81b5835c"
@@ -136,7 +136,7 @@ Reexport = "0.2, 1"
 RuntimeGeneratedFunctions = "0.5.9"
 SCCNonlinearSolve = "1.0.0"
 SciMLBase = "2.75"
-SciMLStructures = "1.0"
+SciMLStructures = "1.7"
 Serialization = "1"
 Setfield = "0.7, 0.8, 1"
 SimpleNonlinearSolve = "0.1.0, 1, 2"

docs/src/tutorials/initialization.md

@@ -536,3 +536,14 @@ sol[α * x - β * x * y]
 ```@example init
 plot(sol)
 ```
+
+## Summary of Initialization API
+
+```@docs; canonical=false
+Initial
+isinitial
+generate_initializesystem
+initialization_equations
+guesses
+defaults
+```

ext/MTKChainRulesCoreExt.jl

@@ -79,7 +79,11 @@ function ChainRulesCore.rrule(
     end
     newbuf = MTK.remake_buffer(indp, oldbuf, idxs, vals)
     tunable_idxs = reduce(
-        vcat, (idx.idx for idx in idxs if idx.portion isa MTK.SciMLStructures.Tunable))
+        vcat, (idx.idx for idx in idxs if idx.portion isa MTK.SciMLStructures.Tunable);
+        init = Union{Int, AbstractVector{Int}}[])
+    initials_idxs = reduce(
+        vcat, (idx.idx for idx in idxs if idx.portion isa MTK.SciMLStructures.Initials);
+        init = Union{Int, AbstractVector{Int}}[])
     disc_idxs = subset_idxs(idxs, MTK.SciMLStructures.Discrete(), oldbuf.discrete)
     const_idxs = subset_idxs(idxs, MTK.SciMLStructures.Constants(), oldbuf.constant)
     nn_idxs = subset_idxs(idxs, MTK.NONNUMERIC_PORTION, oldbuf.nonnumeric)
@@ -91,10 +95,12 @@ function ChainRulesCore.rrule(
             indp′ = NoTangent()
 
             tunable = selected_tangents(buf′.tunable, tunable_idxs)
+            initials = selected_tangents(buf′.initials, initials_idxs)
             discrete = selected_tangents(buf′.discrete, disc_idxs)
             constant = selected_tangents(buf′.constant, const_idxs)
             nonnumeric = selected_tangents(buf′.nonnumeric, nn_idxs)
-            oldbuf′ = Tangent{typeof(oldbuf)}(; tunable, discrete, constant, nonnumeric)
+            oldbuf′ = Tangent{typeof(oldbuf)}(;
+                tunable, initials, discrete, constant, nonnumeric)
             idxs′ = NoTangent()
             vals′ = map(i -> MTK._ducktyped_parameter_values(buf′, i), idxs)
             return f′, indp′, oldbuf′, idxs′, vals′

src/ModelingToolkit.jl

@@ -290,7 +290,7 @@ export toexpr, get_variables
 export simplify, substitute
 export build_function
 export modelingtoolkitize
-export generate_initializesystem, Initial
+export generate_initializesystem, Initial, isinitial
 
 export alg_equations, diff_equations, has_alg_equations, has_diff_equations
 export get_alg_eqs, get_diff_eqs, has_alg_eqs, has_diff_eqs

src/inputoutput.jl

@@ -250,6 +250,9 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
     f = build_function_wrapper(sys, rhss, args...; p_start = 3 + implicit_dae,
         p_end = length(p) + 2 + implicit_dae)
     f = eval_or_rgf.(f; eval_expression, eval_module)
+    f = GeneratedFunctionWrapper{(
+        3 + implicit_dae, length(args) - length(p) + 1, is_split(sys))}(f...)
+    f = f, f
     ps = setdiff(parameters(sys), inputs, disturbance_inputs)
     (; f, dvs, ps, io_sys = sys)
 end

src/systems/abstractsystem.jl

@@ -717,6 +717,22 @@ function add_initialization_parameters(sys::AbstractSystem)
     return sys
 end
 
+"""
+Returns true if the parameter `p` is of the form `Initial(x)`.
+"""
+function isinitial(p)
+    p = unwrap(p)
+    if iscall(p)
+        operation(p) isa Initial && return true
+        if operation(p) === getindex
+            operation(arguments(p)[1]) isa Initial && return true
+        end
+    else
+        return false
+    end
+    return false
+end
+
 """
 $(TYPEDSIGNATURES)
 
@@ -757,38 +773,21 @@ function complete(sys::AbstractSystem; split = true, flatten = true)
         if !isempty(all_ps)
             # reorder parameters by portions
             ps_split = reorder_parameters(sys, all_ps)
+            # if there are tunables, they will all be in `ps_split[1]`
+            # and the arrays will have been scalarized
+            ordered_ps = eltype(all_ps)[]
             # if there are no tunables, vcat them
-            if isempty(get_index_cache(sys).tunable_idx)
-                ordered_ps = reduce(vcat, ps_split)
-            else
-                # if there are tunables, they will all be in `ps_split[1]`
-                # and the arrays will have been scalarized
-                ordered_ps = eltype(all_ps)[]
-                i = 1
-                # go through all the tunables
-                while i <= length(ps_split[1])
-                    sym = ps_split[1][i]
-                    # if the sym is not a scalarized array symbolic OR it was already scalarized,
-                    # just push it as-is
-                    if !iscall(sym) || operation(sym) != getindex ||
-                       any(isequal(sym), all_ps)
-                        push!(ordered_ps, sym)
-                        i += 1
-                        continue
-                    end
-                    # the next `length(sym)` symbols should be scalarized versions of the same
-                    # array symbolic
-                    if !allequal(first(arguments(x))
-                    for x in view(ps_split[1], i:(i + length(sym) - 1)))
-                        error("This should not be possible. Please open an issue in ModelingToolkit.jl with an MWE and stacktrace.")
-                    end
-                    arrsym = first(arguments(sym))
-                    push!(ordered_ps, arrsym)
-                    i += length(arrsym)
-                end
-                ordered_ps = vcat(
-                    ordered_ps, reduce(vcat, ps_split[2:end]; init = eltype(ordered_ps)[]))
+            if !isempty(get_index_cache(sys).tunable_idx)
+                unflatten_parameters!(ordered_ps, ps_split[1], all_ps)
+                ps_split = Base.tail(ps_split)
+            end
+            # unflatten initial parameters
+            if !isempty(get_index_cache(sys).initials_idx)
+                unflatten_parameters!(ordered_ps, ps_split[1], all_ps)
+                ps_split = Base.tail(ps_split)
             end
+            ordered_ps = vcat(
+                ordered_ps, reduce(vcat, ps_split; init = eltype(ordered_ps)[]))
             @set! sys.ps = ordered_ps
         end
     elseif has_index_cache(sys)
@@ -800,6 +799,39 @@ function complete(sys::AbstractSystem; split = true, flatten = true)
     isdefined(sys, :complete) ? (@set! sys.complete = true) : sys
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Given a flattened array of parameters `params` and a collection of all (unscalarized)
+parameters in the system `all_ps`, unscalarize the elements in `params` and append
+to `buffer` in the same order as they are present in `params`. Effectively, if
+`params = [p[1], p[2], p[3], q]` then this is equivalent to `push!(buffer, p, q)`.
+"""
+function unflatten_parameters!(buffer, params, all_ps)
+    i = 1
+    # go through all the tunables
+    while i <= length(params)
+        sym = params[i]
+        # if the sym is not a scalarized array symbolic OR it was already scalarized,
+        # just push it as-is
+        if !iscall(sym) || operation(sym) != getindex ||
+           any(isequal(sym), all_ps)
+            push!(buffer, sym)
+            i += 1
+            continue
+        end
+        # the next `length(sym)` symbols should be scalarized versions of the same
+        # array symbolic
+        if !allequal(first(arguments(x))
+        for x in view(params, i:(i + length(sym) - 1)))
+            error("This should not be possible. Please open an issue in ModelingToolkit.jl with an MWE and stacktrace.")
+        end
+        arrsym = first(arguments(sym))
+        push!(buffer, arrsym)
+        i += length(arrsym)
+    end
+end
+
 for prop in [:eqs
              :tag
              :noiseeqs
@@ -846,6 +878,7 @@ for prop in [:eqs
              :assertions
              :solved_unknowns
              :split_idxs
+             :ignored_connections
              :parent
              :is_dde
              :tstops
@@ -1362,6 +1395,75 @@ function assertions(sys::AbstractSystem)
     return merge(asserts, namespaced_asserts)
 end
 
+const HierarchyVariableT = Vector{Union{BasicSymbolic, Symbol}}
+const HierarchySystemT = Vector{Union{AbstractSystem, Symbol}}
+"""
+The type returned from `as_hierarchy`.
+"""
+const HierarchyT = Union{HierarchyVariableT, HierarchySystemT}
+
+"""
+    $(TYPEDSIGNATURES)
+
+The inverse operation of `as_hierarchy`.
+"""
+function from_hierarchy(hierarchy::HierarchyT)
+    namefn = hierarchy[1] isa AbstractSystem ? nameof : getname
+    foldl(@view hierarchy[2:end]; init = hierarchy[1]) do sys, name
+        rename(sys, Symbol(name, NAMESPACE_SEPARATOR, namefn(sys)))
+    end
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Represent a namespaced system (or variable) `sys` as a hierarchy. Return a vector, where
+the first element is the unnamespaced system (variable) and subsequent elements are
+`Symbol`s representing the parents of the unnamespaced system (variable) in order from
+inner to outer.
+"""
+function as_hierarchy(sys::Union{AbstractSystem, BasicSymbolic})::HierarchyT
+    namefn = sys isa AbstractSystem ? nameof : getname
+    # get the hierarchy
+    hierarchy = namespace_hierarchy(namefn(sys))
+    # rename the system with unnamespaced name
+    newsys = rename(sys, hierarchy[end])
+    # and remove it from the list
+    pop!(hierarchy)
+    # reverse it to go from inner to outer
+    reverse!(hierarchy)
+    # concatenate
+    T = sys isa AbstractSystem ? AbstractSystem : BasicSymbolic
+    return Union{Symbol, T}[newsys; hierarchy]
+end
+
+"""
+    $(TYPEDSIGNATURES)
+
+Get the connections to ignore for `sys` and its subsystems. The returned value is a
+`Tuple` similar in structure to the `ignored_connections` field. Each system (variable)
+in the first (second) element of the tuple is also passed through `as_hierarchy`.
+"""
+function ignored_connections(sys::AbstractSystem)
+    has_ignored_connections(sys) || return (HierarchySystemT[], HierarchyVariableT[])
+
+    ics = get_ignored_connections(sys)
+    if ics === nothing
+        ics = (HierarchySystemT[], HierarchyVariableT[])
+    end
+    # turn into hierarchies
+    ics = (map(as_hierarchy, ics[1]), map(as_hierarchy, ics[2]))
+    systems = get_systems(sys)
+    # for each subsystem, get its ignored connections, add the name of the subsystem
+    # to the hierarchy and concatenate corresponding buffers of the result
+    result = mapreduce(Broadcast.BroadcastFunction(vcat), systems; init = ics) do subsys
+        sub_ics = ignored_connections(subsys)
+        (map(Base.Fix2(push!, nameof(subsys)), sub_ics[1]),
+            map(Base.Fix2(push!, nameof(subsys)), sub_ics[2]))
+    end
+    return (Vector{HierarchySystemT}(result[1]), Vector{HierarchyVariableT}(result[2]))
+end
+
 """
 $(TYPEDSIGNATURES)
 

src/systems/analysis_points.jl

@@ -412,6 +412,27 @@ function get_analysis_variable(var, name, iv; perturb = true)
     return pvar, default
 end
 
+function with_analysis_point_ignored(sys::AbstractSystem, ap::AnalysisPoint)
+    has_ignored_connections(sys) || return sys
+    ignored = get_ignored_connections(sys)
+    if ignored === nothing
+        ignored = (ODESystem[], BasicSymbolic[])
+    else
+        ignored = copy.(ignored)
+    end
+    if ap.outputs === nothing
+        error("Empty analysis point")
+    end
+    for x in ap.outputs
+        if x isa ODESystem
+            push!(ignored[1], x)
+        else
+            push!(ignored[2], unwrap(x))
+        end
+    end
+    return @set sys.ignored_connections = ignored
+end
+
 #### PRIMITIVE TRANSFORMATIONS
 
 const DOC_WILL_REMOVE_AP = """
@@ -469,7 +490,9 @@ function apply_transformation(tf::Break, sys::AbstractSystem)
         ap = breaksys_eqs[ap_idx].rhs
         deleteat!(breaksys_eqs, ap_idx)
 
-        tf.add_input || return sys, ()
+        breaksys = with_analysis_point_ignored(breaksys, ap)
+
+        tf.add_input || return breaksys, ()
 
         ap_ivar = ap_var(ap.input)
         new_var, new_def = get_analysis_variable(ap_ivar, nameof(ap), get_iv(sys))
@@ -511,7 +534,7 @@ function apply_transformation(tf::GetInput, sys::AbstractSystem)
         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_sys_eqs = get_eqs(ap_sys)
         ap = ap_sys_eqs[ap_idx].rhs
 
         # input variable
@@ -570,6 +593,7 @@ function apply_transformation(tf::PerturbOutput, sys::AbstractSystem)
         ap = ap_sys_eqs[ap_idx].rhs
         # remove analysis point
         deleteat!(ap_sys_eqs, ap_idx)
+        ap_sys = with_analysis_point_ignored(ap_sys, ap)
 
         # add equations involving new variable
         ap_ivar = ap_var(ap.input)
@@ -634,7 +658,7 @@ function apply_transformation(tf::AddVariable, sys::AbstractSystem)
         ap_idx = analysis_point_index(ap_sys, tf.ap)
         ap_idx === nothing &&
             error("Analysis point $(nameof(tf.ap)) not found in system $(nameof(sys)).")
-        ap_sys_eqs = copy(get_eqs(ap_sys))
+        ap_sys_eqs = get_eqs(ap_sys)
         ap = ap_sys_eqs[ap_idx].rhs
 
         # add equations involving new variable

src/systems/codegen_utils.jl

@@ -287,7 +287,7 @@ end
         # The user provided a single buffer/tuple for the parameter object, so wrap that
         # one in a tuple
         fargs = ntuple(Val(length(args))) do i
-            i == paramidx ? :((args[$i],)) : :(args[$i])
+            i == paramidx ? :((args[$i], nothing)) : :(args[$i])
         end
         return :($f($(fargs...)))
     end