Better API for discrete controller simulation

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

Author: YingboMa

src/systems/abstractsystem.jl

@@ -208,7 +208,8 @@ for prop in [:eqs
              :torn_matching
              :tearing_state
              :substitutions
-             :metadata]
+             :metadata
+             :discrete_subsystems]
     fname1 = Symbol(:get_, prop)
     fname2 = Symbol(:has_, prop)
     @eval begin

src/systems/clock_inference.jl

@@ -147,19 +147,22 @@ function split_system(ci::ClockInference)
         @set! ts_i.structure.eq_to_diff = eq_to_diff
         tss[id] = ts_i
     end
-    return tss, inputs, continuous_id
+    return tss, inputs, continuous_id, id_to_clock
 end
 
-function generate_discrete_affect(syss, inputs, continuous_id; checkbounds = true,
+function generate_discrete_affect(syss, inputs, continuous_id, id_to_clock;
+                                  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))
     offset = length(appended_parameters)
     affect_funs = []
     svs = []
+    clocks = TimeDomain[]
     for (i, (sys, input)) in enumerate(zip(syss, inputs))
         i == continuous_id && continue
+        push!(clocks, id_to_clock[i])
         subs = get_substitutions(sys)
         assignments = map(s -> Assignment(s.lhs, s.rhs), subs.subs)
         let_body = SetArray(!checkbounds, out, rhss(equations(sys)))
