fix: run affect! for hybrid systems during problem initialization

Author: AayushSabharwal

src/systems/clock_inference.jl

@@ -157,7 +157,6 @@ function generate_discrete_affect(
         Dict{Any, Int}(reverse(en) for en in enumerate(appended_parameters))
     end
     affect_funs = []
-    init_funs = []
     svs = []
     clocks = TimeDomain[]
     for (i, (sys, input)) in enumerate(zip(syss, inputs))
@@ -249,27 +248,6 @@ function generate_discrete_affect(
             end
         end
         empty_disc = isempty(disc_range)
-        disc_init = if use_index_cache
-            :(function (p, t)
-                d2c_obs = $disc_to_cont_obs
-                disc_state = Tuple($(parameter_values)(p, i) for i in $disc_range)
-                result = d2c_obs(disc_state, p..., t)
-                for (val, i) in zip(result, $disc_to_cont_idxs)
-                    # prevent multiple updates to dependents
-                    _set_parameter_unchecked!(p, val, i; update_dependent = false)
-                end
-                discretes, repack, _ = $(SciMLStructures.canonicalize)(
-                    $(SciMLStructures.Discrete()), p)
-                repack(discretes) # to force recalculation of dependents
-            end)
-        else
-            :(function (p, t)
-                d2c_obs = $disc_to_cont_obs
-                d2c_view = view(p, $disc_to_cont_idxs)
-                disc_state = view(p, $disc_range)
-                copyto!(d2c_view, d2c_obs(disc_state, p, t))
-            end)
-        end
 
         # @show disc_to_cont_idxs
         # @show cont_to_disc_idxs
@@ -357,20 +335,15 @@ function generate_discrete_affect(
         end)
         sv = SavedValues(Float64, Vector{Float64})
         push!(affect_funs, affect!)
-        push!(init_funs, disc_init)
         push!(svs, sv)
     end
     if eval_expression
         affects = map(affect_funs) do a
             drop_expr(@RuntimeGeneratedFunction(eval_module, toexpr(LiteralExpr(a))))
         end
-        inits = map(init_funs) do a
-            drop_expr(@RuntimeGeneratedFunction(eval_module, toexpr(LiteralExpr(a))))
-        end
     else
         affects = map(a -> toexpr(LiteralExpr(a)), affect_funs)
-        inits = map(a -> toexpr(LiteralExpr(a)), init_funs)
     end
     defaults = Dict{Any, Any}(v => 0.0 for v in Iterators.flatten(inputs))
-    return affects, inits, clocks, svs, appended_parameters, defaults
+    return affects, clocks, svs, appended_parameters, defaults
 end

src/systems/diffeqs/abstractodesystem.jl

@@ -1039,9 +1039,9 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
         t = tspan !== nothing ? tspan[1] : tspan,
         check_length, warn_initialize_determined, kwargs...)
     cbs = process_events(sys; callback, kwargs...)
-    inits = []
+    affects = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
+        affects, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
         discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
             if clock isa Clock
                 PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt)
@@ -1062,6 +1062,12 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
         else
             cbs = CallbackSet(cbs, discrete_cbs...)
         end
+        # initialize by running affects
+        dummy_saveval = (; t = [], saveval = [])
+        for affect! in affects
+            affect!(
+                (; u = u0, p = p, t = tspan !== nothing ? tspan[1] : tspan), dummy_saveval)
+        end
     else
         svs = nothing
     end
@@ -1075,13 +1081,7 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
     if svs !== nothing
         kwargs1 = merge(kwargs1, (disc_saved_values = svs,))
     end
-    prob = ODEProblem{iip}(f, u0, tspan, p, pt; kwargs1..., kwargs...)
-    if !isempty(inits)
-        for init in inits
-            init(prob.p, tspan[1])
-        end
-    end
-    prob
+    ODEProblem{iip}(f, u0, tspan, p, pt; kwargs1..., kwargs...)
 end
 get_callback(prob::ODEProblem) = prob.kwargs[:callback]
 

test/clock.jl

@@ -105,15 +105,15 @@ eqs = [yd ~ Sample(t, dt)(y)
 ss = structural_simplify(sys);
 
 Tf = 1.0
-prob = ODEProblem(ss, [x => 0.0, y => 0.0], (0.0, Tf),
+prob = ODEProblem(ss, [x => 0.1], (0.0, Tf),
     [kp => 1.0; ud(k - 1) => 2.0; ud(k - 2) => 2.0])
-@test sort(vcat(prob.p...)) == [0, 1.0, 2.0, 2.0, 2.0] # yd, Hold(ud), kp, ud(k - 1)
+@test sort(vcat(prob.p...)) == [0.1, 1.0, 2.0, 2.1, 2.1] # yd, kp, ud(k-2), ud(k-1), Hold(ud)
 sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
 
 ss_nosplit = structural_simplify(sys; split = false)
-prob_nosplit = ODEProblem(ss_nosplit, [x => 0.0, y => 0.0], (0.0, Tf),
+prob_nosplit = ODEProblem(ss_nosplit, [x => 0.1], (0.0, Tf),
     [kp => 1.0; ud(k - 1) => 2.0; ud(k - 2) => 2.0])
-@test sort(prob_nosplit.p) == [0, 1.0, 2.0, 2.0, 2.0] # yd, Hold(ud), kp, ud(k - 1)
+@test sort(prob_nosplit.p) == [0.1, 1.0, 2.0, 2.1, 2.1] # yd, kp, ud(k-2), ud(k-1), Hold(ud)
 sol_nosplit = solve(prob_nosplit, Tsit5(), kwargshandle = KeywordArgSilent)
 # For all inputs in parameters, just initialize them to 0.0, and then set them
 # in the callback.
@@ -141,7 +141,7 @@ end
 saved_values = SavedValues(Float64, Vector{Float64})
 cb = PeriodicCallback(Base.Fix2(affect!, saved_values), 0.1)
 #                                           kp   ud
-prob = ODEProblem(foo!, [0.0], (0.0, Tf), [1.0, 2.0, 2.0], callback = cb)
+prob = ODEProblem(foo!, [0.1], (0.0, Tf), [1.0, 2.1, 2.0], callback = cb)
 sol2 = solve(prob, Tsit5())
 @test sol.u == sol2.u
 @test sol_nosplit.u == sol2.u
@@ -433,8 +433,8 @@ prob = ODEProblem(ssys,
     [model.plant.x => 0.0; model.controller.kp => 2.0; model.controller.ki => 2.0],
     (0.0, Tf))
 
-@test prob.ps[Hold(ssys.holder.input.u)] == 1 # constant output * kp issue https://github.com/SciML/ModelingToolkit.jl/issues/2356
-@test prob.ps[ssys.controller.x(k - 1)] == 0 # c2d
+@test prob.ps[Hold(ssys.holder.input.u)] == 2 # constant output * kp issue https://github.com/SciML/ModelingToolkit.jl/issues/2356
+@test prob.ps[ssys.controller.x(k - 1)] == 1 # c2d
 @test prob.ps[Sample(d)(ssys.sampler.input.u)] == 0 # disc state
 sol = solve(prob,
     Tsit5(),