Merge pull request #3573 from vyudu/ss_simplify

[v10] refactor: change inputs/outputs handling in `structural_simplify`

Author: AayushSabharwal

src/inputoutput.jl

@@ -163,7 +163,7 @@ has_var(ex, x) = x ∈ Set(get_variables(ex))
     (f_oop, f_ip), x_sym, p_sym, io_sys = generate_control_function(
             sys::AbstractODESystem,
             inputs             = unbound_inputs(sys),
-            disturbance_inputs = nothing;
+            disturbance_inputs = disturbances(sys);
             implicit_dae       = false,
             simplify           = false,
         )
@@ -179,9 +179,6 @@ The return values also include the chosen state-realization (the remaining unkno
 
 If `disturbance_inputs` is an array of variables, the generated dynamics function will preserve any state and dynamics associated with disturbance inputs, but the disturbance inputs themselves will (by default) not be included as inputs to the generated function. The use case for this is to generate dynamics for state observers that estimate the influence of unmeasured disturbances, and thus require unknown variables for the disturbance model, but without disturbance inputs since the disturbances are not available for measurement. To add an input argument corresponding to the disturbance inputs, either include the disturbance inputs among the control inputs, or set `disturbance_argument=true`, in which case an additional input argument `w` is added to the generated function `(x,u,p,t,w)->rhs`.
 
-!!! note "Un-simplified system"
-    This function expects `sys` to be un-simplified, i.e., `structural_simplify` or `@mtkbuild` should not be called on the system before passing it into this function. `generate_control_function` calls a special version of `structural_simplify` internally.
-
 # Example
 
 ```
@@ -200,16 +197,18 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
         simplify = false,
         eval_expression = false,
         eval_module = @__MODULE__,
+        check_simplified = true,
         kwargs...)
+    # Remove this when the ControlFunction gets merged.
+    if check_simplified && !iscomplete(sys)
+        error("A completed `ODESystem` is required. Call `complete` or `structural_simplify` on the system before creating the control function.")
+    end
     isempty(inputs) && @warn("No unbound inputs were found in system.")
-
     if disturbance_inputs !== nothing
         # add to inputs for the purposes of io processing
         inputs = [inputs; disturbance_inputs]
     end
 
-    sys, _ = io_preprocessing(sys, inputs, []; simplify, kwargs...)
-
     dvs = unknowns(sys)
     ps = parameters(sys; initial_parameters = true)
     ps = setdiff(ps, inputs)
@@ -257,8 +256,11 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
     (; f, dvs, ps, io_sys = sys)
 end
 
-function inputs_to_parameters!(state::TransformationState, io)
-    check_bound = io === nothing
+"""
+Turn input variables into parameters of the system.
+"""
+function inputs_to_parameters!(state::TransformationState, inputsyms)
+    check_bound = inputsyms === nothing
     @unpack structure, fullvars, sys = state
     @unpack var_to_diff, graph, solvable_graph = structure
     @assert solvable_graph === nothing
@@ -287,7 +289,7 @@ function inputs_to_parameters!(state::TransformationState, io)
             push!(new_fullvars, v)
         end
     end
-    ninputs == 0 && return (state, 1:0)
+    ninputs == 0 && return state
 
     nvars = ndsts(graph) - ninputs
     new_graph = BipartiteGraph(nsrcs(graph), nvars, Val(false))
@@ -316,24 +318,11 @@ function inputs_to_parameters!(state::TransformationState, io)
     @set! sys.unknowns = setdiff(unknowns(sys), keys(input_to_parameters))
     ps = parameters(sys)
 
-    if io !== nothing
-        inputs, = io
-        # Change order of new parameters to correspond to user-provided order in argument `inputs`
-        d = Dict{Any, Int}()
-        for (i, inp) in enumerate(new_parameters)
-            d[inp] = i
-        end
-        permutation = [d[i] for i in inputs]
-        new_parameters = new_parameters[permutation]
-    end
-
     @set! sys.ps = [ps; new_parameters]
-
     @set! state.sys = sys
     @set! state.fullvars = new_fullvars
     @set! state.structure = structure
-    base_params = length(ps)
-    return state, (base_params + 1):(base_params + length(new_parameters)) # (1:length(new_parameters)) .+ base_params
+    return state
 end
 
 """
@@ -359,7 +348,7 @@ function get_disturbance_system(dist::DisturbanceModel{<:ODESystem})
 end
 
 """
-    (f_oop, f_ip), augmented_sys, dvs, p = add_input_disturbance(sys, dist::DisturbanceModel, inputs = nothing)
+    (f_oop, f_ip), augmented_sys, dvs, p = add_input_disturbance(sys, dist::DisturbanceModel, inputs = Any[])
 
 Add a model of an unmeasured disturbance to `sys`. The disturbance model is an instance of [`DisturbanceModel`](@ref).
 
@@ -408,13 +397,13 @@ model_outputs = [model.inertia1.w, model.inertia2.w, model.inertia1.phi, model.i
 
 `f_oop` will have an extra state corresponding to the integrator in the disturbance model. This state will not be affected by any input, but will affect the dynamics from where it enters, in this case it will affect additively from `model.torque.tau.u`.
 """
-function add_input_disturbance(sys, dist::DisturbanceModel, inputs = nothing; kwargs...)
+function add_input_disturbance(sys, dist::DisturbanceModel, inputs = Any[]; kwargs...)
     t = get_iv(sys)
     @variables d(t)=0 [disturbance = true]
     @variables u(t)=0 [input = true] # New system input
     dsys = get_disturbance_system(dist)
 
-    if inputs === nothing
+    if isempty(inputs)
         all_inputs = [u]
     else
         i = findfirst(isequal(dist.input), inputs)
@@ -429,8 +418,9 @@ function add_input_disturbance(sys, dist::DisturbanceModel, inputs = nothing; kw
            dist.input ~ u + dsys.output.u[1]]
     augmented_sys = ODESystem(eqs, t, systems = [dsys], name = gensym(:outer))
     augmented_sys = extend(augmented_sys, sys)
+    ssys = structural_simplify(augmented_sys, inputs = all_inputs, disturbance_inputs = [d])
 
-    (f_oop, f_ip), dvs, p, io_sys = generate_control_function(augmented_sys, all_inputs,
+    (f_oop, f_ip), dvs, p, io_sys = generate_control_function(ssys, all_inputs,
         [d]; kwargs...)
     (f_oop, f_ip), augmented_sys, dvs, p, io_sys
 end

src/linearization.jl

@@ -58,9 +58,8 @@ function linearization_function(sys::AbstractSystem, inputs,
     outputs = mapreduce(vcat, outputs; init = []) do var
         symbolic_type(var) == ArraySymbolic() ? collect(var) : [var]
     end
-    ssys, diff_idxs, alge_idxs, input_idxs = io_preprocessing(sys, inputs, outputs;
-        simplify,
-        kwargs...)
+    ssys = structural_simplify(sys; inputs, outputs, simplify, kwargs...)
+    diff_idxs, alge_idxs = eq_idxs(ssys)
     if zero_dummy_der
         dummyder = setdiff(unknowns(ssys), unknowns(sys))
         defs = Dict(x => 0.0 for x in dummyder)
@@ -87,9 +86,9 @@ function linearization_function(sys::AbstractSystem, inputs,
 
     p = parameter_values(prob)
     t0 = current_time(prob)
-    inputvals = [p[idx] for idx in input_idxs]
+    inputvals = [prob.ps[i] for i in inputs]
 
-    hp_fun = let fun = h, setter = setp_oop(sys, input_idxs)
+    hp_fun = let fun = h, setter = setp_oop(sys, inputs)
         function hpf(du, input, u, p, t)
             p = setter(p, input)
             fun(du, u, p, t)
@@ -113,7 +112,7 @@ function linearization_function(sys::AbstractSystem, inputs,
         # observed function is a `GeneratedFunctionWrapper` with iip component
         h_jac = PreparedJacobian{true}(h, similar(prob.u0, size(outputs)), autodiff,
             prob.u0, DI.Constant(p), DI.Constant(t0))
-        pf_fun = let fun = prob.f, setter = setp_oop(sys, input_idxs)
+        pf_fun = let fun = prob.f, setter = setp_oop(sys, inputs)
             function pff(du, input, u, p, t)
                 p = setter(p, input)
                 SciMLBase.ParamJacobianWrapper(fun, t, u)(du, p)
@@ -127,12 +126,23 @@ function linearization_function(sys::AbstractSystem, inputs,
     end
 
     lin_fun = LinearizationFunction(
-        diff_idxs, alge_idxs, input_idxs, length(unknowns(sys)),
+        diff_idxs, alge_idxs, inputs, length(unknowns(sys)),
         prob, h, u0 === nothing ? nothing : similar(u0), uf_jac, h_jac, pf_jac,
         hp_jac, initializealg, initialization_kwargs)
     return lin_fun, sys
 end
 
+"""
+Return the set of indexes of differential equations and algebraic equations in the simplified system.
+"""
+function eq_idxs(sys::AbstractSystem)
+    eqs = equations(sys)
+    alge_idxs = findall(!isdiffeq, eqs)
+    diff_idxs = setdiff(1:length(eqs), alge_idxs)
+
+    diff_idxs, alge_idxs
+end
+
 """
     $(TYPEDEF)
 
@@ -192,7 +202,7 @@ A callable struct which linearizes a system.
 $(TYPEDFIELDS)
 """
 struct LinearizationFunction{
-    DI <: AbstractVector{Int}, AI <: AbstractVector{Int}, II, P <: ODEProblem,
+    DI <: AbstractVector{Int}, AI <: AbstractVector{Int}, I, P <: ODEProblem,
     H, C, J1, J2, J3, J4, IA <: SciMLBase.DAEInitializationAlgorithm, IK}
     """
     The indexes of differential equations in the linearized system.
@@ -206,7 +216,7 @@ struct LinearizationFunction{
     The indexes of parameters in the linearized system which represent
     input variables.
     """
-    input_idxs::II
+    inputs::I
     """
     The number of unknowns in the linearized system.
     """
@@ -281,6 +291,7 @@ function (linfun::LinearizationFunction)(u, p, t)
     end
 
     fun = linfun.prob.f
+    input_vals = [linfun.prob.ps[i] for i in linfun.inputs]
     if u !== nothing # Handle systems without unknowns
         linfun.num_states == length(u) ||
             error("Number of unknown variables ($(linfun.num_states)) does not match the number of input unknowns ($(length(u)))")
@@ -294,15 +305,15 @@ function (linfun::LinearizationFunction)(u, p, t)
         end
         fg_xz = linfun.uf_jac(u, DI.Constant(p), DI.Constant(t))
         h_xz = linfun.h_jac(u, DI.Constant(p), DI.Constant(t))
-        fg_u = linfun.pf_jac([p[idx] for idx in linfun.input_idxs],
+        fg_u = linfun.pf_jac(input_vals,
             DI.Constant(u), DI.Constant(p), DI.Constant(t))
     else
         linfun.num_states == 0 ||
             error("Number of unknown variables (0) does not match the number of input unknowns ($(length(u)))")
         fg_xz = zeros(0, 0)
-        h_xz = fg_u = zeros(0, length(linfun.input_idxs))
+        h_xz = fg_u = zeros(0, length(linfun.inputs))
     end
-    h_u = linfun.hp_jac([p[idx] for idx in linfun.input_idxs],
+    h_u = linfun.hp_jac(input_vals,
         DI.Constant(u), DI.Constant(p), DI.Constant(t))
     (f_x = fg_xz[linfun.diff_idxs, linfun.diff_idxs],
         f_z = fg_xz[linfun.diff_idxs, linfun.alge_idxs],
@@ -482,9 +493,8 @@ function linearize_symbolic(sys::AbstractSystem, inputs,
         outputs; simplify = false, allow_input_derivatives = false,
         eval_expression = false, eval_module = @__MODULE__,
         kwargs...)
-    sys, diff_idxs, alge_idxs, input_idxs = io_preprocessing(
-        sys, inputs, outputs; simplify,
-        kwargs...)
+    sys = structural_simplify(sys; inputs, outputs, simplify, kwargs...)
+    diff_idxs, alge_idxs = eq_idxs(sys)
     sts = unknowns(sys)
     t = get_iv(sys)
     ps = parameters(sys; initial_parameters = true)
@@ -545,10 +555,14 @@ function linearize_symbolic(sys::AbstractSystem, inputs,
     (; A, B, C, D, f_x, f_z, g_x, g_z, f_u, g_u, h_x, h_z, h_u), sys
 end
 
-function markio!(state, orig_inputs, inputs, outputs; check = true)
+"""
+Modify the variable metadata of system variables to indicate which ones are inputs, outputs, and disturbances. Needed for `inputs`, `outputs`, `disturbances`, `unbound_inputs`, `unbound_outputs` to return the proper subsets.
+"""
+function markio!(state, orig_inputs, inputs, outputs, disturbances; check = true)
     fullvars = get_fullvars(state)
     inputset = Dict{Any, Bool}(i => false for i in inputs)
     outputset = Dict{Any, Bool}(o => false for o in outputs)
+    disturbanceset = Dict{Any, Bool}(d => false for d in disturbances)
     for (i, v) in enumerate(fullvars)
         if v in keys(inputset)
             if v in keys(outputset)
@@ -570,6 +584,13 @@ function markio!(state, orig_inputs, inputs, outputs; check = true)
             v = setio(v, false, false)
             fullvars[i] = v
         end
+
+        if v in keys(disturbanceset)
+            v = setio(v, true, false)
+            v = setdisturbance(v, true)
+            disturbanceset[v] = true
+            fullvars[i] = v
+        end
     end
     if check
         ikeys = keys(filter(!last, inputset))
@@ -578,11 +599,16 @@ function markio!(state, orig_inputs, inputs, outputs; check = true)
                 "Some specified inputs were not found in system. The following variables were not found ",
                 ikeys)
         end
+        dkeys = keys(filter(!last, disturbanceset))
+        if !isempty(dkeys)
+            error(
+                "Specified disturbance inputs were not found in system. The following variables were not found ",
+                ikeys)
+        end
+        (all(values(outputset)) || error(
+            "Some specified outputs were not found in system. The following Dict indicates the found variables ",
+            outputset))
     end
-    check && (all(values(outputset)) ||
-     error(
-        "Some specified outputs were not found in system. The following Dict indicates the found variables ",
-        outputset))
     state, orig_inputs
 end
 

src/systems/abstractsystem.jl

@@ -2484,21 +2484,6 @@ function eliminate_constants(sys::AbstractSystem)
     return sys
 end
 
-function io_preprocessing(sys::AbstractSystem, inputs,
-        outputs; simplify = false, kwargs...)
-    sys, input_idxs = structural_simplify(sys, (inputs, outputs); simplify, kwargs...)
-
-    eqs = equations(sys)
-    alg_start_idx = findfirst(!isdiffeq, eqs)
-    if alg_start_idx === nothing
-        alg_start_idx = length(eqs) + 1
-    end
-    diff_idxs = 1:(alg_start_idx - 1)
-    alge_idxs = alg_start_idx:length(eqs)
-
-    sys, diff_idxs, alge_idxs, input_idxs
-end
-
 @latexrecipe function f(sys::AbstractSystem)
     return latexify(equations(sys))
 end

src/systems/analysis_points.jl

@@ -958,7 +958,6 @@ function linearization_function(sys::AbstractSystem,
     end
 
     sys = handle_loop_openings(sys, map(AnalysisPoint, collect(loop_openings)))
-
     return linearization_function(system_modifier(sys), input_vars, output_vars; kwargs...)
 end
 

src/systems/clock_inference.jl

@@ -1,7 +1,11 @@
 struct ClockInference{S}
+    """Tearing state."""
     ts::S
+    """The time domain (discrete clock, continuous) of each equation."""
     eq_domain::Vector{TimeDomain}
+    """The output time domain (discrete clock, continuous) of each variable."""
     var_domain::Vector{TimeDomain}
+    """The set of variables with concrete domains."""
     inferred::BitSet
 end
 
@@ -67,6 +71,9 @@ function substitute_sample_time(ex, dt)
     end
 end
 
+"""
+Update the equation-to-time domain mapping by inferring the time domain from the variables.
+"""
 function infer_clocks!(ci::ClockInference)
     @unpack ts, eq_domain, var_domain, inferred = ci
     @unpack var_to_diff, graph = ts.structure
@@ -132,6 +139,9 @@ function is_time_domain_conversion(v)
         input_timedomain(o) != output_timedomain(o)
 end
 
+"""
+For multi-clock systems, create a separate system for each clock in the system, along with associated equations. Return the updated tearing state, and the sets of clocked variables associated with each time domain.
+"""
 function split_system(ci::ClockInference{S}) where {S}
     @unpack ts, eq_domain, var_domain, inferred = ci
     fullvars = get_fullvars(ts)
@@ -143,11 +153,14 @@ function split_system(ci::ClockInference{S}) where {S}
     cid_to_eq = Vector{Int}[]
     var_to_cid = Vector{Int}(undef, ndsts(graph))
     cid_to_var = Vector{Int}[]
+    # cid_counter = number of clocks
     cid_counter = Ref(0)
     for (i, d) in enumerate(eq_domain)
         cid = let cid_counter = cid_counter, id_to_clock = id_to_clock,
             continuous_id = continuous_id
 
+            # Fill the clock_to_id dict as you go,
+            # ContinuousClock() => 1, ...
             get!(clock_to_id, d) do
                 cid = (cid_counter[] += 1)
                 push!(id_to_clock, d)