Support constants in jump systems.

Author: lamorton

src/systems/jumps/jumpsystem.jl

@@ -157,13 +157,23 @@ function JumpSystem(eqs, iv, states, ps;
 end
 
 function generate_rate_function(js::JumpSystem, rate)
+    consts = collect_constants(rate)
+    if !isempty(consts) # The SymbolicUtils._build_function method of this case doesn't support preprocessing
+        csubs = Dict(c => getdefault(c) for c in consts)
+        rate = substitute(rate, csubs)
+    end
     rf = build_function(rate, states(js), parameters(js),
                         get_iv(js),
                         conv = states_to_sym(states(js)),
                         expression = Val{true})
 end
 
 function generate_affect_function(js::JumpSystem, affect, outputidxs)
+    consts = collect_constants(affect)
+    if !isempty(consts) # The SymbolicUtils._build_function method of this case doesn't support preprocessing
+        csubs = Dict(c => getdefault(c) for c in consts)
+        affect = substitute(affect, csubs)
+    end
     compile_affect(affect, js, states(js), parameters(js); outputidxs = outputidxs,
                    expression = Val{true}, checkvars = false)
 end

src/utils.jl

@@ -521,6 +521,13 @@ function collect_constants(eqs::AbstractArray{T}) where {T} # For generate_tgrad
     return constants
 end
 
+collect_constants(x::Num) = collect_constants(unwrap(x))
+function collect_constants(expr::Symbolic{T}) where {T} # For jump system affect / rate 
+    constants = Symbolic[]
+    collect_constants!(constants,expr)
+    return constants
+end
+
 function collect_constant!(constants, var)
     if isconstant(var)
         push!(constants, var)

test/discretesystem.jl

@@ -11,14 +11,15 @@ end;
 
 # Independent and dependent variables and parameters
 @parameters t c nsteps δt β γ
+@constants h = 1
 D = Difference(t; dt = 0.1)
 @variables S(t) I(t) R(t)
 
-infection = rate_to_proportion(β * c * I / (S + I + R), δt) * S
-recovery = rate_to_proportion(γ, δt) * I
+infection = rate_to_proportion(β * c * I / (S * h + I + R), δt * h) * S
+recovery = rate_to_proportion(γ * h, δt) * I
 
 # Equations
-eqs = [D(S) ~ S - infection,
+eqs = [D(S) ~ S - infection * h,
     D(I) ~ I + infection - recovery,
     D(R) ~ R + recovery]
 
@@ -99,7 +100,7 @@ D2 = Difference(t; dt = 2)
 # Equations
 eqs = [
     D1(x(t)) ~ 0.4x(t) + 0.3x(t - 1.5) + 0.1x(t - 3),
-    D2(y(t)) ~ 0.3y(t) + 0.7y(t - 2) + 0.1z,
+    D2(y(t)) ~ 0.3y(t) + 0.7y(t - 2) + 0.1z * h,
 ]
 
 # System
@@ -119,7 +120,7 @@ linearized_eqs = [eqs
 # observed variable handling
 @variables t x(t) RHS(t)
 @parameters τ
-@named fol = DiscreteSystem([D(x) ~ (1 - x) / τ]; observed = [RHS ~ (1 - x) / τ])
+@named fol = DiscreteSystem([D(x) ~ (1 - x) / τ]; observed = [RHS ~ (1 - x) / τ * h])
 @test isequal(RHS, @nonamespace fol.RHS)
 RHS2 = RHS
 @unpack RHS = fol

test/jumpsystem.jl

@@ -3,9 +3,10 @@ MT = ModelingToolkit
 
 # basic MT SIR model with tweaks
 @parameters β γ t
+@constants h=1
 @variables S(t) I(t) R(t)
-rate₁ = β * S * I
-affect₁ = [S ~ S - 1, I ~ I + 1]
+rate₁ = β * S * I * h
+affect₁ = [S ~ S - 1 * h, I ~ I + 1]
 rate₂ = γ * I + t
 affect₂ = [I ~ I - 1, R ~ R + 1]
 j₁ = ConstantRateJump(rate₁, affect₁)
@@ -52,8 +53,8 @@ jump2.affect!(integrator)
 @test all(integrator.u .== mtintegrator.u)
 
 # test MT can make and solve a jump problem
-rate₃ = γ * I
-affect₃ = [I ~ I - 1, R ~ R + 1]
+rate₃ = γ * I * h
+affect₃ = [I ~ I * h - 1, R ~ R + 1]
 j₃ = ConstantRateJump(rate₃, affect₃)
 @named js2 = JumpSystem([j₁, j₃], t, [S, I, R], [β, γ])
 u₀ = [999, 1, 0];
@@ -83,6 +84,7 @@ sol = solve(jprob, SSAStepper(), saveat = tspan[2] / 10)
 @test all(2 .* sol[S] .== sol[S2])
 
 # test save_positions is working
+
 jprob = JumpProblem(js2, dprob, Direct(), save_positions = (false, false))
 sol = solve(jprob, SSAStepper(), saveat = 1.0)
 @test all((sol.t) .== collect(0.0:tspan[2]))
@@ -192,7 +194,7 @@ sol = solve(jprob, SSAStepper(), tstops = [1000.0],
 
 # observed variable handling
 @variables OBS(t)
-@named js5 = JumpSystem([maj1, maj2], t, [S], [β, γ]; observed = [OBS ~ 2 * S])
+@named js5 = JumpSystem([maj1, maj2], t, [S], [β, γ]; observed = [OBS ~ 2 * S * h])
 OBS2 = OBS
 @test isequal(OBS2, @nonamespace js5.OBS)
 @unpack OBS = js5