add discrete events to more systems

Author: isaacsas

Project.toml

@@ -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

@@ -440,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)
@@ -453,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

@@ -527,10 +527,10 @@ symbolically calculating numerical enhancements.
 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; kwargs...)
+    cbs = process_events(sys; callback)
     sparsenoise === nothing && (sparsenoise = get(kwargs, :sparse, false))
 
     noiseeqs = get_noiseeqs(sys)

src/systems/jumps/jumpsystem.jl

@@ -1,19 +1,21 @@
 const JumpType = Union{VariableRateJump, ConstantRateJump, MassActionJump}
 
+# modifies the expression representating an affect function to
+# call reset_aggregated_jumps!(integrator).
 # assumes iip
-function add_jump_resetting!(expr, integrator)
+function _reset_aggregator!(expr, integrator)
     if expr isa Symbol
         error("Error, encountered a symbol. This should not happen.")
     end
 
     if (expr.head == :function)
-        add_jump_resetting!(expr.args[end], integrator)
+        _reset_aggregator!(expr.args[end], integrator)
     else
         if expr.args[end] == :nothing
             expr.args[end] = :(reset_aggregated_jumps!($integrator))
             push!(expr.args, :nothing)
         else
-            add_jump_resetting!(expr.args[end], integrator)
+            _reset_aggregator!(expr.args[end], integrator)
         end
     end
 
@@ -73,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
 
@@ -95,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"))
@@ -127,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)
@@ -325,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])
@@ -354,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

test/root_equations.jl

@@ -1,6 +1,8 @@
-using ModelingToolkit, OrdinaryDiffEq, StochasticDiffEq, Test
+using ModelingToolkit, OrdinaryDiffEq, StochasticDiffEq, JumpProcesses, Test
 using ModelingToolkit: SymbolicContinuousCallback, SymbolicContinuousCallbacks, NULL_AFFECT,
                        get_callback
+using StableRNGs
+rng = StableRNG(12345)
 
 @parameters t
 @variables x(t) = 0
@@ -461,7 +463,6 @@ let
     end
     cb2‵‵ = [2.0] => (affect!, [], [k], nothing)
     @named ssys4 = SDESystem(eqs, Equation[], t, [A], [k, t1], discrete_events = [cb1, cb2‵‵])
-    oprob4 = ODEProblem(ssys4, u0, tspan, p)
     testsol(ssys4, u0, p, tspan; tstops = [1.0])
 
     # mixing with symbolic condition in the func affect
@@ -480,3 +481,68 @@ let
     sol = testsol(ssys7, u0, p, (0.0, 10.0); tstops = [1.0, 2.0], skipparamtest = true)
     @test isapprox(sol(10.0)[1], .1; atol=1e-10, rtol=1e-10)
 end
+
+let rng = rng
+    function testsol(jsys, u0, p, tspan; tstops = Float64[], skipparamtest = false,
+                                         N = 40000, kwargs...)
+        dprob = DiscreteProblem(jsys, u0, tspan, p)
+        jprob = JumpProblem(jsys, dprob, Direct(); kwargs...)
+        sol = solve(jprob, SSAStepper(); tstops = tstops)
+        @test (sol(1.000000000001)[1] - sol(0.99999999999)[1]) == 1
+        !skipparamtest && (@test dprob.p[1] == 1.0)
+        @test sol(40.0)[1] == 0
+        sol
+    end
+
+    @parameters k t1 t2
+    @variables t A(t) B(t)
+
+    cond1 = (t == t1)
+    affect1 = [A ~ A + 1]
+    cb1 = cond1 => affect1
+    cond2 = (t == t2)
+    affect2 = [k ~ 1.0]
+    cb2 = cond2 => affect2
+
+    eqs = [MassActionJump(k, [A => 1], [A => -1])]
+    @named jsys = JumpSystem(eqs, t, [A], [k, t1, t2], discrete_events = [cb1, cb2])
+    u0 = [A => 1]
+    p = [k => 0.0, t1 => 1.0, t2 => 2.0]
+    tspan = (0.0, 40.0)
+    testsol(jsys, u0, p, tspan; tstops = [1.0, 2.0], rng)
+
+    cond1a = (t == t1)
+    affect1a = [A ~ A + 1, B ~ A]
+    cb1a = cond1a => affect1a
+    @named jsys1 = JumpSystem(eqs, t, [A, B], [k, t1, t2], discrete_events = [cb1a, cb2])
+    u0′ = [A => 1, B => 0]
+    sol = testsol(jsys1, u0′, p, tspan; tstops = [1.0, 2.0], check_length = false, rng)
+    @test sol(1.000000001, idxs = B) == 2
+
+    # same as above - but with set-time event syntax
+    cb1‵ = [1.0] => affect1 # needs to be a Vector for the event to happen only once
+    cb2‵ = [2.0] => affect2
+    @named jsys‵ = JumpSystem(eqs, t, [A], [k], discrete_events = [cb1‵, cb2‵])
+    testsol(jsys‵, u0, [p[1]], tspan; rng)
+
+    # mixing discrete affects
+    @named jsys3 = JumpSystem(eqs, t, [A], [k, t1, t2], discrete_events = [cb1, cb2‵])
+    testsol(jsys3, u0, p, tspan; tstops = [1.0], rng)
+
+    # mixing with a func affect
+    function affect!(integrator, u, p, ctx)
+        integrator.p[p.k] = 1.0
+        reset_aggregated_jumps!(integrator)
+        nothing
+    end
+    cb2‵‵ = [2.0] => (affect!, [], [k], nothing)
+    @named jsys4 = JumpSystem(eqs, t, [A], [k, t1], discrete_events = [cb1, cb2‵‵])
+    testsol(jsys4, u0, p, tspan; tstops = [1.0], rng)
+
+    # mixing with symbolic condition in the func affect
+    cb2‵‵‵ = (t == t2) => (affect!, [], [k], nothing)
+    @named jsys5 = JumpSystem(eqs, t, [A], [k, t1, t2], discrete_events = [cb1, cb2‵‵‵])
+    testsol(jsys5, u0, p, tspan; tstops = [1.0, 2.0], rng)
+    @named jsys6 = JumpSystem(eqs, t, [A], [k, t1, t2], discrete_events = [cb2‵‵‵, cb1])
+    testsol(jsys6, u0, p, tspan; tstops = [1.0, 2.0], rng)
+end