Compute inputs for each clock

Co-authored-by: Fredrik Bagge Carlson <[email protected]>

Author: YingboMa

src/inputoutput.jl

@@ -302,12 +302,13 @@ function inputs_to_parameters!(state::TransformationState, io)
     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 io.inputs]
+        permutation = [d[i] for i in inputs]
         new_parameters = new_parameters[permutation]
     end
 

src/systems/abstractsystem.jl

@@ -1038,20 +1038,6 @@ function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false
     structural_simplify!(state, io; simplify, kwargs...)
 end
 
-function structural_simplify!(state::TearingState, io = nothing; simplify = false,
-                              kwargs...)
-    has_io = io !== nothing
-    has_io && markio!(state, io...)
-    state, input_idxs = inputs_to_parameters!(state, io)
-    sys, ag = alias_elimination!(state; kwargs...)
-    #check_consistency(state, ag)
-    sys = dummy_derivative(sys, state, ag; simplify)
-    fullstates = [map(eq -> eq.lhs, observed(sys)); states(sys)]
-    @set! sys.observed = topsort_equations(observed(sys), fullstates)
-    invalidate_cache!(sys)
-    return has_io ? (sys, input_idxs) : sys
-end
-
 function eliminate_constants(sys::AbstractSystem)
     if has_eqs(sys)
         eqs = get_eqs(sys)
@@ -1066,7 +1052,7 @@ end
 
 function io_preprocessing(sys::AbstractSystem, inputs,
                           outputs; simplify = false, kwargs...)
-    sys, input_idxs = structural_simplify(sys, (; inputs, outputs); simplify, kwargs...)
+    sys, input_idxs = structural_simplify(sys, (inputs, outputs); simplify, kwargs...)
 
     eqs = equations(sys)
     alg_start_idx = findfirst(!isdiffeq, eqs)
@@ -1169,9 +1155,13 @@ function markio!(state, inputs, outputs; check = true)
             fullvars[i] = v
         end
     end
-    check && (all(values(inputset)) ||
-     error("Some specified inputs were not found in system. The following Dict indicates the found variables ",
-           inputset))
+    if check
+        ikeys = keys(filter(!last, inputset))
+        if !isempty(ikeys)
+            error("Some specified inputs were not found in system. The following variables were not found ",
+                  ikeys)
+        end
+    end
     check && (all(values(outputset)) ||
      error("Some specified outputs were not found in system. The following Dict indicates the found variables ",
            outputset))

src/systems/clock_inference.jl

@@ -58,8 +58,10 @@ function infer_clocks!(ci::ClockInference)
         vd = var_domain[v]
         eqs = 𝑑neighbors(graph, v)
         isempty(eqs) && continue
-        eq = first(eqs)
-        eq_domain[eq] = vd
+        #eq = first(eqs)
+        for eq in eqs
+            eq_domain[eq] = vd
+        end
     end
 
     return ci
@@ -80,38 +82,42 @@ function split_system(ci::ClockInference)
     @unpack ts, eq_domain, var_domain, inferred = ci
     @unpack fullvars = ts
     @unpack graph = ts.structure
-    continuous_id = 0
+    continuous_id = Ref(0)
     clock_to_id = Dict{TimeDomain, Int}()
     id_to_clock = TimeDomain[]
     eq_to_cid = Vector{Int}(undef, nsrcs(graph))
     cid_to_eq = Vector{Int}[]
     var_to_cid = Vector{Int}(undef, ndsts(graph))
     cid_to_var = Vector{Int}[]
-    cid = 0
+    cid_counter = Ref(0)
     for (i, d) in enumerate(eq_domain)
-        cid = get!(clock_to_id, d) do
-            cid += 1
-            push!(id_to_clock, d)
-            if d isa Continuous
-                continuous_id = cid
+        cid = let cid_counter = cid_counter, id_to_clock = id_to_clock,
+            continuous_id = continuous_id
+
+            get!(clock_to_id, d) do
+                cid = (cid_counter[] += 1)
+                push!(id_to_clock, d)
+                if d isa Continuous
+                    continuous_id[] = cid
+                end
+                cid
             end
-            cid
         end
         eq_to_cid[i] = cid
         resize_or_push!(cid_to_eq, i, cid)
     end
-    input_discrete = Int[]
-    inputs = []
+    input_idxs = map(_ -> Int[], 1:cid_counter[])
+    inputs = map(_ -> Any[], 1:cid_counter[])
     for (i, d) in enumerate(var_domain)
         cid = get(clock_to_id, d, 0)
         @assert cid!==0 "Internal error!"
         var_to_cid[i] = cid
         v = fullvars[i]
         #TODO: remove Inferred*
