Merge pull request #1714 from isaacsas/disc-cb-three

Add events to SDESystems and JumpSystems

Author: YingboMa

Project.toml

@@ -59,7 +59,7 @@ ForwardDiff = "0.10.3"
 Graphs = "1.5.2"
 IfElse = "0.1"
 JuliaFormatter = "1"
-JumpProcesses = "9"
+JumpProcesses = "9.1"
 LabelledArrays = "1.3"
 Latexify = "0.11, 0.12, 0.13, 0.14, 0.15"
 MacroTools = "0.5"
@@ -91,6 +91,7 @@ OptimizationOptimJL = "36348300-93cb-4f02-beb5-3c3902f8871e"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReferenceTests = "324d217c-45ce-50fc-942e-d289b448e8cf"
+StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SteadyStateDiffEq = "9672c7b4-1e72-59bd-8a11-6ac3964bc41f"
@@ -99,4 +100,4 @@ Sundials = "c3572dad-4567-51f8-b174-8c6c989267f4"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["AmplNLWriter", "BenchmarkTools", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]
+test = ["AmplNLWriter", "BenchmarkTools", "ForwardDiff", "Ipopt", "Ipopt_jll", "ModelingToolkitStandardLibrary", "Optimization", "OptimizationOptimJL", "OptimizationMOI", "OrdinaryDiffEq", "Random", "ReferenceTests", "SafeTestsets", "StableRNGs", "Statistics", "SteadyStateDiffEq", "Test", "StochasticDiffEq", "Sundials"]

src/systems/callbacks.jl

@@ -241,17 +241,14 @@ end
 ################################# compilation functions ####################################
 
 # handles ensuring that affect! functions work with integrator arguments
-function add_integrator_header(out = :u)
-    integrator = gensym(:MTKIntegrator)
-
+function add_integrator_header(integrator = gensym(:MTKIntegrator), out = :u)
     expr -> Func([DestructuredArgs(expr.args, integrator, inds = [:u, :p, :t])], [],
                  expr.body),
     expr -> Func([DestructuredArgs(expr.args, integrator, inds = [out, :u, :p, :t])], [],
                  expr.body)
 end
 
-function condition_header()
-    integrator = gensym(:MTKIntegrator)
+function condition_header(integrator = gensym(:MTKIntegrator))
     expr -> Func([expr.args[1], expr.args[2],
                      DestructuredArgs(expr.args[3:end], integrator, inds = [:p])], [], expr.body)
 end
@@ -296,7 +293,8 @@ Notes
 - `kwargs` are passed through to `Symbolics.build_function`.
 """
 function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothing,
-                        expression = Val{true}, checkvars = true, kwargs...)
+                        expression = Val{true}, checkvars = true,
+                        postprocess_affect_expr! = nothing, kwargs...)
     if isempty(eqs)
         if expression == Val{true}
             return :((args...) -> ())
@@ -337,10 +335,17 @@ function compile_affect(eqs::Vector{Equation}, sys, dvs, ps; outputidxs = nothin
             p = ps
         end
         t = get_iv(sys)
-        rf_oop, rf_ip = build_function(rhss, u, p, t; expression = expression,
-                                       wrap_code = add_integrator_header(outvar),
+        integ = gensym(:MTKIntegrator)
+        getexpr = (postprocess_affect_expr! === nothing) ? expression : Val{true}
+        rf_oop, rf_ip = build_function(rhss, u, p, t; expression = getexpr,
+                                       wrap_code = add_integrator_header(integ, outvar),
                                        outputidxs = update_inds,
                                        kwargs...)
+        # applied user-provided function to the generated expression
+        if postprocess_affect_expr! !== nothing
+            postprocess_affect_expr!(rf_ip, integ)
+            (expression == Val{false}) && (return @RuntimeGeneratedFunction(rf_ip))
+        end
         rf_ip
     end
 end
@@ -435,10 +440,11 @@ function compile_affect(affect::FunctionalAffect, sys, dvs, ps; kwargs...)
     compile_user_affect(affect, sys, dvs, ps; kwargs...)
 end
 
-function generate_timed_callback(cb, sys, dvs, ps; kwargs...)
+function generate_timed_callback(cb, sys, dvs, ps; postprocess_affect_expr! = nothing,
+                                 kwargs...)
     cond = condition(cb)
     as = compile_affect(affects(cb), sys, dvs, ps; expression = Val{false},
-                        kwargs...)
+                        postprocess_affect_expr!, kwargs...)
     if cond isa AbstractVector
         # Preset Time
         return PresetTimeCallback(cond, as)
@@ -448,13 +454,15 @@ function generate_timed_callback(cb, sys, dvs, ps; kwargs...)
     end
 end
 
-function generate_discrete_callback(cb, sys, dvs, ps; kwargs...)
+function generate_discrete_callback(cb, sys, dvs, ps; postprocess_affect_expr! = nothing,
+                                    kwargs...)
     if is_timed_condition(cb)
-        return generate_timed_callback(cb, sys, dvs, ps, kwargs...)
+        return generate_timed_callback(cb, sys, dvs, ps; postprocess_affect_expr!,
+                                       kwargs...)
     else
         c = compile_condition(cb, sys, dvs, ps; expression = Val{false}, kwargs...)
         as = compile_affect(affects(cb), sys, dvs, ps; expression = Val{false},
-                            kwargs...)
+                            postprocess_affect_expr!, kwargs...)
         return DiscreteCallback(c, as)
     end
 end

src/systems/diffeqs/sdesystem.jl

@@ -85,18 +85,30 @@ struct SDESystem <: AbstractODESystem
     type: type of the system
     """
     connector_type::Any
