Merge pull request #2146 from SciML/myb/clock_generic

More generic clock split

Author: YingboMa

src/systems/clock_inference.jl

@@ -67,18 +67,18 @@ end
 function resize_or_push!(v, val, idx)
     n = length(v)
     if idx > n
-        for i in (n + 1):idx
+        for _ in (n + 1):idx
             push!(v, Int[])
         end
         resize!(v, idx)
     end
     push!(v[idx], val)
 end
 
-function split_system(ci::ClockInference)
+function split_system(ci::ClockInference{S}) where {S}
     @unpack ts, eq_domain, var_domain, inferred = ci
     fullvars = get_fullvars(ts)
-    @unpack graph, var_to_diff = ts.structure
+    @unpack graph = ts.structure
     continuous_id = Ref(0)
     clock_to_id = Dict{TimeDomain, Int}()
     id_to_clock = TimeDomain[]
@@ -121,26 +121,10 @@ function split_system(ci::ClockInference)
         resize_or_push!(cid_to_var, i, cid)
     end
 
-    eqs = equations(ts)
-    tss = similar(cid_to_eq, TearingState)
+    tss = similar(cid_to_eq, S)
     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)))
+        ts_i = system_subset(ts, ieqs)
         @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
     end
     return tss, inputs, continuous_id, id_to_clock
@@ -213,16 +197,20 @@ function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
                         @unpack u, p, t = integrator
                         c2d_obs = $cont_to_disc_obs
                         d2c_obs = $disc_to_cont_obs
+                        # Like Sample
                         c2d_view = view(p, $cont_to_disc_idxs)
+                        # Like Hold
                         d2c_view = view(p, $disc_to_cont_idxs)
                         disc_state = view(p, $disc_range)
                         disc = $disc
-                        # Write continuous info to discrete
-                        # Write discrete info to continuous
+                        # Write continuous into to discrete: handles `Sample`
                         copyto!(c2d_view, c2d_obs(integrator.u, p, t))
+                        # Write discrete into to continuous
+                        # get old discrete states
                         copyto!(d2c_view, d2c_obs(disc_state, p, t))
                         push!(saved_values.t, t)
                         push!(saved_values.saveval, $save_vec)
+                        # update discrete states
                         $empty_disc || disc(disc_state, disc_state, p, t)
                     end)
         sv = SavedValues(Float64, Vector{Float64})

src/systems/systemstructure.jl

@@ -31,7 +31,7 @@ export initialize_system_structure, find_linear_equations
 export isdiffvar, isdervar, isalgvar, isdiffeq, isalgeq, algeqs, is_only_discrete
 export dervars_range, diffvars_range, algvars_range
 export DiffGraph, complete!
-export get_fullvars
+export get_fullvars, system_subset
 
 struct DiffGraph <: Graphs.AbstractGraph{Int}
     primal_to_diff::Vector{Union{Int, Nothing}}
@@ -205,6 +205,28 @@ mutable struct TearingState{T <: AbstractSystem} <: AbstractTearingState{T}
 end
 
 TransformationState(sys::AbstractSystem) = TearingState(sys)
+function system_subset(ts::TearingState, ieqs::Vector{Int})
+    eqs = equations(ts)
+    @set! ts.sys.eqs = eqs[ieqs]
+    @set! ts.structure = system_subset(ts.structure, ieqs)
+    ts
+end
+
+function system_subset(structure::SystemStructure, ieqs::Vector{Int})
+    @unpack graph, eq_to_diff = structure
+    fadj = Vector{Int}[]
+    eq_to_diff = DiffGraph(length(ieqs))
+    ne = 0
+    for (j, eq_i) in enumerate(ieqs)
+        ivars = copy(graph.fadjlist[eq_i])
+        ne += length(ivars)
+        push!(fadj, ivars)
+        eq_to_diff[j] = structure.eq_to_diff[eq_i]
+    end
+    @set! structure.graph = complete(BipartiteGraph(ne, fadj, ndsts(graph)))
+    @set! structure.eq_to_diff = eq_to_diff
+    structure
+end
 
 function Base.show(io::IO, state::TearingState)
     print(io, "TearingState of ", typeof(state.sys))