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

AayushSabharwal · AayushSabharwal · commit ef728cc588e2 · 2024-07-29T11:15:17.000+05:30
diff --git a/src/systems/clock_inference.jl b/src/systems/clock_inference.jl
@@ -203,19 +203,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
@@ -231,11 +226,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
@@ -255,7 +246,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,
@@ -281,56 +272,16 @@ function generate_discrete_affect(
         disc_range = [parameter_index(osys, sym) for sym in unknowns(sys)]
         save_expr = :($(SciMLBase.save_discretes!)(integrator, $i))
         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`
@@ -353,71 +304,32 @@ function generate_discrete_affect(
                         $(_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)
 
         push!(affect_funs, affect!)
-        push!(init_funs, disc_init)
-        push!(svs, sv)
     end
     if eval_expression
         affects = map(a -> eval_module.eval(toexpr(LiteralExpr(a))), affect_funs)
-        inits = map(a -> eval_module.eval(toexpr(LiteralExpr(a))), init_funs)
     else
         affects = map(affect_funs) do a
             drop_expr(RuntimeGeneratedFunction(
                 eval_module, eval_module, toexpr(LiteralExpr(a))))
         end
-        inits = map(init_funs) do a
-            drop_expr(RuntimeGeneratedFunction(
-                eval_module, eval_module, toexpr(LiteralExpr(a))))
-        end
     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
diff --git a/src/systems/diffeqs/abstractodesystem.jl b/src/systems/diffeqs/abstractodesystem.jl
@@ -1008,14 +1008,13 @@ function DiffEqBase.ODEProblem{iip, specialize}(sys::AbstractODESystem, u0map =
     cbs = process_events(sys; callback, eval_expression, eval_module, kwargs...)
     inits = []
     if has_discrete_subsystems(sys) && (dss = get_discrete_subsystems(sys)) !== nothing
-        affects, inits, clocks, svs = ModelingToolkit.generate_discrete_affect(
+        affects, clocks = ModelingToolkit.generate_discrete_affect(
             sys, dss...; eval_expression, eval_module)
-        discrete_cbs = map(affects, clocks, svs) do affect, clock, sv
+        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
@@ -1031,8 +1030,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())
@@ -1041,17 +1038,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]
 
diff --git a/test/parameter_dependencies.jl b/test/parameter_dependencies.jl
@@ -173,22 +173,19 @@ end
     @test_skip begin
         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)
+            [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())
 
         @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])
+            [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)
+            [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())
         @test integ.ps[kp] == 1.0
         @test integ.ps[kq] == 2.0
         step!(integ, 0.6)
