feat: use new discrete saving, only allow split=true hybrid systems

Author: AayushSabharwal

src/systems/clock_inference.jl

@@ -200,19 +200,14 @@ function generate_discrete_affect(
     @static if VERSION < v"1.7"
         error("The `generate_discrete_affect` function requires at least Julia 1.7")
     end
-    use_index_cache = has_index_cache(osys) && get_index_cache(osys) !== nothing
+    has_index_cache(osys) && get_index_cache(osys) !== nothing ||
+        error("Hybrid systems require `split = true`")
     out = Sym{Any}(:out)
     appended_parameters = full_parameters(syss[continuous_id])
     offset = length(appended_parameters)
-    param_to_idx = if use_index_cache
-        Dict{Any, ParameterIndex}(p => parameter_index(osys, p)
-        for p in appended_parameters)
-    else
-        Dict{Any, Int}(reverse(en) for en in enumerate(appended_parameters))
-    end
+    param_to_idx = Dict{Any, ParameterIndex}(p => parameter_index(osys, p)
+    for p in appended_parameters)
     affect_funs = []
-    init_funs = []
-    svs = []
     clocks = TimeDomain[]
     for (i, (sys, input)) in enumerate(zip(syss, inputs))
         i == continuous_id && continue
@@ -228,11 +223,7 @@ function generate_discrete_affect(
             push!(fullvars, s)
         end
         needed_disc_to_cont_obs = []
-        if use_index_cache
-            disc_to_cont_idxs = ParameterIndex[]
-        else
-            disc_to_cont_idxs = Int[]
-        end
+        disc_to_cont_idxs = ParameterIndex[]
         for v in inputs[continuous_id]
             _v = arguments(v)[1]
             if _v in fullvars
@@ -252,7 +243,7 @@ function generate_discrete_affect(
         end
         append!(appended_parameters, input)
         cont_to_disc_obs = build_explicit_observed_function(
-            use_index_cache ? osys : syss[continuous_id],
+            osys,
             needed_cont_to_disc_obs,
             throw = false,
             expression = true,
@@ -274,76 +265,20 @@ function generate_discrete_affect(
             ],
             [],
             let_block) |> toexpr
-        if use_index_cache
-            cont_to_disc_idxs = [parameter_index(osys, sym) for sym in input]
-            disc_range = [parameter_index(osys, sym) for sym in unknowns(sys)]
-        else
-            cont_to_disc_idxs = (offset + 1):(offset += ni)
-            input_offset = offset
-            disc_range = (offset + 1):(offset += ns)
-        end
-        save_vec = Expr(:ref, :Float64)
-        if use_index_cache
-            for unk in unknowns(sys)
-                idx = parameter_index(osys, unk)
-                push!(save_vec.args, :($(parameter_values)(p, $idx)))
-            end
-        else
-            for i in 1:ns
-                push!(save_vec.args, :(p[$(input_offset + i)]))
-            end
-        end
+        cont_to_disc_idxs = [parameter_index(osys, sym) for sym in input]
+        disc_range = [parameter_index(osys, sym) for sym in unknowns(sys)]
+        save_expr = :($(SciMLBase.save_discretes!)(integrator, $(i - 1)))
         empty_disc = isempty(disc_range)
-        disc_init = if use_index_cache
-            :(function (u, p, t)
-                c2d_obs = $cont_to_disc_obs
-                d2c_obs = $disc_to_cont_obs
-                result = c2d_obs(u, p..., t)
-                for (val, i) in zip(result, $cont_to_disc_idxs)
-                    $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
-                end
-
-                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 (u, p, t)
-                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_unknowns = view(p, $disc_range)
-                copyto!(c2d_view, c2d_obs(u, p, t))
-                copyto!(d2c_view, d2c_obs(disc_unknowns, p, t))
-            end)
-        end
 
         # @show disc_to_cont_idxs
         # @show cont_to_disc_idxs
         # @show disc_range
-        affect! = :(function (integrator, saved_values)
+        affect! = :(function (integrator)
             @unpack u, p, t = integrator
             c2d_obs = $cont_to_disc_obs
             d2c_obs = $disc_to_cont_obs
-            $(
-                if use_index_cache
-                :(disc_unknowns = [$(parameter_values)(p, i) for i in $disc_range])
-            else
-                quote
-                    c2d_view = view(p, $cont_to_disc_idxs)
-                    d2c_view = view(p, $disc_to_cont_idxs)
-                    disc_unknowns = view(p, $disc_range)
-                end
-            end
-            )
             # TODO: find a way to do this without allocating
+            disc_unknowns = [$(parameter_values)(p, i) for i in $disc_range]
             disc = $disc
 
             # Write continuous into to discrete: handles `Sample`
@@ -355,77 +290,41 @@ function generate_discrete_affect(
             # d2c comes last
             # @show t
             # @show "incoming", p
-            $(
-                if use_index_cache
+            result = c2d_obs(integrator.u, p..., t)
+            for (val, i) in zip(result, $cont_to_disc_idxs)
+                $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
+            end
+            $(if !empty_disc
                 quote
-                    result = c2d_obs(integrator.u, p..., t)
-                    for (val, i) in zip(result, $cont_to_disc_idxs)
+                    disc(disc_unknowns, integrator.u, p..., t)
+                    for (val, i) in zip(disc_unknowns, $disc_range)
                         $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
                     end
                 end
-            else
-                :(copyto!(c2d_view, c2d_obs(integrator.u, p, t)))
-            end
-            )
+            end)
             # @show "after c2d", p
-            $(
-                if use_index_cache
-                quote
-                    if !$empty_disc
-                        disc(disc_unknowns, integrator.u, p..., t)
-                        for (val, i) in zip(disc_unknowns, $disc_range)
-                            $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
-                        end
-                    end
-                end
-            else
-                :($empty_disc || disc(disc_unknowns, disc_unknowns, p, t))
-            end
-            )
             # @show "after state update", p
-            $(
-                if use_index_cache
-                quote
-                    result = d2c_obs(disc_unknowns, p..., t)
-                    for (val, i) in zip(result, $disc_to_cont_idxs)
-                        $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
-                    end
-                end
-            else
-                :(copyto!(d2c_view, d2c_obs(disc_unknowns, p, t)))
+            result = d2c_obs(disc_unknowns, p..., t)
+            for (val, i) in zip(result, $disc_to_cont_idxs)
+                $(_set_parameter_unchecked!)(p, val, i; update_dependent = false)
             end
-            )
 
-            push!(saved_values.t, t)
-            push!(saved_values.saveval, $save_vec)
+            $save_expr
 
             # @show "after d2c", p
-            $(
-                if use_index_cache
-                quote
-                    discretes, repack, _ = $(SciMLStructures.canonicalize)(
-                        $(SciMLStructures.Discrete()), p)
-                    repack(discretes)
-                end
-            end
-            )
+            discretes, repack, _ = $(SciMLStructures.canonicalize)(
+                $(SciMLStructures.Discrete()), p)
+            repack(discretes)
         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, appended_parameters, defaults
 end

src/systems/diffeqs/abstractodesystem.jl

@@ -1086,15 +1086,13 @@ 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 = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(sys, dss...)
-        discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
+        affects, clocks = ModelingToolkit.generate_discrete_affect(sys, dss...)
+        discrete_cbs = map(affects, clocks) do affect, clock
             if clock isa Clock
-                PeriodicCallback(DiscreteSaveAffect(affect, sv), clock.dt;
+                PeriodicCallback(affect, clock.dt;
                     final_affect = true, initial_affect = true)
             elseif clock isa SolverStepClock
-                affect = DiscreteSaveAffect(affect, sv)
                 DiscreteCallback(Returns(true), affect,
                     initialize = (c, u, t, integrator) -> affect(integrator))
             else
@@ -1110,8 +1108,6 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
         else
             cbs = CallbackSet(cbs, discrete_cbs...)
         end
-    else
-        svs = nothing
     end
     kwargs = filter_kwargs(kwargs)
     pt = something(get_metadata(sys), StandardODEProblem())
@@ -1120,17 +1116,8 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
     if cbs !== nothing
         kwargs1 = merge(kwargs1, (callback = cbs,))
     end
-    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.u0, prob.p, tspan[1])
-        end
-    end
-    prob
+    return ODEProblem{iip}(f, u0, tspan, p, pt; kwargs1..., kwargs...)
 end
 get_callback(prob::ODEProblem) = prob.kwargs[:callback]
 

test/clock.jl

@@ -111,20 +111,12 @@ Tf = 1.0
 prob = ODEProblem(ss, [x => 0.1], (0.0, Tf),
     [kp => 1.0; ud(k - 1) => 2.1; ud(k - 2) => 2.0])
 # create integrator so callback is evaluated at t=0 and we can test correct param values
-int = init(prob, Tsit5(); kwargshandle = KeywordArgSilent)
+int = init(prob, Tsit5())
 @test sort(vcat(int.p...)) == [0.1, 1.0, 2.1, 2.1, 2.1] # yd, kp, ud(k-1), ud, Hold(ud)
 prob = ODEProblem(ss, [x => 0.1], (0.0, Tf),
     [kp => 1.0; ud(k - 1) => 2.1; ud(k - 2) => 2.0]) # recreate problem to empty saved values
-sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
+sol = solve(prob, Tsit5())
 
-ss_nosplit = structural_simplify(sys; split = false)
-prob_nosplit = ODEProblem(ss_nosplit, [x => 0.1], (0.0, Tf),
-    [kp => 1.0; ud(k - 1) => 2.1; ud(k - 2) => 2.0])
-int = init(prob_nosplit, Tsit5(); kwargshandle = KeywordArgSilent)
-@test sort(int.p) == [0.1, 1.0, 2.1, 2.1, 2.1] # yd, kp, ud(k-1), ud, Hold(ud)
-prob_nosplit = ODEProblem(ss_nosplit, [x => 0.1], (0.0, Tf),
-    [kp => 1.0; ud(k - 1) => 2.1; ud(k - 2) => 2.0]) # recreate problem to empty saved values
-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.
 
@@ -155,11 +147,8 @@ cb = PeriodicCallback(
 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
-@test saved_values.t == sol.prob.kwargs[:disc_saved_values][1].t
-@test saved_values.t == sol_nosplit.prob.kwargs[:disc_saved_values][1].t
-@test saved_values.saveval == sol.prob.kwargs[:disc_saved_values][1].saveval
-@test saved_values.saveval == sol_nosplit.prob.kwargs[:disc_saved_values][1].saveval
+@test saved_values.t ≈ sol.discretes[1].t
+@test saved_values.saveval == sol.ps[[ud, ud(k - 1)]]
 
 @info "Testing multi-rate hybrid system"
 dt = 0.1
@@ -284,13 +273,10 @@ ci, varmap = infer_clocks(cl)
 @test varmap[u] == Continuous()
 
 ss = structural_simplify(cl)
-ss_nosplit = structural_simplify(cl; split = false)
 
 if VERSION >= v"1.7"
     prob = ODEProblem(ss, [x => 0.0], (0.0, 1.0), [kp => 1.0])
-    prob_nosplit = ODEProblem(ss_nosplit, [x => 0.0], (0.0, 1.0), [kp => 1.0])
-    sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
-    sol_nosplit = solve(prob_nosplit, Tsit5(), kwargshandle = KeywordArgSilent)
+    sol = solve(prob, Tsit5())
 
     function foo!(dx, x, p, t)
         kp, ud1, ud2 = p
@@ -321,7 +307,6 @@ if VERSION >= v"1.7"
     sol2 = solve(prob, Tsit5())
 
     @test sol.u≈sol2.u atol=1e-6
-    @test sol_nosplit.u≈sol2.u atol=1e-6
 end
 
 ##
@@ -440,13 +425,12 @@ y = res.y[:]
 prob = ODEProblem(ssys,
     [model.plant.x => 0.0; model.controller.kp => 2.0; model.controller.ki => 2.0],
     (0.0, Tf))
-int = init(prob, Tsit5(); kwargshandle = KeywordArgSilent)
+int = init(prob, Tsit5())
 @test_broken int.ps[Hold(ssys.holder.input.u)] == 2 # constant output * kp issue https://github.com/SciML/ModelingToolkit.jl/issues/2356
 @test int.ps[ssys.controller.x] == 1 # c2d
 @test int.ps[Sample(d)(ssys.sampler.input.u)] == 0 # disc state
 sol = solve(prob,
     Tsit5(),
-    kwargshandle = KeywordArgSilent,
     abstol = 1e-8,
     reltol = 1e-8)
 @test_skip begin
@@ -508,9 +492,9 @@ end
 
 @mtkbuild model = FirstOrderWithStepCounter()
 prob = ODEProblem(model, [], (0.0, 10.0))
-sol = solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
+sol = solve(prob, Tsit5())
 
-@test sol.prob.kwargs[:disc_saved_values][1].t == sol.t[1:2:end] # Test that the discrete-tiem system executed at every step of the continuous solver. The solver saves each time step twice, one state value before discrete affect and one after.
+@test sol.discretes[1].t == sol.t[1:2:end] # Test that the discrete-tiem system executed at every step of the continuous solver. The solver saves each time step twice, one state value before discrete affect and one after.
 @test_nowarn ModelingToolkit.build_explicit_observed_function(
     model, model.counter.ud)(sol.u[1], prob.p..., sol.t[1])
 
@@ -519,12 +503,12 @@ eqs = [D(y) ~ Hold(x)
        x ~ x(k - 1) + x(k - 2)]
 @mtkbuild sys = ODESystem(eqs, t)
 prob = ODEProblem(sys, [], (0.0, 10.0))
-int = init(prob, Tsit5(); kwargshandle = KeywordArgSilent)
+int = init(prob, Tsit5())
 @test int.ps[x] == 2.0
 @test int.ps[x(k - 1)] == 1.0
 
 @test_throws ErrorException ODEProblem(sys, [], (0.0, 10.0), [x => 2.0])
 prob = ODEProblem(sys, [], (0.0, 10.0), [x(k - 1) => 2.0])
-int = init(prob, Tsit5(); kwargshandle = KeywordArgSilent)
+int = init(prob, Tsit5())
 @test int.ps[x] == 3.0
 @test int.ps[x(k - 1)] == 2.0

test/parameter_dependencies.jl

@@ -86,18 +86,18 @@ end
     Tf = 1.0
     prob = ODEProblem(sys, [x => 0.0, y => 0.0], (0.0, Tf),
         [kp => 1.0; z(k - 1) => 3.0; yd(k - 1) => 0.0; z(k - 2) => 4.0; yd(k - 2) => 2.0])
-    @test_nowarn solve(prob, Tsit5(); kwargshandle = KeywordArgSilent)
+    @test_nowarn solve(prob, Tsit5())
 
     @mtkbuild sys = ODESystem(eqs, t; parameter_dependencies = [kq => 2kp],
         discrete_events = [[0.5] => [kp ~ 2.0]])
     prob = ODEProblem(sys, [x => 0.0, y => 0.0], (0.0, Tf),
         [kp => 1.0; z(k - 1) => 3.0; yd(k - 1) => 0.0; z(k - 2) => 4.0; yd(k - 2) => 2.0])
     @test prob.ps[kp] == 1.0
     @test prob.ps[kq] == 2.0
-    @test_nowarn solve(prob, Tsit5(), kwargshandle = KeywordArgSilent)
+    @test_nowarn solve(prob, Tsit5())
     prob = ODEProblem(sys, [x => 0.0, y => 0.0], (0.0, Tf),
         [kp => 1.0; z(k - 1) => 3.0; yd(k - 1) => 0.0; z(k - 2) => 4.0; yd(k - 2) => 2.0])
-    integ = init(prob, Tsit5(), kwargshandle = KeywordArgSilent)
+    integ = init(prob, Tsit5())
     @test integ.ps[kp] == 1.0
     @test integ.ps[kq] == 2.0
     step!(integ, 0.6)