Merge pull request #1673 from dodoplus/func-affect

Func affect

Author: ChrisRackauckas

docs/src/basics/Composition.md

@@ -337,3 +337,84 @@ plot(sol(tv)[y], sol(tv)[x], line_z=tv)
 vline!([-1.5, 1.5], l=(:black, 5), primary=false)
 hline!([0], l=(:black, 5), primary=false)
 ```
+
+### Generalized affect support
+In some instances, a more flexible response to events is needed, which cannot be encapsulated by
+an equation. For example, a component may implement complex behavior that it is inconvenient or
+impossible to capture in equations.
+
+ModelingToolkit therefore supports Julia functions as affects: instead of an equation, an affect is
+defined as a `tuple`:
+```
+[x ~ 0] => (affect!, [v, x], [p, q], ctx)
+```
+
+where, `affect!` is a Julia function with the signature: `affect!(integ, u, p, ctx)`; `[u,v]` and `[p,q]` 
+are the states (variables) and parameters that are accessed by `affect!`, respectively; and `ctx` is a 
+context that is passed to `affect!` as the `ctx` argument.
+
+`affect!` receives the `DiffEqs` integrator as its first argument, which can then be used to access states
+and parameters that are provided in the `u` and `p` arguments (implemented as `NamedTuple`s):
+
+```
+function affect!(integ, u, v, ctx)
+    # integ.t is the current time
+    # integ.u[u.v] is the value of the state `v` above
+    # integ.p[p.q] is the value of the parameter `q` above
+end
+```
+
+When accessing variables of a sub-system, it could be useful to rename them (alternatively, an affect function
+may be reused in different contexts):
+```
+[x ~ 0] => (affect!, [resistor₊v => :v, x], [p, q => :p2], ctx)
+```
+
+Here, `resistor₊v` is passed as `v` while `q` has been renamed `p2`.
+
+As an example, here is the bouncing ball example from `DiffEqs` using ModelingToolkit:
+
+```@example events
+sts = @variables y(t), v(t)
+par = @parameters g = 9.8
+bb_eqs = [D(y) ~ v
+          D(v) ~ -g]
+
+function bb_affect!(integ, u, p, ctx)
+    integ.u[u.v] = -integ.u[u.v]
+end
+
+@named bb_model = ODESystem(bb_eqs, t, sts, par,
+                            continuous_events = [[y ~ 0] => (bb_affect!, [v], [], nothing)])
+
+bb_sys = structural_simplify(bb_model)
+u0 = [v => 0.0, y => 50.0]
+
+bb_prob = ODEProblem(bb_sys, u0, (0, 15.0))
+bb_sol = solve(bb_prob, Tsit5())
+
+plot(bb_sol)
+```
+
+### Discrete events support
+In addition to continuous events, discrete events are also supported.
+
+TBD
+
+Two important sub-classes of discrete events are periodic and set-time events. A periodic event is triggered at
+fixed intervals (e.g. every Δt seconds). To specify a periodic interval, pass the interval as the condition for
+the event:
+
+```
+discrete_events=[1.0 => [v ~ -v]]
+```
+
+will change the sign of `v` at t=1.0, 2.0, ...
+
+Alternatively, the event may be triggered at specific set times:
+```
+discrete_events=[[1.0, 4.0] => [v ~ -v]]
+```
+
+will change the sign of `v` *only* at t=1.0, 4.0.
+

src/structural_transformation/codegen.jl

@@ -1,7 +1,6 @@
 using LinearAlgebra
 
-using ModelingToolkit: isdifferenceeq, has_continuous_events, generate_rootfinding_callback,
-                       generate_difference_cb, merge_cb
+using ModelingToolkit: isdifferenceeq, process_events
 
 const MAX_INLINE_NLSOLVE_SIZE = 8
 
@@ -522,18 +521,11 @@ function ODAEProblem{iip}(sys,
                                        use_union)
 
     has_difference = any(isdifferenceeq, eqs)
-    if has_continuous_events(sys)
-        event_cb = generate_rootfinding_callback(sys; kwargs...)
-    else
-        event_cb = nothing
-    end
-    difference_cb = has_difference ? generate_difference_cb(sys; kwargs...) : nothing
-    cb = merge_cb(event_cb, difference_cb)
-    cb = merge_cb(cb, callback)
+    cbs = process_events(sys; callback, has_difference, kwargs...)
 
-    if cb === nothing
+    if cbs === nothing
         ODEProblem{iip}(fun, u0, tspan, p; kwargs...)
     else
-        ODEProblem{iip}(fun, u0, tspan, p; callback = cb, kwargs...)
+        ODEProblem{iip}(fun, u0, tspan, p; callback = cbs, kwargs...)
     end
 end

src/systems/abstractsystem.jl

@@ -1348,6 +1348,7 @@ function Base.hash(sys::AbstractSystem, s::UInt)
     end
     s = foldr(hash, get_observed(sys), init = s)
     s = foldr(hash, get_continuous_events(sys), init = s)
+    s = foldr(hash, get_discrete_events(sys), init = s)
     s = hash(independent_variables(sys), s)
     return s
 end
@@ -1375,16 +1376,17 @@ function extend(sys::AbstractSystem, basesys::AbstractSystem; name::Symbol = nam
     sts = union(get_states(basesys), get_states(sys))
     ps = union(get_ps(basesys), get_ps(sys))
     obs = union(get_observed(basesys), get_observed(sys))
-    evs = union(get_continuous_events(basesys), get_continuous_events(sys))
+    cevs = union(get_continuous_events(basesys), get_continuous_events(sys))
+    devs = union(get_discrete_events(basesys), get_discrete_events(sys))
     defs = merge(get_defaults(basesys), get_defaults(sys)) # prefer `sys`
     syss = union(get_systems(basesys), get_systems(sys))
 
     if length(ivs) == 0
         T(eqs, sts, ps, observed = obs, defaults = defs, name = name, systems = syss,
-          continuous_events = evs)
+          continuous_events = cevs, discrete_events = devs)
     elseif length(ivs) == 1
         T(eqs, ivs[1], sts, ps, observed = obs, defaults = defs, name = name,
-          systems = syss, continuous_events = evs)
+          systems = syss, continuous_events = cevs, discrete_events = devs)
     end
 end