feat: allow saving discrete variables in symbolic callbacks

Author: AayushSabharwal

src/systems/callbacks.jl

@@ -394,14 +394,15 @@ function callback_save_header(sys::AbstractSystem, cb)
     save_idxs = get(ic.callback_to_clocks, cb, Int[])
     isempty(save_idxs) && return (identity, identity)
 
-    wrapper = function(expr)
-        return Func(expr.args, [], LiteralExpr(quote
-            $(expr.body)
-            save_idxs = $(save_idxs)
-            for idx in save_idxs
-                $(SciMLBase.save_discretes!)($(expr.args[1]), idx)
-            end
-        end))
+    wrapper = function (expr)
+        return Func(expr.args, [],
+            LiteralExpr(quote
+                $(expr.body)
+                save_idxs = $(save_idxs)
+                for idx in save_idxs
+                    $(SciMLBase.save_discretes!)($(expr.args[1]), idx)
+                end
+            end))
     end
 
     return wrapper, wrapper
@@ -705,7 +706,9 @@ function compile_user_affect(affect::FunctionalAffect, cb, sys, dvs, ps; kwargs.
     else
         save_idxs = Int[]
     end
-    let u = u, p = p, user_affect = func(affect), ctx = context(affect), save_idxs = save_idxs
+    let u = u, p = p, user_affect = func(affect), ctx = context(affect),
+        save_idxs = save_idxs
+
         function (integ)
             user_affect(integ, u, p, ctx)
             for idx in save_idxs

src/systems/jumps/jumpsystem.jl

@@ -221,7 +221,8 @@ function generate_affect_function(js::JumpSystem, affect, outputidxs)
         csubs = Dict(c => getdefault(c) for c in consts)
         affect = substitute(affect, csubs)
     end
-    compile_affect(affect, nothing, js, unknowns(js), parameters(js); outputidxs = outputidxs,
+    compile_affect(
+        affect, nothing, js, unknowns(js), parameters(js); outputidxs = outputidxs,
         expression = Val{true}, checkvars = false)
 end
 

test/symbolic_events.jl

@@ -876,12 +876,13 @@ end
     cb2 = [x ~ 0.5] => (save_affect!, [], [b], [b], nothing)
     cb3 = 1.0 => [c ~ t]
 
-    @mtkbuild sys = ODESystem(D(x) ~ cos(x), t, [x], [a, b, c]; continuous_events = [cb1, cb2], discrete_events = [cb3])
+    @mtkbuild sys = ODESystem(D(x) ~ cos(t), t, [x], [a, b, c];
+        continuous_events = [cb1, cb2], discrete_events = [cb3])
     prob = ODEProblem(sys, [x => 1.0], (0.0, 2pi), [a => 1.0, b => 2.0, c => 0.0])
     @test sort(canonicalize(Discrete(), prob.p)[1]) == [0.0, 1.0, 2.0]
     sol = solve(prob, Tsit5())
 
     @test sol[a] == [-1.0]
     @test sol[b] == [5.0, 5.0]
-    @test sol[c] == [1.0, 2.0, 3.0]
+    @test sol[c] == [1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
 end