can make JumpProblem, but wrong answers on SIR

Author: isaacsas

src/ModelingToolkit.jl

@@ -5,7 +5,7 @@ using StaticArrays, LinearAlgebra, SparseArrays
 using Latexify, Unitful, ArrayInterface
 using MacroTools
 using UnPack: @unpack
-using DiffEqJump: VariableRateJump, ConstantRateJump, MassActionJump
+using DiffEqJump
 
 using Base.Threads
 import MacroTools: splitdef, combinedef, postwalk, striplines
@@ -104,7 +104,7 @@ include("build_function.jl")
 export ODESystem, ODEFunction
 export SDESystem, SDEFunction
 export JumpSystem
-export ODEProblem, SDEProblem, NonlinearProblem, OptimizationProblem
+export ODEProblem, SDEProblem, NonlinearProblem, OptimizationProblem, JumpProblem
 export NonlinearSystem, OptimizationSystem
 export ode_order_lowering
 export PDESystem

src/systems/jumps/jumpsystem.jl

@@ -25,7 +25,7 @@ generate_affect_function(js, affect) = build_function(affect, states(js),
                                                       independent_variable(js),
                                                       expression=Val{false},
                                                       headerfun=add_integrator_header)[2]
-function assemble_vrj(js, vrj)
+function assemble_vrj(js, vrj) 
     rate   = generate_rate_function(js, vrj.rate)
     affect = generate_affect_function(js, vrj.affect!)
     VariableRateJump(rate, affect)
@@ -37,23 +37,26 @@ function assemble_crj(js, crj)
     ConstantRateJump(rate, affect)
 end
 
-function assemble_maj(maj, states_to_idxs, ps_to_idxs; scale_rate=false)
-
-    # mass action scaled_rates need to be a Number, but 
-    # operations are numbers, so can't check the type directly
-    @assert !isa(maj.scaled_rates, Union{Operation,Variable})
-
-    rstype = fieldtypes(eltype(maj.reactant_stoch))[2]
-    rs     = Vector{Pair{valtype(states_to_idxs),rstype}}()
-    for (spec,stoich) in maj.reactant_stoch
-        push!(rs, states_to_idxs[spec] => stoich)
+"""
+```julia
+function DiffEqBase.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
+```
+
+Generates a JumpProblem from a JumpSystem.
+"""
+function DiffEqJump.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
+    vrjs = Vector{VariableRateJump}()
+    crjs = Vector{ConstantRateJump}()
+    for j in equations(js)
+        if j isa ConstantRateJump
+            push!(crjs, assemble_crj(js, j))
+        elseif j isa VariableRateJump
+            push!(vrjs, assemble_vrj(js, j))
+        else
+            (j isa MassActionJump) && error("Generation of JumpProblems with MassActionJumps is not yet supported.")
+        end
     end
-
-    nstype = fieldtypes(eltype(maj.net_stoch))[2]
-    ns     = Vector{Pair{valtype(states_to_idxs),nstype}}()
-    for (spec,stoich) in maj.net_stoch
-        push!(ns, states_to_idxs[spec] => stoich)
-    end
-
-    MassActionJump(maj.scaled_rates, rs, ns, scale_rates=scale_rate)
+    ((prob isa DiscreteProblem) && !isempty(vrjs)) && error("Use continuous problems such as an ODEProblem or a SDEProblem with VariableRateJumps")    
+    jset = JumpSet(Tuple(vrjs), Tuple(crjs), nothing, nothing)
+    JumpProblem(prob, aggregator, jset)
 end

test/jumpsystem.jl

@@ -1,4 +1,4 @@
-using ModelingToolkit, DiffEqJump, Test, LinearAlgebra
+using ModelingToolkit, DiffEqBase, DiffEqJump, Test, LinearAlgebra
 MT = ModelingToolkit
 
 # basic SIR model with tweaks
@@ -12,7 +12,7 @@ j₁      = ConstantRateJump(rate₁,affect₁)
 j₂      = VariableRateJump(rate₂,affect₂)
 js      = JumpSystem([j₁,j₂], t, [S,I,R], [β,γ])
 mtjump1 = MT.assemble_crj(js, j₁)
-mtjump2 = MT.assemble_crj(js, j₂)
+mtjump2 = MT.assemble_vrj(js, j₂)
 
 # doc version
 rate1(u,p,t) = (0.1/1000.0)*u[1]*u[2]
@@ -31,16 +31,16 @@ jump2 = VariableRateJump(rate2,affect2!)
 # test crjs
 u = [100, 9, 5]
 p = (0.1/1000,0.01)
-t = 1.0
+tf = 1.0
 mutable struct TestInt
     u
     p
     t
 end
-mtintegrator = TestInt(u,p,t)
-integrator   = TestInt(u,p,t)
-@test abs(mtjump1.rate(u,p,t) - jump1.rate(u,p,t)) < 10*eps()
-@test abs(mtjump2.rate(u,p,t) - jump2.rate(u,p,t)) < 10*eps()
+mtintegrator = TestInt(u,p,tf)
+integrator   = TestInt(u,p,tf)
+@test abs(mtjump1.rate(u,p,tf) - jump1.rate(u,p,tf)) < 10*eps()
+@test abs(mtjump2.rate(u,p,tf) - jump2.rate(u,p,tf)) < 10*eps()
 mtjump1.affect!(mtintegrator)
 jump1.affect!(integrator)
 @test norm(integrator.u - mtintegrator.u) < 10*eps()
@@ -50,3 +50,15 @@ jump2.affect!(integrator)
 @test norm(integrator.u - mtintegrator.u) < 10*eps()
 
 
+# test can make and solve a jump problem
+rate₂   = γ*I
+affect₂ = [I ~ I - 1, R ~ R + 1]
+j₃ = ConstantRateJump(rate₂,affect₂)
+js2 = JumpSystem([j₁,j₃], t, [S,I,R], [β,γ])
+u₀ = [999,1,0]; p = (0.1/1000,0.01); tspan = (0.,250.)
+dprob = DiscreteProblem(u₀,tspan,p)
+jprob = JumpProblem(js2, dprob, Direct())
+sol = solve(jprob, SSAStepper())
+
+using Plots
+plot(sol)