refactor: remove input_idxs output

Author: vyudu

src/inputoutput.jl

@@ -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
 
 ```
@@ -201,17 +198,17 @@ function generate_control_function(sys::AbstractODESystem, inputs = unbound_inpu
         eval_expression = false,
         eval_module = @__MODULE__,
         kwargs...)
-    isempty(inputs) && @warn("No unbound inputs were found in system.")
 
+    # Remove this when the ControlFunction gets merged.
+    if !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
 
-    if !iscomplete(sys)
-        sys, _ = io_preprocessing(sys, inputs, []; simplify, kwargs...)
-    end
-
     dvs = unknowns(sys)
     ps = parameters(sys; initial_parameters = true)
     ps = setdiff(ps, inputs)
@@ -257,8 +254,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

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)

src/systems/systems.jl

@@ -27,12 +27,15 @@ topological sort of the observed equations in `sys`.
 + `sort_eqs=true` controls whether equations are sorted lexicographically before simplification or not.
 """
 function structural_simplify(
-        sys::AbstractSystem, io = nothing; additional_passes = [], simplify = false, split = true,
+        sys::AbstractSystem; additional_passes = [], simplify = false, split = true,
         allow_symbolic = false, allow_parameter = true, conservative = false, fully_determined = true,
+        inputs = nothing, outputs = nothing,
+        disturbance_inputs = nothing,
         kwargs...)
     isscheduled(sys) && throw(RepeatedStructuralSimplificationError())
-    newsys′ = __structural_simplify(sys, io; simplify,
+    newsys′ = __structural_simplify(sys; simplify,
         allow_symbolic, allow_parameter, conservative, fully_determined,
+        inputs, outputs, disturbance_inputs,
         kwargs...)
     if newsys′ isa Tuple
         @assert length(newsys′) == 2
@@ -70,8 +73,9 @@ function __structural_simplify(sys::SDESystem, args...; kwargs...)
     return __structural_simplify(ODESystem(sys), args...; kwargs...)
 end
 
-function __structural_simplify(
-        sys::AbstractSystem, io = nothing; simplify = false, sort_eqs = true,
+function __structural_simplify(sys::AbstractSystem; simplify = false,
+        inputs = Any[], outputs = Any[],
+        disturbance_inputs = Any[],
         kwargs...)
     sys = expand_connections(sys)
     state = TearingState(sys; sort_eqs)
@@ -90,7 +94,8 @@ function __structural_simplify(
         end
     end
     if isempty(brown_vars)
-        return structural_simplify!(state, io; simplify, kwargs...)
+        return structural_simplify!(
+            state; simplify, inputs, outputs, disturbance_inputs, kwargs...)
     else
         Is = Int[]
         Js = Int[]
@@ -122,8 +127,8 @@ function __structural_simplify(
                               for (i, v) in enumerate(fullvars)
                               if !iszero(new_idxs[i]) &&
                                  invview(var_to_diff)[i] === nothing]
-        # TODO: IO is not handled.
-        ode_sys = structural_simplify(sys, io; simplify, kwargs...)
+        ode_sys = structural_simplify(
+            sys; simplify, inputs, outputs, disturbance_inputs, kwargs...)
         eqs = equations(ode_sys)
         sorted_g_rows = zeros(Num, length(eqs), size(g, 2))
         for (i, eq) in enumerate(eqs)

src/systems/systemstructure.jl

@@ -657,29 +657,21 @@ function Base.show(io::IO, mime::MIME"text/plain", ms::MatchedSystemStructure)
     printstyled(io, " SelectedState")
 end
 
-# TODO: clean up
-function merge_io(io, inputs)
-    isempty(inputs) && return io
-    if io === nothing
-        io = (inputs, [])
-    else
-        io = ([inputs; io[1]], io[2])
-    end
-    return io
-end
-
-function structural_simplify!(state::TearingState, io = nothing; simplify = false,
+function structural_simplify!(state::TearingState; simplify = false,
         check_consistency = true, fully_determined = true, warn_initialize_determined = true,
+        inputs = nothing, outputs = nothing, 
+        disturbance_inputs = nothing,
         kwargs...)
     if state.sys isa ODESystem
         ci = ModelingToolkit.ClockInference(state)
         ci = ModelingToolkit.infer_clocks!(ci)
         time_domains = merge(Dict(state.fullvars .=> ci.var_domain),
             Dict(default_toterm.(state.fullvars) .=> ci.var_domain))
-        tss, inputs, continuous_id, id_to_clock = ModelingToolkit.split_system(ci)
-        cont_io = merge_io(io, inputs[continuous_id])
-        sys, input_idxs = _structural_simplify!(tss[continuous_id], cont_io; simplify,
+        tss, clocked_inputs, continuous_id, id_to_clock = ModelingToolkit.split_system(ci)
+        cont_inputs = [inputs; clocked_inputs[continuous_id]]
+        sys = _structural_simplify!(tss[continuous_id]; simplify,
             check_consistency, fully_determined,
+            cont_inputs, outputs, disturbance_inputs,
             kwargs...)
         if length(tss) > 1
             if continuous_id > 0
@@ -695,8 +687,9 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
                     discrete_subsystems[i] = sys
                     continue
                 end
-                dist_io = merge_io(io, inputs[i])
-                ss, = _structural_simplify!(state, dist_io; simplify, check_consistency,
+                disc_inputs = [inputs; clocked_inputs[i]]
+                ss, = _structural_simplify!(state; simplify, check_consistency,
+                    inputs = disc_inputs, outputs, disturbance_inputs,
                     fully_determined, kwargs...)
                 append!(appended_parameters, inputs[i], unknowns(ss))
                 discrete_subsystems[i] = ss
@@ -713,31 +706,29 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
               for sym in get_ps(sys)]
         @set! sys.ps = ps
     else
-        sys, input_idxs = _structural_simplify!(state, io; simplify, check_consistency,
+        sys = _structural_simplify!(state; simplify, check_consistency,
+            inputs, outputs, disturbance_inputs,
             fully_determined, kwargs...)
     end
-    has_io = io !== nothing
-    return has_io ? (sys, input_idxs) : sys
+    return sys
 end
 
-function _structural_simplify!(state::TearingState, io; simplify = false,
+function _structural_simplify!(state::TearingState; simplify = false,
         check_consistency = true, fully_determined = true, warn_initialize_determined = false,
         dummy_derivative = true,
+        inputs = nothing, outputs = nothing,
+        disturbance_inputs = nothing,
         kwargs...)
     if fully_determined isa Bool
         check_consistency &= fully_determined
     else
         check_consistency = true
     end
-    has_io = io !== nothing
+    has_io = inputs !== nothing || outputs !== nothing
     orig_inputs = Set()
     if has_io
-        ModelingToolkit.markio!(state, orig_inputs, io...)
-    end
-    if io !== nothing
-        state, input_idxs = ModelingToolkit.inputs_to_parameters!(state, io)
-    else
-        input_idxs = 0:-1 # Empty range
+        ModelingToolkit.markio!(state, orig_inputs, inputs, outputs)
+        state = ModelingToolkit.inputs_to_parameters!(state, inputs)
     end
     sys, mm = ModelingToolkit.alias_elimination!(state; kwargs...)
     if check_consistency
@@ -761,5 +752,5 @@ function _structural_simplify!(state::TearingState, io; simplify = false,
     fullunknowns = [observables(sys); unknowns(sys)]
     @set! sys.observed = ModelingToolkit.topsort_equations(observed(sys), fullunknowns)
 
-    ModelingToolkit.invalidate_cache!(sys), input_idxs
+    ModelingToolkit.invalidate_cache!(sys)
 end