WIP: work toward merging clock processing with the common interface

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

Author: YingboMa

src/systems/clock_inference.jl

@@ -150,7 +150,8 @@ function split_system(ci::ClockInference)
     return tss, inputs, continuous_id
 end
 
-function generate_discrete_affect(syss, inputs, continuous_id, check_bounds = true)
+function generate_discrete_affect(syss, inputs, continuous_id; checkbounds = true,
+                                  eval_module = @__MODULE__, eval_expression = true)
     out = Sym{Any}(:out)
     appended_parameters = parameters(syss[continuous_id])
     param_to_idx = Dict{Any, Int}(reverse(en) for en in enumerate(appended_parameters))
@@ -161,7 +162,7 @@ function generate_discrete_affect(syss, inputs, continuous_id, check_bounds = tr
         i == continuous_id && continue
         subs = get_substitutions(sys)
         assignments = map(s -> Assignment(s.lhs, s.rhs), subs.subs)
-        let_body = SetArray(!check_bounds, out, rhss(equations(sys)))
+        let_body = SetArray(!checkbounds, out, rhss(equations(sys)))
         let_block = Let(assignments, let_body, false)
         needed_cont_to_disc_obs = map(v -> arguments(v)[1], input)
         # TODO: filter the needed ones
@@ -190,27 +191,37 @@ function generate_discrete_affect(syss, inputs, continuous_id, check_bounds = tr
         cont_to_disc_idxs = (offset + 1):(offset += ni)
         input_offset = offset
         disc_range = (offset + 1):(offset += ns)
-        affect! = quote
-            function affect!(integrator, saved_values)
-                @unpack u, p, t = integrator
-                c2d_obs = $cont_to_disc_obs
-                d2c_obs = $disc_to_cont_obs
-                c2d_view = view(p, $cont_to_disc_idxs)
-                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
-                copyto!(c2d_view, c2d_obs(integrator.u, p, t))
-                copyto!(d2c_view, d2c_obs(disc_state, p, t))
-                push!(saved_values.t, t)
-                push!(saved_values.saveval, Base.@ntuple $ns i->p[$input_offset + i])
-                disc(disc_state, disc_state, p, t)
-            end
+        save_tuple = Expr(:tuple)
+        for i in 1:ns
+            push!(save_tuple.args, :(p[$(input_offset + i)]))
         end
+        affect! = :(function (integrator, saved_values)
+                        @unpack u, p, t = integrator
+                        c2d_obs = $cont_to_disc_obs
+                        d2c_obs = $disc_to_cont_obs
+                        c2d_view = view(p, $cont_to_disc_idxs)
+                        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
+                        copyto!(c2d_view, c2d_obs(integrator.u, p, t))
+                        copyto!(d2c_view, d2c_obs(disc_state, p, t))
+                        push!(saved_values.t, t)
+                        push!(saved_values.saveval, $save_tuple)
+                        disc(disc_state, disc_state, p, t)
+                    end)
         sv = SavedValues(Float64, NTuple{ns, Float64})
         push!(affect_funs, affect!)
         push!(svs, sv)
     end
-    return map(a -> toexpr(LiteralExpr(a)), affect_funs), svs, appended_parameters
+    if eval_expression
+        affects = map(affect_funs) do a
+            @RuntimeGeneratedFunction(eval_module, toexpr(LiteralExpr(a)))
+        end
+    else
+        affects = map(a -> toexpr(LiteralExpr(a)), affect_funs)
+    end
+    defaults = Dict{Any, Any}(v => 0.0 for v in Iterators.flatten(inputs))
+    return affects, svs, appended_parameters, defaults
 end

src/systems/diffeqs/odesystem.jl

@@ -126,13 +126,17 @@ struct ODESystem <: AbstractODESystem
     complete: if a model `sys` is complete, then `sys.x` no longer performs namespacing.
     """
     complete::Bool
+    """
+    discrete_subsystems: a list of discrete subsystems
+    """
+    discrete_subsystems::Union{Nothing, Tuple{Vector{ODESystem}, Vector{Any}, Int}}
 
     function ODESystem(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad,
                        jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults,
                        torn_matching, connector_type, preface, cevents,
                        devents, metadata = nothing, tearing_state = nothing,
-                       substitutions = nothing, complete = false;
-                       checks::Union{Bool, Int} = true)
+                       substitutions = nothing, complete = false,
+                       discrete_subsystems = nothing; checks::Union{Bool, Int} = true)
         if checks == true || (checks & CheckComponents) > 0
             check_variables(dvs, iv)
             check_parameters(ps, iv)
@@ -145,7 +149,7 @@ struct ODESystem <: AbstractODESystem
         new(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad, jac,
             ctrl_jac, Wfact, Wfact_t, name, systems, defaults, torn_matching,
             connector_type, preface, cevents, devents, metadata, tearing_state,
-            substitutions, complete)
+            substitutions, complete, discrete_subsystems)
     end
 end
 

src/systems/systemstructure.jl

@@ -452,8 +452,49 @@ function Base.show(io::IO, mime::MIME"text/plain", ms::MatchedSystemStructure)
 end
 
 # TODO: clean up
+function merge_io(io, inputs)
+    isempty(inputs) && return io
+    if io === nothing
+        io = (inputs, [])
+    else
+        io = ([inputs; io[1]], io[2])
+    end
+    return io
+end
+
 function structural_simplify!(state::TearingState, io = nothing; simplify = false,
                               check_consistency = true, kwargs...)
+    if state.sys isa ODESystem
+        ci = ModelingToolkit.ClockInference(state)
+        ModelingToolkit.infer_clocks!(ci)
+        tss, inputs, continuous_id = ModelingToolkit.split_system(ci)
+        cont_io = merge_io(io, inputs[continuous_id])
+        sys = _structural_simplify!(tss[continous_id], cont_io; simplify, check_consistency,
+                                    kwargs...)
+        if length(tss) > 1
+            # TODO: rename it to something else
+            discrete_subsystems = Vector{ODESystem}(undef, length(tss))
+            for (i, state) in enumerate(tss)
+                if i == continuous_id
+                    discrete_subsystems[i] = sys
+                    continue
+                end
+                dist_io = merge_io(io, inputs[i])
+                ss = _structural_simplify!(state, dist_io; simplify, check_consistency,
+                                           kwargs...)
+                push!(discrete_subsystems, ss)
+            end
+            @set! sys.discrete_subsystems = discrete_subsystems, inputs, continuous_id
+        end
+    else
+        sys, input_idxs = _structural_simplify!(state, io; simplify, check_consistency,
+                                                kwargs...)
+    end
+    return has_io ? (sys, input_idxs) : sys
+end
+
+function _structural_simplify!(state::TearingState, io; 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)
@@ -464,8 +505,7 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
     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
+    ModelingToolkit.invalidate_cache!(sys), input_idxs
 end
 
 end # module

test/clock.jl

@@ -61,12 +61,13 @@ By inference:
 =>   Shift(x, 0, dt) := (Shift(x, -1, dt) + dt) / (1 - dt) # Discrete system
 =#
 
+using ModelingToolkit.SystemStructures
 ci, varmap = infer_clocks(sys)
 eqmap = ci.eq_domain
 tss, inputs = ModelingToolkit.split_system(deepcopy(ci))
-sss, = ModelingToolkit.structural_simplify!(deepcopy(tss[1]), (inputs[1], ()))
+sss, = SystemStructures._structural_simplify!(deepcopy(tss[1]), (inputs[1], ()))
 @test equations(sss) == [D(x) ~ u - x]
-sss, = ModelingToolkit.structural_simplify!(deepcopy(tss[2]), (inputs[2], ()))
+sss, = SystemStructures._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)]
 
@@ -112,23 +113,31 @@ eqs = [yd ~ Sample(t, dt)(y)
        ]
 @named sys = ODESystem(eqs)
 ci, varmap = infer_clocks(sys)
-tss, inputs = ModelingToolkit.split_system(deepcopy(ci))
-syss = map(i -> ModelingToolkit.structural_simplify!(deepcopy(tss[i]), (inputs[i], ()))[1],
+tss, inputs, continuous_id = ModelingToolkit.split_system(deepcopy(ci))
+syss = map(i -> SystemStructures._structural_simplify!(deepcopy(tss[i]), (inputs[i], ()))[1],
            eachindex(tss))
 sys1, sys2 = syss
 @test length(states(sys2)) == 2
 z, z_t = states(sys2)
 S = Shift(t, 1)
 @test full_equations(sys2) == [S(z) ~ z_t; S(z_t) ~ z + Sample(t, dt)(y)]
 # TODO: set Hold(ud)
-prob = ODEProblem(sys1, [x => 0.0, y => 0.0], (0.0, 1.0), [kp => 1.0, Hold(ud) => 0.0]);
+affects, svs, pp, defaults = ModelingToolkit.generate_discrete_affect(syss, inputs,
+                                                                      continuous_id);
+@set! sys1.ps = pp
+prob = ODEProblem(sys1, [x => 0.0, y => 0.0], (0.0, 1.0), [pp .=> 0.0; kp => 1.0]);
 using OrdinaryDiffEq, DiffEqCallbacks
-exprs, svs, pp = ModelingToolkit.generate_discrete_affect(syss, inputs, 1);
-prob = remake(prob, p = zeros(Float64, length(pp)));
-prob.p[1] = 1.0;
-gen_affect! = Base.Fix2(eval(exprs[1]), svs[1]);
+struct DiscreteSaveAffect{F, S} <: Function
+    f::F
+    s::S
+end
+(d::DiscreteSaveAffect)(args...) = d.f(args..., d.s)
+gen_affect! = DiscreteSaveAffect(affects[1], svs[1]);
 cb = PeriodicCallback(gen_affect!, 0.1);
-sol2 = solve(prob, Tsit5(), callback = cb);
+prob = remake(prob, callback = cb);
+sol2 = solve(prob, Tsit5());
+# For all inputs in parameters, just initialize them to 0.0, and then set them
+# in the callback.
 
 # kp is the only real parameter
 function foo!(du, u, p, t)