+    """
+    continuous_events: A `Vector{SymbolicContinuousCallback}` that model events.
+    The integrator will use root finding to guarantee that it steps at each zero crossing.
+    """
+    continuous_events::Vector{SymbolicContinuousCallback}
+    """
+    discrete_events: A `Vector{SymbolicDiscreteCallback}` that models events. Symbolic
+    analog to `SciMLBase.DiscreteCallback` that exectues an affect when a given condition is
+    true at the end of an integration step.
+    """
+    discrete_events::Vector{SymbolicDiscreteCallback}
 
     function SDESystem(deqs, neqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac,
-                       ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type;
-                       checks::Bool = true)
+                       ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type,
+                       cevents, devents; checks::Bool = true)
         if checks
             check_variables(dvs, iv)
             check_parameters(ps, iv)
             check_equations(deqs, iv)
+            check_equations(equations(cevents), iv)
             all_dimensionless([dvs; ps; iv]) || check_units(deqs, neqs)
         end
         new(deqs, neqs, iv, dvs, ps, var_to_name, ctrls, observed, tgrad, jac, ctrl_jac,
-            Wfact, Wfact_t, name, systems, defaults, connector_type)
+            Wfact, Wfact_t, name, systems, defaults, connector_type, cevents, devents)
     end
 end
 
@@ -109,7 +121,9 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
                    defaults = _merge(Dict(default_u0), Dict(default_p)),
                    name = nothing,
                    connector_type = nothing,
-                   checks = true)
+                   checks = true,
+                   continuous_events = nothing,
+                   discrete_events = nothing)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
     deqs = scalarize(deqs)
@@ -139,9 +153,12 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
     ctrl_jac = RefValue{Any}(EMPTY_JAC)
     Wfact = RefValue(EMPTY_JAC)
     Wfact_t = RefValue(EMPTY_JAC)
+    cont_callbacks = SymbolicContinuousCallbacks(continuous_events)
+    disc_callbacks = SymbolicDiscreteCallbacks(discrete_events)
+
     SDESystem(deqs, neqs, iv′, dvs′, ps′, var_to_name, ctrl′, observed, tgrad, jac,
               ctrl_jac, Wfact, Wfact_t, name, systems, defaults, connector_type,
-              checks = checks)
+              cont_callbacks, disc_callbacks; checks = checks)
 end
 
 function SDESystem(sys::ODESystem, neqs; kwargs...)
@@ -491,9 +508,10 @@ end
 function DiffEqBase.SDEProblem{iip}(sys::SDESystem, u0map, tspan,
                                     parammap = DiffEqBase.NullParameters();
                                     sparsenoise = nothing, check_length = true,
-                                    kwargs...) where {iip}
+                                    callback = nothing, kwargs...) where {iip}
     f, u0, p = process_DEProblem(SDEFunction{iip}, sys, u0map, parammap; check_length,
                                  kwargs...)
+    cbs = process_events(sys; callback)
     sparsenoise === nothing && (sparsenoise = get(kwargs, :sparse, false))
 
     noiseeqs = get_noiseeqs(sys)
