Simplify callback interface, fix references

Author: BenChung

docs/src/basics/Events.md

@@ -416,12 +416,12 @@ in between. To do this, we create two continuous callbacks:
 using Setfield
 furnace_disable = ModelingToolkit.SymbolicContinuousCallback(
     [temp ~ furnace_off_threshold],
-    ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x, o, i, c
+    ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x, o, c, i
         @set! x.furnace_on = false
     end)
 furnace_enable = ModelingToolkit.SymbolicContinuousCallback(
     [temp ~ furnace_on_threshold],
-    ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x, o, i, c
+    ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x, o, c, i
         @set! x.furnace_on = true
     end)
 ```
@@ -454,18 +454,16 @@ In our example, each event merely changes whether the furnace is on or off. Acco
 evaluate this before calling our function to fill out all of the numerical values, then apply them back to the system
 once our affect function returns. Furthermore, it will check that it is possible to do this assignment.
 
-The function given to `ImperativeAffect` needs to have one of four signatures, checked in this order:
-
-  - `f(modified::NamedTuple, observed::NamedTuple, ctx, integrator)::NamedTuple` if the function needs the low-level integrator,
-  - `f(modified::NamedTuple, observed::NamedTuple, ctx)::NamedTuple` if the function needs the user-defined context,
-  - `f(modified::NamedTuple, observed::NamedTuple)::NamedTuple` if the function also reads observed values from the system,
-  - `f(modified::NamedTuple)::NamedTuple` if the function only writes values (unknowns or parameters) to the system.
-    The `do` block in the example implicitly constructs said function inline. For exposition, we use the full version (e.g. `x, o, i, c`) but this could be simplified to merely `x`.
+The function given to `ImperativeAffect` needs to have the signature:
 
+```julia 
+    f(modified::NamedTuple, observed::NamedTuple, ctx, integrator)::NamedTuple
+```
 The function `f` will be called with `observed` and `modified` `NamedTuple`s that are derived from their respective `NamedTuple` definitions.
 In our example, if `furnace_on` is `false`, then the value of the `x` that's passed in as `modified` will be `(furnace_on = false)`.
 The modified values should be passed out in the same format: to set `furnace_on` to `true` we need to return a tuple `(furnace_on = true)`.
-We use Setfield to do this in the example, recreating the result tuple before returning it.
+The examples does this with Setfield, recreating the result tuple before returning it; the returned tuple may optionally be missing values as 
+well, in which case those values will not be written back to the problem.
 
 Accordingly, we can now interpret the `ImperativeAffect` definitions to mean that when `temp = furnace_off_threshold` we
 will write `furnace_on = false` back to the system, and when `temp = furnace_on_threshold` we will write `furnace_on = true` back
@@ -543,7 +541,7 @@ In our encoder, we interpret this as occlusion or nonocclusion of the sensor, up
 
 ```@example events
 qAevt = ModelingToolkit.SymbolicContinuousCallback([cos(100 * theta) ~ 0],
-    ModelingToolkit.ImperativeAffect((; qA, hA, hB, cnt), (; qB)) do x, o, i, c
+    ModelingToolkit.ImperativeAffect((; qA, hA, hB, cnt), (; qB)) do x, o, c, i
         @set! x.hA = x.qA
         @set! x.hB = o.qB
         @set! x.qA = 1
@@ -567,7 +565,7 @@ Instead, we can use right root finding:
 
 ```@example events
 qBevt = ModelingToolkit.SymbolicContinuousCallback([cos(100 * theta - π / 2) ~ 0],
-    ModelingToolkit.ImperativeAffect((; qB, hA, hB, cnt), (; qA, theta)) do x, o, i, c
+    ModelingToolkit.ImperativeAffect((; qB, hA, hB, cnt), (; qA, theta)) do x, o, c, i
         @set! x.hA = o.qA
         @set! x.hB = x.qB
         @set! x.qB = clamp(sign(cos(100 * o.theta - π / 2)), 0.0, 1.0)

src/systems/callbacks.jl

@@ -76,12 +76,11 @@ end
 
 `ImperativeAffect` is a helper for writing affect functions that will compute observed values and
 ensure that modified values are correctly written back into the system. The affect function `f` needs to have
-one of four signatures:
-* `f(modified::NamedTuple)::NamedTuple` if the function only writes values (unknowns or parameters) to the system,
-* `f(modified::NamedTuple, observed::NamedTuple)::NamedTuple` if the function also reads observed values from the system,
-* `f(modified::NamedTuple, observed::NamedTuple, ctx)::NamedTuple` if the function needs the user-defined context,
-* `f(modified::NamedTuple, observed::NamedTuple, ctx, integrator)::NamedTuple` if the function needs the low-level integrator.
-These will be checked in reverse order (that is, the four-argument version first, than the 3, etc).
+the signature 
+
+```
+    f(modified::NamedTuple, observed::NamedTuple, ctx, integrator)::NamedTuple
+```
 
 The function `f` will be called with `observed` and `modified` `NamedTuple`s that are derived from their respective `NamedTuple` definitions.
 Each  declaration`NamedTuple` should map an expression to a symbol; for example if we pass `observed=(; x = a + b)` this will alias the result of executing `a+b` in the system as `x`
@@ -1046,7 +1045,7 @@ function compile_user_affect(affect::ImperativeAffect, cb, sys, dvs, ps; kwargs.
     Implementation sketch:
         generate observed function (oop), should save to a component array under obs_syms
         do the same stuff as the normal FA for pars_syms
-        call the affect method - test if it's OOP or IP using applicable
+        call the affect method
         unpack and apply the resulting values
     =#
     function check_dups(syms, exprs) # = (syms_dedup, exprs_dedup)
@@ -1135,22 +1134,10 @@ function compile_user_affect(affect::ImperativeAffect, cb, sys, dvs, ps; kwargs.
                 integ.u, integ.p, integ.t))
 
             # let the user do their thing
-            modvals = if applicable(
-                user_affect, upd_component_array, obs_component_array, ctx, integ)
-                user_affect(upd_component_array, obs_component_array, ctx, integ)
-            elseif applicable(user_affect, upd_component_array, obs_component_array, ctx)
-                user_affect(upd_component_array, obs_component_array, ctx)
-            elseif applicable(user_affect, upd_component_array, obs_component_array)
-                user_affect(upd_component_array, obs_component_array)
-            elseif applicable(user_affect, upd_component_array)
-                user_affect(upd_component_array)
-            else
-                @error "User affect function $user_affect needs to implement one of the supported ImperativeAffect callback forms; see the ImperativeAffect docstring for more details"
-                user_affect(upd_component_array, obs_component_array, integ, ctx) # this WILL error but it'll give a more sensible message
-            end
-
+            upd_vals = user_affect(upd_component_array, obs_component_array, ctx, integ)
+            
             # write the new values back to the integrator
-            _generated_writeback(integ, upd_funs, modvals)
+            _generated_writeback(integ, upd_funs, upd_vals)
 
             for idx in save_idxs
                 SciMLBase.save_discretes!(integ, idx)

test/symbolic_events.jl

@@ -1027,60 +1027,9 @@ end
 
     furnace_off = ModelingToolkit.SymbolicContinuousCallback(
         [temp ~ furnace_off_threshold],
-        ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x, o, i
-            @set! x.furnace_on = false
-        end)
-    furnace_enable = ModelingToolkit.SymbolicContinuousCallback(
-        [temp ~ furnace_on_threshold],
-        ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x, o, i
-            @set! x.furnace_on = true
-        end)
-    @named sys = ODESystem(
-        eqs, t, [temp], params; continuous_events = [furnace_off, furnace_enable])
-    ss = structural_simplify(sys)
-    prob = ODEProblem(ss, [temp => 0.0, furnace_on => true], (0.0, 100.0))
-    sol = solve(prob, Tsit5(); dtmax = 0.01)
-    @test all(sol[temp][sol.t .> 1.0] .<= 0.79) && all(sol[temp][sol.t .> 1.0] .>= 0.49)
-
-    furnace_off = ModelingToolkit.SymbolicContinuousCallback(
-        [temp ~ furnace_off_threshold],
-        ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x, o
-            @set! x.furnace_on = false
-        end)
-    furnace_enable = ModelingToolkit.SymbolicContinuousCallback(
-        [temp ~ furnace_on_threshold],
-        ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x, o
-            @set! x.furnace_on = true
-        end)
-    @named sys = ODESystem(
-        eqs, t, [temp], params; continuous_events = [furnace_off, furnace_enable])
-    ss = structural_simplify(sys)
-    prob = ODEProblem(ss, [temp => 0.0, furnace_on => true], (0.0, 100.0))
-    sol = solve(prob, Tsit5(); dtmax = 0.01)
-    @test all(sol[temp][sol.t .> 1.0] .<= 0.79) && all(sol[temp][sol.t .> 1.0] .>= 0.49)
-
-    furnace_off = ModelingToolkit.SymbolicContinuousCallback(
-        [temp ~ furnace_off_threshold],
-        ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x
-            @set! x.furnace_on = false
-        end)
-    furnace_enable = ModelingToolkit.SymbolicContinuousCallback(
-        [temp ~ furnace_on_threshold],
-        ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x
-            @set! x.furnace_on = true
-        end)
-    @named sys = ODESystem(
-        eqs, t, [temp], params; continuous_events = [furnace_off, furnace_enable])
-    ss = structural_simplify(sys)
-    prob = ODEProblem(ss, [temp => 0.0, furnace_on => true], (0.0, 100.0))
-    sol = solve(prob, Tsit5(); dtmax = 0.01)
-    @test all(sol[temp][sol.t .> 1.0] .<= 0.79) && all(sol[temp][sol.t .> 1.0] .>= 0.49)
-
-    furnace_off = ModelingToolkit.SymbolicContinuousCallback(
-        [temp ~ furnace_off_threshold],
-        ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x
+        ModelingToolkit.ImperativeAffect(modified = (; furnace_on)) do x, o, c, i
             @set! x.furnace_on = false
-        end; initialize = ModelingToolkit.ImperativeAffect(modified = (; temp)) do x
+        end; initialize = ModelingToolkit.ImperativeAffect(modified = (; temp)) do x, o, c, i
             @set! x.temp = 0.2
         end)
     furnace_enable = ModelingToolkit.SymbolicContinuousCallback(
@@ -1180,7 +1129,7 @@ end
         end
     end
     qAevt = ModelingToolkit.SymbolicContinuousCallback([cos(100 * theta) ~ 0],
-        ModelingToolkit.ImperativeAffect((; qA, hA, hB, cnt), (; qB)) do x, o, i, c
+        ModelingToolkit.ImperativeAffect((; qA, hA, hB, cnt), (; qB)) do x, o, c, i
             @set! x.hA = x.qA
             @set! x.hB = o.qB
             @set! x.qA = 1
@@ -1196,7 +1145,7 @@ end
             x
         end; rootfind = SciMLBase.RightRootFind)
     qBevt = ModelingToolkit.SymbolicContinuousCallback([cos(100 * theta - π / 2) ~ 0],
-        ModelingToolkit.ImperativeAffect((; qB, hA, hB, cnt), (; qA)) do x, o, i, c
+        ModelingToolkit.ImperativeAffect((; qB, hA, hB, cnt), (; qA)) do x, o, c, i
             @set! x.hA = o.qA
             @set! x.hB = x.qB
             @set! x.qB = 1