Early work on the new discrete backend for MTK

Author: BenChung

src/systems/clock_inference.jl

@@ -100,7 +100,7 @@ function infer_clocks!(ci::ClockInference)
         c = BitSet(c′)
         idxs = intersect(c, inferred)
         isempty(idxs) && continue
-        if !allequal(var_domain[i] for i in idxs)
+        if !allequal(iscontinuous(var_domain[i]) for i in idxs)
             display(fullvars[c′])
             throw(ClockInferenceException("Clocks are not consistent in connected component $(fullvars[c′])"))
         end
@@ -155,6 +155,9 @@ function split_system(ci::ClockInference{S}) where {S}
     cid_to_var = Vector{Int}[]
     # cid_counter = number of clocks
     cid_counter = Ref(0)
+
+    # populates clock_to_id and id_to_clock
+    # checks if there is a continuous_id (for some reason? clock to id does this too)
     for (i, d) in enumerate(eq_domain)
         cid = let cid_counter = cid_counter, id_to_clock = id_to_clock,
             continuous_id = continuous_id
@@ -174,9 +177,13 @@ function split_system(ci::ClockInference{S}) where {S}
         resize_or_push!(cid_to_eq, i, cid)
     end
     continuous_id = continuous_id[]
+    # for each clock partition what are the input (indexes/vars)
     input_idxs = map(_ -> Int[], 1:cid_counter[])
     inputs = map(_ -> Any[], 1:cid_counter[])
+    # var_domain corresponds to fullvars/all variables in the system
     nvv = length(var_domain)
+    # put variables into the right clock partition
+    # keep track of inputs to each partition
     for i in 1:nvv
         d = var_domain[i]
         cid = get(clock_to_id, d, 0)
@@ -190,6 +197,7 @@ function split_system(ci::ClockInference{S}) where {S}
         resize_or_push!(cid_to_var, i, cid)
     end
 
+    # breaks the system up into a continous and 0 or more discrete systems
     tss = similar(cid_to_eq, S)
     for (id, ieqs) in enumerate(cid_to_eq)
         ts_i = system_subset(ts, ieqs)
@@ -199,6 +207,7 @@ function split_system(ci::ClockInference{S}) where {S}
         end
         tss[id] = ts_i
     end
+    # put the continous system at the back
     if continuous_id != 0
         tss[continuous_id], tss[end] = tss[end], tss[continuous_id]
         inputs[continuous_id], inputs[end] = inputs[end], inputs[continuous_id]

src/systems/systems.jl

@@ -36,7 +36,7 @@ function mtkcompile(
     isscheduled(sys) && throw(RepeatedStructuralSimplificationError())
     newsys′ = __mtkcompile(sys; simplify,
         allow_symbolic, allow_parameter, conservative, fully_determined,
-        inputs, outputs, disturbance_inputs,
+        inputs, outputs, disturbance_inputs, additional_passes,
         kwargs...)
     if newsys′ isa Tuple
         @assert length(newsys′) == 2
@@ -292,3 +292,8 @@ function map_variables_to_equations(sys::AbstractSystem; rename_dummy_derivative
 
     return mapping
 end
+
+"""
+Mark whether an extra pass `p` can support compiling discrete systems.
+"""
+discrete_compile_pass(p) = false

src/systems/systemstructure.jl

@@ -820,19 +820,40 @@ function mtkcompile!(state::TearingState; simplify = false,
     time_domains = merge(Dict(state.fullvars .=> ci.var_domain),
         Dict(default_toterm.(state.fullvars) .=> ci.var_domain))
     tss, clocked_inputs, continuous_id, id_to_clock = ModelingToolkit.split_system(ci)
+    if continuous_id == 0
+        # do a trait check here - handle fully discrete system
+        additional_passes = get(kwargs, :additional_passes, nothing)
+        if !isnothing(additional_passes) && any(discrete_compile_pass, additional_passes)
+            # take the first discrete compilation pass given for now
+            discrete_pass_idx = findfirst(discrete_compile_pass, additional_passes)
+            discrete_compile = additional_passes[discrete_pass_idx]
+            deleteat!(additional_passes, discrete_pass_idx)
+            return discrete_compile(tss, clocked_inputs)
+        end
+        throw(HybridSystemNotSupportedException("""
+        Discrete systems with multiple clocks are not supported with the standard \
+        MTK compiler.
+        """))
+    end
     if length(tss) > 1
-        if continuous_id == 0
-            throw(HybridSystemNotSupportedException("""
-            Discrete systems with multiple clocks are not supported with the standard \
-            MTK compiler.
-            """))
-        else
-            throw(HybridSystemNotSupportedException("""
-            Hybrid continuous-discrete systems are currently not supported with \
-            the standard MTK compiler. This system requires JuliaSimCompiler.jl, \
-            see https://help.juliahub.com/juliasimcompiler/stable/
-            """))
+        # simplify as normal
+        sys = _mtkcompile!(tss[continuous_id]; simplify,
+            inputs = [inputs; clocked_inputs[continuous_id]], outputs, disturbance_inputs,
+            check_consistency, fully_determined,
+            kwargs...)
+        if !isnothing(additional_passes) && any(discrete_compile_pass, additional_passes)
+            discrete_pass_idx = findfirst(discrete_compile_pass, additional_passes)
+            discrete_compile = additional_passes[discrete_pass_idx]
+            deleteat!(additional_passes, discrete_pass_idx)
+            # in the case of a hybrid system, the discrete_compile pass should take the currents of sys.discrete_subsystems
+            # and modifies discrete_subsystems to bea tuple of the io and anything else, while adding or manipulating the rest of sys as needed
+            return discrete_compile(sys, tss[2:end], inputs)
         end
+        throw(HybridSystemNotSupportedException("""
+        Hybrid continuous-discrete systems are currently not supported with \
+        the standard MTK compiler. This system requires JuliaSimCompiler.jl, \
+        see https://help.juliahub.com/juliasimcompiler/stable/
+        """))
     end
     if get_is_discrete(state.sys) ||
        continuous_id == 1 && any(Base.Fix2(isoperator, Shift), state.fullvars)