@@ -191,9 +194,9 @@ function generate_discrete_affect(syss, inputs, continuous_id; checkbounds = tru
         cont_to_disc_idxs = (offset + 1):(offset += ni)
         input_offset = offset
         disc_range = (offset + 1):(offset += ns)
-        save_tuple = Expr(:tuple)
+        save_vec = Expr(:ref, :Float64)
         for i in 1:ns
-            push!(save_tuple.args, :(p[$(input_offset + i)]))
+            push!(save_vec.args, :(p[$(input_offset + i)]))
         end
         affect! = :(function (integrator, saved_values)
                         @unpack u, p, t = integrator
@@ -208,10 +211,10 @@ function generate_discrete_affect(syss, inputs, continuous_id; checkbounds = tru
                         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)
+                        push!(saved_values.saveval, $save_vec)
                         disc(disc_state, disc_state, p, t)
                     end)
-        sv = SavedValues(Float64, NTuple{ns, Float64})
+        sv = SavedValues(Float64, Vector{Float64})
         push!(affect_funs, affect!)
         push!(svs, sv)
     end
@@ -223,5 +226,5 @@ function generate_discrete_affect(syss, inputs, continuous_id; checkbounds = tru
         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
+    return affects, clocks, svs, appended_parameters, defaults
 end

src/systems/diffeqs/abstractodesystem.jl

@@ -686,6 +686,12 @@ function DiffEqBase.ODEProblem{false}(sys::AbstractODESystem, args...; kwargs...
     ODEProblem{false, SciMLBase.FullSpecialize}(sys, args...; kwargs...)
 end
 
+struct DiscreteSaveAffect{F, S} <: Function
+    f::F
+    s::S
+end
+(d::DiscreteSaveAffect)(args...) = d.f(args..., d.s)
+
 function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map = [],
                                                 tspan = get_tspan(sys),
                                                 parammap = DiffEqBase.NullParameters();
@@ -698,13 +704,35 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
                                  has_difference = has_difference,
                                  check_length, kwargs...)
     cbs = process_events(sys; callback, has_difference, kwargs...)
+    if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
+        affects, clocks, svs = ModelingToolkit.generate_discrete_affect(dss...)
+        discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
+            if clock isa Clock
+                PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
+            else
+                error("$clock is not a supported clock type.")
+            end
+        end
+        if cbs === nothing
+            if length(discrete_cbs) == 1
+                cbs = only(discrete_cbs)
+            else
+                cbs = CallbackSet(discrete_cbs...)
+            end
+        else
+            cbs = CallbackSet(cbs, discrete_cbs)
+        end
+    else
+        svs = nothing
+    end
     kwargs = filter_kwargs(kwargs)
     pt = something(get_metadata(sys), StandardODEProblem())
 
     if cbs === nothing
-        ODEProblem{iip}(f, u0, tspan, p, pt; kwargs...)
+        ODEProblem{iip}(f, u0, tspan, p, pt; disc_saved_values = svs, kwargs...)
     else
-        ODEProblem{iip}(f, u0, tspan, p, pt; callback = cbs, kwargs...)
+        ODEProblem{iip}(f, u0, tspan, p, pt; callback = cbs, disc_saved_values = svs,
+                        kwargs...)
     end
 end
 get_callback(prob::ODEProblem) = prob.kwargs[:callback]

src/systems/diffeqs/odesystem.jl

@@ -129,7 +129,7 @@ struct ODESystem <: AbstractODESystem
     """
     discrete_subsystems: a list of discrete subsystems
     """
-    discrete_subsystems::Union{Nothing, Tuple{Vector{ODESystem}, Vector{Any}, Int}}
+    discrete_subsystems::Any
 
     function ODESystem(tag, deqs, iv, dvs, ps, tspan, var_to_name, ctrls, observed, tgrad,
                        jac, ctrl_jac, Wfact, Wfact_t, name, systems, defaults,

src/systems/systemstructure.jl

@@ -467,24 +467,33 @@ function structural_simplify!(state::TearingState, io = nothing; simplify = fals
     if state.sys isa ODESystem
         ci = ModelingToolkit.ClockInference(state)
         ModelingToolkit.infer_clocks!(ci)
-        tss, inputs, continuous_id = ModelingToolkit.split_system(ci)
+        tss, inputs, continuous_id, id_to_clock = ModelingToolkit.split_system(ci)
         cont_io = merge_io(io, inputs[continuous_id])
-        sys, input_idxs = _structural_simplify!(tss[continuous_id], cont_io; simplify, check_consistency,
-                                    kwargs...)
+        sys, input_idxs = _structural_simplify!(tss[continuous_id], cont_io; simplify,
+                                                check_consistency,
+                                                kwargs...)
         if length(tss) > 1
             # TODO: rename it to something else
             discrete_subsystems = Vector{ODESystem}(undef, length(tss))
+            # Note that the appended_parameters must agree with
+            # `generate_discrete_affect`!
+            appended_parameters = parameters(sys)
             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)
+                ss, = _structural_simplify!(state, dist_io; simplify, check_consistency,
+                                            kwargs...)
+                append!(appended_parameters, inputs[i], states(ss))
+                discrete_subsystems[i] = ss
             end
-            @set! sys.discrete_subsystems = discrete_subsystems, inputs, continuous_id
+            @set! sys.discrete_subsystems = discrete_subsystems, inputs, continuous_id,
+                                            id_to_clock
+            @set! sys.ps = appended_parameters
+            @set! sys.defaults = merge(ModelingToolkit.defaults(sys),
+                                       Dict(v => 0.0 for v in Iterators.flatten(inputs)))
         end
     else
         sys, input_idxs = _structural_simplify!(state, io; simplify, check_consistency,

test/clock.jl

@@ -112,30 +112,9 @@ eqs = [yd ~ Sample(t, dt)(y)
        =#
        ]
 @named sys = ODESystem(eqs)
-ci, varmap = infer_clocks(sys)
-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)
-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
-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);
-prob = remake(prob, callback = cb);
-sol2 = solve(prob, Tsit5());
+ss = structural_simplify(sys)
+prob = ODEProblem(ss, [x => 0.0, y => 0.0], (0.0, 1.0), [kp => 1.0; z => 0.0; D(z) => 0.0])
+sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
 # For all inputs in parameters, just initialize them to 0.0, and then set them
 # in the callback.
 
@@ -153,17 +132,19 @@ function affect!(integrator, saved_values)
     r = 1.0
     ud = kp * (r - yd) + z
     push!(saved_values.t, integrator.t)
-    push!(saved_values.saveval, (integrator.p[3], integrator.p[4]))
+    push!(saved_values.saveval, [integrator.p[4], integrator.p[3]])
     integrator.p[2] = ud
     integrator.p[3] = z + yd
     integrator.p[4] = z_t
     nothing
 end
-saved_values = SavedValues(Float64, Tuple{Float64, Float64});
-cb = PeriodicCallback(Base.Fix2(affect!, saved_values), 0.1);
-prob = ODEProblem(foo!, [0.0], (0.0, 1.0), [1.0, 0.0, 0.0, 0.0], callback = cb);
-sol = solve(prob, Tsit5());
-@test sol.u ≈ sol2.u
+saved_values = SavedValues(Float64, Vector{Float64});
+cb = PeriodicCallback(Base.Fix2(affect!, saved_values), 0.1)
+prob = ODEProblem(foo!, [0.0], (0.0, 1.0), [1.0, 0.0, 0.0, 0.0], callback = cb)
+sol2 = solve(prob, Tsit5())
+@test sol.u == sol2.u
+@test saved_values.t == sol.prob.kwargs[:disc_saved_values][1].t
+@test saved_values.saveval == sol.prob.kwargs[:disc_saved_values][1].saveval
 
 @info "Testing multi-rate hybrid system"
 dt = 0.1