-        if cid == continuous_id && istree(v) && (o = operation(v)) isa Operator &&
-           !(input_timedomain(o) isa Continuous)
-            push!(input_discrete, i)
-            push!(inputs, fullvars[i])
+        if istree(v) && (o = operation(v)) isa Operator &&
+           input_timedomain(o) != output_timedomain(o)
+            push!(input_idxs[cid], i)
+            push!(inputs[cid], fullvars[i])
         end
         resize_or_push!(cid_to_var, i, cid)
     end
@@ -123,20 +129,23 @@ function split_system(ci::ClockInference)
         ts_i = ts
         fadj = Vector{Int}[]
         eqs_i = Equation[]
+        eq_to_diff = DiffGraph(length(ieqs))
         var_set_i = BitSet(vars)
         ne = 0
-        for eq_i in ieqs
+        for (j, eq_i) in enumerate(ieqs)
             vars = copy(graph.fadjlist[eq_i])
             ne += length(vars)
             push!(fadj, vars)
             push!(eqs_i, eqs[eq_i])
+            eq_to_diff[j] = ts_i.structure.eq_to_diff[eq_i]
         end
         @set! ts_i.structure.graph = complete(BipartiteGraph(ne, fadj, ndsts(graph)))
         @set! ts_i.sys.eqs = eqs_i
+        @set! ts_i.structure.eq_to_diff = eq_to_diff
         tss[id] = ts_i
         # TODO: just mark past and sample variables as inputs
     end
-    return tss, (; inputs, outputs = ())
+    return tss, inputs
 
     #id_to_clock, cid_to_eq, cid_to_var
 end

src/systems/systemstructure.jl

@@ -25,7 +25,7 @@ function quick_cancel_expr(expr)
                                                                          kws...))(expr)
 end
 
-export SystemStructure, TransformationState, TearingState
+export SystemStructure, TransformationState, TearingState, structural_simplify!
 export initialize_system_structure, find_linear_equations
 export isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq, algeqs
 export dervars_range, diffvars_range, algvars_range
@@ -424,4 +424,19 @@ function Base.show(io::IO, mime::MIME"text/plain", ms::MatchedSystemStructure)
                                                  complete(ms.var_eq_matching, nsrcs(graph))))
 end
 
+# TODO: clean up
+function structural_simplify!(state::TearingState, io = nothing; simplify = false,
+                              kwargs...)
+    has_io = io !== nothing
+    has_io && ModelingToolkit.markio!(state, io...)
+    state, input_idxs = ModelingToolkit.inputs_to_parameters!(state, io)
+    sys, ag = ModelingToolkit.alias_elimination!(state; kwargs...)
+    #check_consistency(state, ag)
+    sys = ModelingToolkit.dummy_derivative(sys, state, ag; simplify)
+    fullstates = [map(eq -> eq.lhs, observed(sys)); states(sys)]
+    @set! sys.observed = ModelingToolkit.topsort_equations(observed(sys), fullstates)
+    ModelingToolkit.invalidate_cache!(sys)
+    return has_io ? (sys, input_idxs) : sys
+end
+
 end # module

test/clock.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, Test
+using ModelingToolkit, Test, Setfield
 
 function infer_clocks(sys)
     ts = TearingState(sys)
@@ -14,6 +14,7 @@ D = Differential(t)
 
 eqs = [yd ~ Sample(t, dt)(y)
        ud ~ kp * (r - yd)
+       r ~ 1.0
 
        # plant (time continuous part)
        u ~ Hold(ud)
@@ -61,19 +62,21 @@ By inference:
 
 ci, varmap = infer_clocks(sys)
 eqmap = ci.eq_domain
-tss, io = ModelingToolkit.split_system(deepcopy(ci))
-ts_c = deepcopy(tss[1])
-@set! ts_c.structure.solvable_graph = nothing
-sss, = ModelingToolkit.structural_simplify!(ts_c, io)
+tss, inputs = ModelingToolkit.split_system(deepcopy(ci))
+sss, = ModelingToolkit.structural_simplify!(deepcopy(tss[1]), (inputs[1], ()))
 @test equations(sss) == [D(x) ~ u - x]
+sss, = ModelingToolkit.structural_simplify!(deepcopy(tss[2]), (inputs[2], ()))
+@test isempty(equations(sss))
+@test observed(sss) == [r ~ 1.0; yd ~ Sample(t, dt)(y); ud ~ kp * (r - yd)]
 
 d = Clock(t, dt)
 # Note that TearingState reorders the equations
 @test eqmap[1] == Continuous()
 @test eqmap[2] == d
 @test eqmap[3] == d
-@test eqmap[4] == Continuous()
+@test eqmap[4] == d
 @test eqmap[5] == Continuous()
+@test eqmap[6] == Continuous()
 
 @test varmap[yd] == d
 @test varmap[ud] == d