Merge pull request #1913 from SciML/myb_fb/simplify_clock

Add `split_system` that splits the system by their time domain

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/structural_transformation/pantelides.jl

@@ -131,6 +131,9 @@ function pantelides!(state::TransformationState, ag::Union{AliasGraph, Nothing}
             pathfound = construct_augmenting_path!(var_eq_matching, graph, eq′,
                                                    v -> varwhitelist[v], vcolor, ecolor)
             pathfound && break # terminating condition
+            if is_only_discrete(state.structure)
+                error("The discrete system has high structural index. This is not supported.")
+            end
             for var in eachindex(vcolor)
                 vcolor[var] || continue
                 if var_to_diff[var] === nothing

src/structural_transformation/partial_state_selection.jl

@@ -154,8 +154,8 @@ function dummy_derivative_graph!(state::TransformationState, jac = nothing,
                                  (ag, diff_va) = (nothing, nothing);
                                  state_priority = nothing, kwargs...)
     state.structure.solvable_graph === nothing && find_solvables!(state; kwargs...)
-    var_eq_matching = complete(pantelides!(state, ag))
     complete!(state.structure)
+    var_eq_matching = complete(pantelides!(state, ag))
     dummy_derivative_graph!(state.structure, var_eq_matching, jac, (ag, diff_va),
                             state_priority)
 end

src/structural_transformation/symbolics_tearing.jl

@@ -92,7 +92,7 @@ function full_equations(sys::AbstractSystem; simplify = false)
     @unpack subs = substitutions
     solved = Dict(eq.lhs => eq.rhs for eq in subs)
     neweqs = map(equations(sys)) do eq
-        if isdiffeq(eq)
+        if istree(eq.lhs) && operation(eq.lhs) isa Union{Shift, Differential}
             return tearing_sub(eq.lhs, solved, simplify) ~ tearing_sub(eq.rhs, solved,
                                                                        simplify)
         else
@@ -262,7 +262,11 @@ function tearing_reassemble(state::TearingState, var_eq_matching, ag = nothing;
 
     if ModelingToolkit.has_iv(state.sys)
         iv = get_iv(state.sys)
-        D = Differential(iv)
+        if is_only_discrete(state.structure)
+            D = Shift(iv, 1)
+        else
+            D = Differential(iv)
+        end
     else
         iv = D = nothing
     end
@@ -628,10 +632,16 @@ function tearing_reassemble(state::TearingState, var_eq_matching, ag = nothing;
     @set! state.structure.var_to_diff = var_to_diff
     @set! state.structure.eq_to_diff = eq_to_diff
     @set! state.fullvars = fullvars = fullvars[invvarsperm]
+    ispresent = let var_to_diff = var_to_diff, graph = graph
+        i -> (!isempty(𝑑neighbors(graph, i)) ||
+              (var_to_diff[i] !== nothing && !isempty(𝑑neighbors(graph, var_to_diff[i]))))
+    end
 
     sys = state.sys
     @set! sys.eqs = neweqs
-    @set! sys.states = [v for (i, v) in enumerate(fullvars) if diff_to_var[i] === nothing]
+    @set! sys.states = Any[v
+                           for (i, v) in enumerate(fullvars)
+                           if diff_to_var[i] === nothing && ispresent(i)]
     removed_obs_set = BitSet(removed_obs)
     var_to_idx = Dict(reverse(en) for en in enumerate(fullvars))
     # Make sure differentiated variables don't appear in observed equations

src/structural_transformation/utils.jl

@@ -76,7 +76,8 @@ function check_consistency(state::TearingState, ag = nothing)
 
     unassigned_var = []
     for (vj, eq) in enumerate(extended_var_eq_matching)
-        if eq === unassigned && (ag === nothing || !haskey(ag, vj))
+        if eq === unassigned && (ag === nothing || !haskey(ag, vj)) &&
+           !isempty(𝑑neighbors(graph, vj))
             push!(unassigned_var, fullvars[vj])
         end
     end

src/systems/abstractsystem.jl

@@ -1031,22 +1031,11 @@ This will convert all `inputs` to parameters and allow them to be unconnected, i
 simplification will allow models where `n_states = n_equations - n_inputs`.
 """
 function structural_simplify(sys::AbstractSystem, io = nothing; simplify = false,
-                             simplify_constants = true, check_consistency = true, kwargs...)
+                             kwargs...)
     sys = expand_connections(sys)
     sys isa DiscreteSystem && return sys
     state = TearingState(sys)
-    has_io = io !== nothing
-    has_io && markio!(state, io...)
-    state, input_idxs = inputs_to_parameters!(state, io)
-    sys, ag = alias_elimination!(state; kwargs...)
-    if check_consistency
-        ModelingToolkit.check_consistency(state, ag)
-    end
-    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
+    structural_simplify!(state, io; simplify, kwargs...)
 end
 
 function eliminate_constants(sys::AbstractSystem)
@@ -1063,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)
@@ -1150,8 +1139,8 @@ end
 
 function markio!(state, inputs, outputs; check = true)
     fullvars = state.fullvars
-    inputset = Dict(inputs .=> false)
-    outputset = Dict(outputs .=> false)
+    inputset = Dict{Any, Bool}(i => false for i in inputs)
+    outputset = Dict{Any, Bool}(o => false for o in outputs)
     for (i, v) in enumerate(fullvars)
         if v in keys(inputset)
             v = setio(v, true, false)
@@ -1166,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

@@ -28,7 +28,7 @@ function infer_clocks!(ci::ClockInference)
     @unpack ts, eq_domain, var_domain, inferred = ci
     @unpack fullvars = ts
     @unpack graph = ts.structure
-    # TODO: add a graph time to do this lazily
+    # TODO: add a graph type to do this lazily
     var_graph = SimpleGraph(ndsts(graph))
     for eq in 𝑠vertices(graph)
         vvs = 𝑠neighbors(graph, eq)
@@ -58,9 +58,97 @@ 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
 end
+
+function resize_or_push!(v, val, idx)
+    n = length(v)
+    if idx > n
+        for i in (n + 1):idx
+            push!(v, Int[])
+        end
+        resize!(v, idx)
+    end
+    push!(v[idx], val)
+end
+
+function split_system(ci::ClockInference)
+    @unpack ts, eq_domain, var_domain, inferred = ci
+    @unpack fullvars = ts
+    @unpack graph, var_to_diff = ts.structure
+    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_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
+
+            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
+        end
+        eq_to_cid[i] = cid
+        resize_or_push!(cid_to_eq, i, cid)
+    end
+    continuous_id = continuous_id[]
+    input_idxs = map(_ -> Int[], 1:cid_counter[])
+    inputs = map(_ -> Any[], 1:cid_counter[])
+    nvv = length(var_domain)
+    for i in 1:nvv
+        d = var_domain[i]
+        cid = get(clock_to_id, d, 0)
+        @assert cid!==0 "Internal error! Variable $(fullvars[i]) doesn't have a inferred time domain."
+        var_to_cid[i] = cid
+        v = fullvars[i]
+        #TODO: remove Inferred*
+        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
+
+    eqs = equations(ts)
+    tss = similar(cid_to_eq, TearingState)
+    for (id, ieqs) in enumerate(cid_to_eq)
+        vars = cid_to_var[id]
+        ts_i = ts
+        fadj = Vector{Int}[]
+        eqs_i = Equation[]
+        eq_to_diff = DiffGraph(length(ieqs))
+        ne = 0
+        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.structure.only_discrete = id != continuous_id
+        @set! ts_i.sys.eqs = eqs_i
+        @set! ts_i.structure.eq_to_diff = eq_to_diff
+        tss[id] = ts_i
+        # TODO: just mark current and sample variables as inputs
+    end
+    return tss, inputs
+
+    #id_to_clock, cid_to_eq, cid_to_var
+end

src/systems/systemstructure.jl

@@ -25,9 +25,9 @@ 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 isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq, algeqs, is_only_discrete
 export dervars_range, diffvars_range, algvars_range
 export DiffGraph, complete!
 
@@ -143,7 +143,9 @@ Base.@kwdef mutable struct SystemStructure
     # or as `torn` to assert that tearing has run.
     graph::BipartiteGraph{Int, Nothing}
     solvable_graph::Union{BipartiteGraph{Int, Nothing}, Nothing}
+    only_discrete::Bool
 end
+is_only_discrete(s::SystemStructure) = s.only_discrete
 isdervar(s::SystemStructure, i) = invview(s.var_to_diff)[i] !== nothing
 function isalgvar(s::SystemStructure, i)
     s.var_to_diff[i] === nothing &&
@@ -258,6 +260,27 @@ function TearingState(sys; quick_cancel = false, check = true)
                 idx = addvar!(dvar)
             end
 
+            dvar = var
+            idx = varidx
+            if ModelingToolkit.isoperator(dvar, ModelingToolkit.Shift)
+                if !(idx in dervaridxs)
+                    push!(dervaridxs, idx)
+                end
+                op = operation(dvar)
+                tt = op.t
+                steps = op.steps
+                v = arguments(dvar)[1]
+                for s in (steps - 1):-1:1
+                    sf = Shift(tt, s)
+                    dvar = sf(v)
+                    idx = addvar!(dvar)
+                    if !(idx in dervaridxs)
+                        push!(dervaridxs, idx)
+                    end
+                end
+                idx = addvar!(v)
+            end
+
             if istree(var) && operation(var) isa Symbolics.Operator &&
                !isdifferential(var) && (it = input_timedomain(var)) !== nothing
                 set_incidence = false
@@ -281,7 +304,7 @@ function TearingState(sys; quick_cancel = false, check = true)
     sorted_fullvars = OrderedSet(fullvars[dervaridxs])
     for dervaridx in dervaridxs
         dervar = fullvars[dervaridx]
-        diffvar = arguments(dervar)[1]
+        diffvar = lower_order_var(dervar)
         if !(diffvar in sorted_fullvars)
             push!(sorted_fullvars, diffvar)
         end
@@ -300,24 +323,12 @@ function TearingState(sys; quick_cancel = false, check = true)
     var_to_diff = DiffGraph(nvars, true)
     for dervaridx in dervaridxs
         dervar = fullvars[dervaridx]
-        diffvar = arguments(dervar)[1]
+        diffvar = lower_order_var(dervar)
         diffvaridx = var2idx[diffvar]
         push!(diffvars, diffvar)
         var_to_diff[diffvaridx] = dervaridx
     end
 
-    #=
-    algvars = setdiff(states(sys), diffvars)
-    for algvar in algvars
-        # it could be that a variable appeared in the states, but never appeared
-        # in the equations.
-        algvaridx = get(var2idx, algvar, 0)
-        #if algvaridx == 0
-        #    check ? throw(InvalidSystemException("The system is missing an equation for $algvar.")) : return nothing
-        #end
-    end
-    =#
-
     graph = BipartiteGraph(neqs, nvars, Val(false))
     for (ie, vars) in enumerate(symbolic_incidence), v in vars
         jv = var2idx[v]
@@ -330,7 +341,23 @@ function TearingState(sys; quick_cancel = false, check = true)
 
     return TearingState(sys, fullvars,
                         SystemStructure(complete(var_to_diff), complete(eq_to_diff),
-                                        complete(graph), nothing), Any[])
+                                        complete(graph), nothing, false), Any[])
+end
+
+function lower_order_var(dervar)
+    if isdifferential(dervar)
+        diffvar = arguments(dervar)[1]
+    else # shift
+        s = operation(dervar)
+        step = s.steps - 1
+        vv = arguments(dervar)[1]
+        if step >= 1
+            diffvar = Shift(s.t, step)(vv)
+        else
+            diffvar = vv
+        end
+    end
+    diffvar
 end
 
 using .BipartiteGraphs: Label, BipartiteAdjacencyList
@@ -424,4 +451,21 @@ 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,
+                              check_consistency = true, 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...)
+    if check_consistency
+        ModelingToolkit.check_consistency(state, ag)
+    end
+    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