Add special handling for discrete only systems

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

Author: YingboMa

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
@@ -631,7 +635,9 @@ function tearing_reassemble(state::TearingState, var_eq_matching, ag = nothing;
 
     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 && !isempty(𝑑neighbors(graph, 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/systems/clock_inference.jl

@@ -81,7 +81,7 @@ end
 function split_system(ci::ClockInference)
     @unpack ts, eq_domain, var_domain, inferred = ci
     @unpack fullvars = ts
-    @unpack graph = ts.structure
+    @unpack graph, var_to_diff = ts.structure
     continuous_id = Ref(0)
     clock_to_id = Dict{TimeDomain, Int}()
     id_to_clock = TimeDomain[]
@@ -106,11 +106,14 @@ function split_system(ci::ClockInference)
         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[])
-    for (i, d) in enumerate(var_domain)
+    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!"
+        @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*
@@ -139,10 +142,11 @@ function split_system(ci::ClockInference)
             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 past and sample variables as inputs
+        # TODO: just mark current and sample variables as inputs
     end
     return tss, inputs
 

src/systems/systemstructure.jl

@@ -27,7 +27,7 @@ end
 
 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 &&
@@ -277,9 +279,6 @@ function TearingState(sys; quick_cancel = false, check = true)
                     end
                 end
                 idx = addvar!(v)
-                #if !(idx in dervaridxs)
-                #    push!(dervaridxs, idx)
-                #end
             end
 
             if istree(var) && operation(var) isa Symbolics.Operator &&
@@ -305,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
@@ -324,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]
@@ -354,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

test/clock.jl

@@ -66,7 +66,8 @@ eqmap = ci.eq_domain
 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], ()))
+sss, = ModelingToolkit.structural_simplify!(deepcopy(tss[2]), (inputs[2], ()),
+                                            check_consistency = false)
 @test isempty(equations(sss))
 @test observed(sss) == [r ~ 1.0; yd ~ Sample(t, dt)(y); ud ~ kp * (r - yd)]
 
@@ -104,15 +105,21 @@ eqs = [yd ~ Sample(t, dt)(y)
        y ~ x
        z(k + 2) ~ z(k) + yd
        #=
-       z(k + 2) ~ z(k)
+       z(k + 2) ~ z(k) + yd
        =>
-       z′(k + 1) ~ z(k)
+       z′(k + 1) ~ z(k) + yd
        z(k + 1)  ~ z′(k)
        =#
        ]
 @named sys = ODESystem(eqs)
 ci, varmap = infer_clocks(sys)
 tss, inputs = ModelingToolkit.split_system(deepcopy(ci))
+sss, = ModelingToolkit.structural_simplify!(deepcopy(tss[2]), (inputs[2], ()),
+                                            check_consistency = false)
+@test length(states(sss)) == 2
+z, z_t = states(sss)
+S = Shift(t, 1)
+@test full_equations(sss) == [S(z) ~ z_t; S(z_t) ~ z + Sample(t, dt)(y)]
 
 @info "Testing multi-rate hybrid system"
 dt = 0.1