Merge pull request #1938 from SciML/myb_fb/clock_codegen

Add affect codegen for hybrid systems

Author: YingboMa

src/systems/abstractsystem.jl

@@ -373,7 +373,7 @@ function renamespace(sys, x)
     sys === nothing && return x
     x = unwrap(x)
     if x isa Symbolic
-        if isdifferential(x)
+        if istree(x) && operation(x) isa Operator
             return similarterm(x, operation(x), Any[renamespace(sys, only(arguments(x)))])
         end
         let scope = getmetadata(x, SymScope, LocalScope())

src/systems/clock_inference.jl

@@ -146,9 +146,71 @@ function split_system(ci::ClockInference)
         @set! ts_i.sys.eqs = eqs_i
         @set! ts_i.structure.eq_to_diff = eq_to_diff
         tss[id] = ts_i
-        # TODO: just mark current and sample variables as inputs
     end
-    return tss, inputs
+    return tss, inputs, continuous_id
+end
 
-    #id_to_clock, cid_to_eq, cid_to_var
+function generate_discrete_affect(syss, inputs, continuous_id, check_bounds = 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 = []
+    for (i, (sys, input)) in enumerate(zip(syss, inputs))
+        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_block = Let(assignments, let_body, false)
+        needed_cont_to_disc_obs = map(v -> arguments(v)[1], input)
+        # TODO: filter the needed ones
+        needed_disc_to_cont_obs = map(v -> arguments(v)[1], inputs[continuous_id])
+        append!(appended_parameters, input, states(sys))
+        disc_to_cont_idxs = map(Base.Fix1(getindex, param_to_idx), inputs[continuous_id])
+        cont_to_disc_obs = build_explicit_observed_function(syss[continuous_id],
+                                                            needed_cont_to_disc_obs,
+                                                            throw = false,
+                                                            expression = true,
+                                                            output_type = SVector)
+        @set! sys.ps = appended_parameters
+        disc_to_cont_obs = build_explicit_observed_function(sys, needed_disc_to_cont_obs,
+                                                            throw = false,
+                                                            expression = true,
+                                                            output_type = SVector)
+        ni = length(input)
+        ns = length(states(sys))
+        disc = Func([
+                        out,
+                        DestructuredArgs(states(sys)),
+                        DestructuredArgs(appended_parameters),
+                        get_iv(sys),
+                    ], [],
+                    let_block)
+        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
+        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
 end

src/systems/diffeqs/odesystem.jl

@@ -289,7 +289,8 @@ i.e. there are no cycles.
 function build_explicit_observed_function(sys, ts;
                                           expression = false,
                                           output_type = Array,
-                                          checkbounds = true)
+                                          checkbounds = true,
+                                          throw = true)
     if (isscalar = !(ts isa AbstractVector))
         ts = [ts]
     end
@@ -336,7 +337,12 @@ function build_explicit_observed_function(sys, ts;
                     subs[s] = s′
                     continue
                 end
-                throw(ArgumentError("$s is neither an observed nor a state variable."))
+                if throw
+                    Base.throw(ArgumentError("$s is neither an observed nor a state variable."))
+                else
+                    # TODO: return variables that don't exist in the system.
+                    return nothing
+                end
             end
             continue
         end

test/clock.jl

@@ -66,8 +66,7 @@ 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], ()),
-                                            check_consistency = false)
+sss, = ModelingToolkit.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)]
 
@@ -96,7 +95,7 @@ d = Clock(t, dt)
 k = ShiftIndex(d)
 
 eqs = [yd ~ Sample(t, dt)(y)
-       ud ~ kp * (r - yd)
+       ud ~ kp * (r - yd) + z(k)
        r ~ 1.0
 
        # plant (time continuous part)
@@ -114,11 +113,48 @@ eqs = [yd ~ Sample(t, dt)(y)
 @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], ()))
-@test length(states(sss)) == 2
-z, z_t = states(sss)
+syss = map(i -> ModelingToolkit.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(sss) == [S(z) ~ z_t; S(z_t) ~ z + Sample(t, dt)(y)]
+@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]);
+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]);
+cb = PeriodicCallback(gen_affect!, 0.1);
+sol2 = solve(prob, Tsit5(), callback = cb);
+
+# kp is the only real parameter
+function foo!(du, u, p, t)
+    x = u[1]
+    ud = p[2]
+    du[1] = -x + ud
+end
+function affect!(integrator, saved_values)
+    kp = integrator.p[1]
+    yd = integrator.u[1]
+    z_t = integrator.p[3]
+    z = integrator.p[4]
+    r = 1.0
+    ud = kp * (r - yd) + z
+    push!(saved_values.t, integrator.t)
+    push!(saved_values.saveval, (integrator.p[3], integrator.p[4]))
+    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
 
 @info "Testing multi-rate hybrid system"
 dt = 0.1