@@ -506,8 +524,8 @@ function DiffEqBase.SDEProblem{iip}(sys::SDESystem, u0map, tspan,
         noise_rate_prototype = zeros(eltype(u0), size(noiseeqs))
     end
 
-    SDEProblem{iip}(f, f.g, u0, tspan, p; noise_rate_prototype = noise_rate_prototype,
-                    kwargs...)
+    SDEProblem{iip}(f, f.g, u0, tspan, p; callback = cbs,
+                    noise_rate_prototype = noise_rate_prototype, kwargs...)
 end
 
 """

src/systems/jumps/jumpsystem.jl

@@ -1,5 +1,27 @@
 const JumpType = Union{VariableRateJump, ConstantRateJump, MassActionJump}
 
+# modifies the expression representating an affect function to
+# call reset_aggregated_jumps!(integrator).
+# assumes iip
+function _reset_aggregator!(expr, integrator)
+    if expr isa Symbol
+        error("Error, encountered a symbol. This should not happen.")
+    end
+
+    if (expr.head == :function)
+        _reset_aggregator!(expr.args[end], integrator)
+    else
+        if expr.args[end] == :nothing
+            expr.args[end] = :(reset_aggregated_jumps!($integrator))
+            push!(expr.args, :nothing)
+        else
+            _reset_aggregator!(expr.args[end], integrator)
+        end
+    end
+
+    nothing
+end
+
 """
 $(TYPEDEF)
 
@@ -53,16 +75,25 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractTimeDependentSystem
     type: type of the system
     """
     connector_type::Any
+    """
+    discrete_events: A `Vector{SymbolicDiscreteCallback}` that models events. Symbolic
+    analog to `SciMLBase.DiscreteCallback` that exectues an affect when a given condition is
+    true at the end of an integration step. *Note, one must make sure to call
+    `reset_aggregated_jumps!(integrator)` if using a custom affect function that changes any
+    state value or parameter.*
+    """
+    discrete_events::Vector{SymbolicDiscreteCallback}
+
     function JumpSystem{U}(ap::U, iv, states, ps, var_to_name, observed, name, systems,
-                           defaults, connector_type;
+                           defaults, connector_type, devents;
                            checks::Bool = true) where {U <: ArrayPartition}
         if checks
             check_variables(states, iv)
             check_parameters(ps, iv)
             all_dimensionless([states; ps; iv]) || check_units(ap, iv)
         end
         new{U}(ap, iv, states, ps, var_to_name, observed, name, systems, defaults,
-               connector_type)
+               connector_type, devents)
     end
 end
 
@@ -75,6 +106,8 @@ function JumpSystem(eqs, iv, states, ps;
                     name = nothing,
                     connector_type = nothing,
                     checks = true,
+                    continuous_events = nothing,
+                    discrete_events = nothing,
                     kwargs...)
     name === nothing &&
         throw(ArgumentError("The `name` keyword must be provided. Please consider using the `@named` macro"))
@@ -107,9 +140,12 @@ function JumpSystem(eqs, iv, states, ps;
     process_variables!(var_to_name, defaults, states)
     process_variables!(var_to_name, defaults, ps)
     isempty(observed) || collect_var_to_name!(var_to_name, (eq.lhs for eq in observed))
+    (continuous_events === nothing) ||
+        error("JumpSystems currently only support discrete events.")
+    disc_callbacks = SymbolicDiscreteCallbacks(discrete_events)
 
     JumpSystem{typeof(ap)}(ap, value(iv), states, ps, var_to_name, observed, name, systems,
-                           defaults, connector_type, checks = checks)
+                           defaults, connector_type, disc_callbacks; checks = checks)
 end
 
 function generate_rate_function(js::JumpSystem, rate)
@@ -305,7 +341,8 @@ jprob = JumpProblem(js, dprob, Direct())
 sol = solve(jprob, SSAStepper())
 ```
 """
-function JumpProcesses.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
+function JumpProcesses.JumpProblem(js::JumpSystem, prob, aggregator; callback = nothing,
+                                   kwargs...)
     statetoid = Dict(value(state) => i for (i, state) in enumerate(states(js)))
     eqs = equations(js)
     invttype = prob.tspan[1] === nothing ? Float64 : typeof(1 / prob.tspan[2])
@@ -334,9 +371,12 @@ function JumpProcesses.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
         jtoj = nothing
     end
 
+    # handle events, making sure to reset aggregators in the generated affect functions
+    cbs = process_events(js; callback, postprocess_affect_expr! = _reset_aggregator!)
+
     JumpProblem(prob, aggregator, jset; dep_graph = jtoj, vartojumps_map = vtoj,
-                jumptovars_map = jtov,
-                scale_rates = false, nocopy = true, kwargs...)
+                jumptovars_map = jtov, scale_rates = false, nocopy = true,
+                callback = cbs, kwargs...)
 end
 
 ### Functions to determine which states a jump depends on