add events for SDESystems

Author: isaacsas

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

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,16 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
     ctrl_jac = RefValue{Any}(EMPTY_JAC)
     Wfact = RefValue(EMPTY_JAC)
     Wfact_t = RefValue(EMPTY_JAC)
+    sysnames = nameof.(systems)
+    if length(unique(sysnames)) != length(sysnames)
+        throw(ArgumentError("System names must be unique."))
+    end
+    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...)
@@ -509,6 +530,7 @@ function DiffEqBase.SDEProblem{iip}(sys::SDESystem, u0map, tspan,
                                     kwargs...) where {iip}
     f, u0, p = process_DEProblem(SDEFunction{iip}, sys, u0map, parammap; check_length,
                                  kwargs...)
+    cbs = process_events(sys; kwargs...)
     sparsenoise === nothing && (sparsenoise = get(kwargs, :sparse, false))
 
     noiseeqs = get_noiseeqs(sys)
@@ -521,8 +543,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
 
 function DiffEqBase.SDEProblem(sys::SDESystem, args...; kwargs...)

src/systems/jumps/jumpsystem.jl

@@ -1,5 +1,25 @@
 const JumpType = Union{VariableRateJump, ConstantRateJump, MassActionJump}
 
+# assumes iip
+function add_jump_resetting!(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)
+    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)
+        end
+    end
+
+    nothing
+end
+
 """
 $(TYPEDEF)
 

test/root_equations.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, OrdinaryDiffEq, Test
+using ModelingToolkit, OrdinaryDiffEq, StochasticDiffEq, Test
 using ModelingToolkit: SymbolicContinuousCallback, SymbolicContinuousCallbacks, NULL_AFFECT,
                        get_callback
 
@@ -407,3 +407,76 @@ let
     sol = testsol(osys7, 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
+    function testsol(ssys, u0, p, tspan; tstops = Float64[], skipparamtest = false, kwargs...)
+        sprob = SDEProblem(ssys, u0, tspan, p; kwargs...)
+        sol = solve(sprob, RI5(); tstops = tstops, abstol = 1e-10, reltol = 1e-10)
+        @test isapprox(sol(1.0000000001)[1] - sol(0.999999999)[1], 1.0; rtol = 1e-4)
+        !skipparamtest && (@test sprob.p[1] == 1.0)
+        @test isapprox(sol(4.0)[1], 2 * exp(-2.0), atol=1e-4)
+        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
+
+    ∂ₜ = Differential(t)
+    eqs = [∂ₜ(A) ~ -k * A]
+    @named ssys = SDESystem(eqs, Equation[], t, [A], [k, t1, t2], discrete_events = [cb1, cb2])
+    u0 = [A => 1.0]
+    p = [k => 0.0, t1 => 1.0, t2 => 2.0]
+    tspan = (0.0, 4.0)
+    testsol(ssys, u0, p, tspan; tstops = [1.0, 2.0])
+
+    cond1a = (t == t1)
+    affect1a = [A ~ A + 1, B ~ A]
+    cb1a = cond1a => affect1a
+    @named ssys1 = SDESystem(eqs, Equation[], t, [A, B], [k, t1, t2], discrete_events = [cb1a, cb2])
+    u0′ = [A => 1.0, B => 0.0]
+    sol = testsol(ssys1, u0′, p, tspan; tstops = [1.0, 2.0], check_length = false)
+    @test sol(1.0000001, idxs = 2) == 2.0
+
+    # 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 ssys‵ = SDESystem(eqs, Equation[], t, [A], [k], discrete_events = [cb1‵, cb2‵])
+    testsol(ssys‵, u0, p, tspan)
+
+    # mixing discrete affects
+    @named ssys3 = SDESystem(eqs, Equation[], t, [A], [k, t1, t2], discrete_events = [cb1, cb2‵])
+    testsol(ssys3, u0, p, tspan; tstops = [1.0])
+
+    # mixing with a func affect
+    function affect!(integrator, u, p, ctx)
+        integrator.p[p.k] = 1.0
+        nothing
+    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
+    cb2‵‵‵ = (t == t2) => (affect!, [], [k], nothing)
+    @named ssys5 = SDESystem(eqs, Equation[], t, [A], [k, t1, t2], discrete_events = [cb1, cb2‵‵‵])
+    testsol(ssys5, u0, p, tspan; tstops = [1.0, 2.0])
+    @named ssys6 = SDESystem(eqs, Equation[], t, [A], [k, t1, t2], discrete_events = [cb2‵‵‵, cb1])
+    testsol(ssys6, u0, p, tspan; tstops = [1.0, 2.0])
+
+    # mix a continuous event too
+    cond3 = A ~ .1
+    affect3 = [k ~ 0.0]
+    cb3 = cond3 => affect3
+    @named ssys7 = SDESystem(eqs, Equation[], t, [A], [k, t1, t2], discrete_events = [cb1, cb2‵‵‵],
+                             continuous_events = [cb3])
+    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