ODEProblem for vrjs and add back test

Author: isaacsas

src/systems/jumps/jumpsystem.jl

@@ -311,7 +311,7 @@ end
 ```julia
 DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan,
                            parammap = DiffEqBase.NullParameters;
-                           use_union = false,
+                           use_union = true,
                            kwargs...)
 ```
 
@@ -331,7 +331,6 @@ dprob = DiscreteProblem(complete(js), u₀map, tspan, parammap)
 """
 function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
         parammap = DiffEqBase.NullParameters();
-        checkbounds = false,
         use_union = true,
         eval_expression = false,
         eval_module = @__MODULE__,
@@ -385,7 +384,7 @@ struct DiscreteProblemExpr{iip} end
 
 function DiscreteProblemExpr{iip}(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
         parammap = DiffEqBase.NullParameters();
-        use_union = false,
+        use_union = true,
         kwargs...) where {iip}
     if !iscomplete(sys)
         error("A completed `JumpSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `DiscreteProblemExpr`")
@@ -412,6 +411,61 @@ function DiscreteProblemExpr{iip}(sys::JumpSystem, u0map, tspan::Union{Tuple, No
     end
 end
 
+"""
+```julia
+DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan,
+                           parammap = DiffEqBase.NullParameters;
+                           use_union = true,
+                           kwargs...)
+```
+
+Generates a blank ODEProblem for a pure jump JumpSystem to utilize as its `prob.prob`. This
+is used in the case where there are no ODEs and no SDEs associated with the system but there
+are jumps with an explicit time dependency (i.e. `VariableRateJump`s). If no jumps have an
+explicit time dependence, i.e. all are `ConstantRateJump`s or `MassActionJump`s then
+`DiscreteProblem` should be preferred for performance reasons.
+
+Continuing the example from the [`JumpSystem`](@ref) definition:
+
+```julia
+using DiffEqBase, JumpProcesses
+u₀map = [S => 999, I => 1, R => 0]
+parammap = [β => 0.1 / 1000, γ => 0.01]
+tspan = (0.0, 250.0)
+oprob = ODEProblem(complete(js), u₀map, tspan, parammap)
+```
+"""
+function DiffEqBase.ODEProblem(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing},
+        parammap = DiffEqBase.NullParameters();
+        use_union = true,
+        eval_expression = false,
+        eval_module = @__MODULE__,
+        kwargs...)
+    if !iscomplete(sys)
+        error("A completed `JumpSystem` is required. Call `complete` or `structural_simplify` on the system before creating a `DiscreteProblem`")
+    end
+    dvs = unknowns(sys)
+    ps = parameters(sys)
+
+    defs = defaults(sys)
+    defs = mergedefaults(defs, parammap, ps)
+    defs = mergedefaults(defs, u0map, dvs)
+
+    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false)
+    if has_index_cache(sys) && get_index_cache(sys) !== nothing
+        p = MTKParameters(sys, parammap, u0map; eval_expression, eval_module)
+    else
+        p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = false, use_union)
+    end
+
+    observedfun = ObservedFunctionCache(sys; eval_expression, eval_module)
+
+    f = (du, u, p, t) -> (du .= 0; nothing)
+    df = ODEFunction(f; sys = sys, observed = observedfun)
+    ODEProblem(df, u0, tspan, p; kwargs...)
+end
+
+
 """
 ```julia
 DiffEqBase.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)

test/jumpsystem.jl

@@ -67,7 +67,7 @@ tspan = (0.0, 250.0);
 u₀map = [S => 999, I => 1, R => 0]
 parammap = [β => 0.1 / 1000, γ => 0.01]
 dprob = DiscreteProblem(js2, u₀map, tspan, parammap)
-jprob = JumpProblem(js2, dprob, Direct(), save_positions = (false, false), rng = rng)
+jprob = JumpProblem(js2, dprob, Direct(), save_positions = (false, false), rng)
 Nsims = 30000
 function getmean(jprob, Nsims; use_stepper = true)
     m = 0.0
@@ -89,12 +89,12 @@ obs = [S2 ~ 2 * S]
 @named js2b = JumpSystem([j₁, j₃], t, [S, I, R], [β, γ], observed = obs)
 js2b = complete(js2b)
 dprob = DiscreteProblem(js2b, u₀map, tspan, parammap)
-jprob = JumpProblem(js2b, dprob, Direct(), save_positions = (false, false), rng = rng)
+jprob = JumpProblem(js2b, dprob, Direct(), save_positions = (false, false), rng)
 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), rng = rng)
+jprob = JumpProblem(js2, dprob, Direct(), save_positions = (false, false), rng)
 sol = solve(jprob, SSAStepper(), saveat = 1.0)
 @test all((sol.t) .== collect(0.0:tspan[2]))
 
@@ -129,7 +129,7 @@ function a2!(integrator)
 end
 j2 = ConstantRateJump(r2, a2!)
 jset = JumpSet((), (j1, j2), nothing, nothing)
-jprob = JumpProblem(prob, Direct(), jset, save_positions = (false, false), rng = rng)
+jprob = JumpProblem(prob, Direct(), jset, save_positions = (false, false), rng)
 m2 = getmean(jprob, Nsims)
 
 # test JumpSystem solution agrees with direct version
@@ -141,17 +141,17 @@ maj2 = MassActionJump(γ, [I => 1], [I => -1, R => 1])
 @named js3 = JumpSystem([maj1, maj2], t, [S, I, R], [β, γ])
 js3 = complete(js3)
 dprob = DiscreteProblem(js3, u₀map, tspan, parammap)
-jprob = JumpProblem(js3, dprob, Direct(), rng = rng)
+jprob = JumpProblem(js3, dprob, Direct(), rng)
 m3 = getmean(jprob, Nsims)
 @test abs(m - m3) / m < 0.01
 
 # maj jump test with various dep graphs
 @named js3b = JumpSystem([maj1, maj2], t, [S, I, R], [β, γ])
 js3b = complete(js3b)
-jprobb = JumpProblem(js3b, dprob, NRM(), rng = rng)
+jprobb = JumpProblem(js3b, dprob, NRM(), rng)
 m4 = getmean(jprobb, Nsims)
 @test abs(m - m4) / m < 0.01
-jprobc = JumpProblem(js3b, dprob, RSSA(), rng = rng)
+jprobc = JumpProblem(js3b, dprob, RSSA(), rng)
 m4 = getmean(jprobc, Nsims)
 @test abs(m - m4) / m < 0.01
 
@@ -161,7 +161,7 @@ maj2 = MassActionJump(γ, [S => 1], [S => -1])
 @named js4 = JumpSystem([maj1, maj2], t, [S], [β, γ])
 js4 = complete(js4)
 dprob = DiscreteProblem(js4, [S => 999], (0, 1000.0), [β => 100.0, γ => 0.01])
-jprob = JumpProblem(js4, dprob, Direct(), rng = rng)
+jprob = JumpProblem(js4, dprob, Direct(), rng)
 m4 = getmean(jprob, Nsims)
 @test abs(m4 - 2.0 / 0.01) * 0.01 / 2.0 < 0.01
 
@@ -171,7 +171,7 @@ maj2 = MassActionJump(γ, [S => 2], [S => -1])
 @named js4 = JumpSystem([maj1, maj2], t, [S], [β, γ])
 js4 = complete(js4)
 dprob = DiscreteProblem(js4, [S => 999], (0, 1000.0), [β => 100.0, γ => 0.01])
-jprob = JumpProblem(js4, dprob, Direct(), rng = rng)
+jprob = JumpProblem(js4, dprob, Direct(), rng)
 sol = solve(jprob, SSAStepper());
 
 # issue #819
@@ -183,28 +183,30 @@ sol = solve(jprob, SSAStepper());
 end
 
 # test if param mapper is setup correctly for callbacks
-@parameters k1 k2 k3
-@variables A(t) B(t)
-maj1 = MassActionJump(k1 * k3, [0 => 1], [A => -1, B => 1])
-maj2 = MassActionJump(k2, [B => 1], [A => 1, B => -1])
-@named js5 = JumpSystem([maj1, maj2], t, [A, B], [k1, k2, k3])
-js5 = complete(js5)
-p = [k1 => 2.0, k2 => 0.0, k3 => 0.5]
-u₀ = [A => 100, B => 0]
-tspan = (0.0, 2000.0)
-dprob = DiscreteProblem(js5, u₀, tspan, p)
-jprob = JumpProblem(js5, dprob, Direct(), save_positions = (false, false), rng = rng)
-@test all(jprob.massaction_jump.scaled_rates .== [1.0, 0.0])
+let
+    @parameters k1 k2 k3
+    @variables A(t) B(t)
+    maj1 = MassActionJump(k1 * k3, [0 => 1], [A => -1, B => 1])
+    maj2 = MassActionJump(k2, [B => 1], [A => 1, B => -1])
+    @named js5 = JumpSystem([maj1, maj2], t, [A, B], [k1, k2, k3])
+    js5 = complete(js5)
+    p = [k1 => 2.0, k2 => 0.0, k3 => 0.5]
+    u₀ = [A => 100, B => 0]
+    tspan = (0.0, 2000.0)
+    dprob = DiscreteProblem(js5, u₀, tspan, p)
+    jprob = JumpProblem(js5, dprob, Direct(); save_positions = (false, false), rng)
+    @test all(jprob.massaction_jump.scaled_rates .== [1.0, 0.0])
 
-pcondit(u, t, integrator) = t == 1000.0
-function paffect!(integrator)
-    integrator.ps[k1] = 0.0
-    integrator.ps[k2] = 1.0
-    reset_aggregated_jumps!(integrator)
+    pcondit(u, t, integrator) = t == 1000.0
+    function paffect!(integrator)
+        integrator.ps[k1] = 0.0
+        integrator.ps[k2] = 1.0
+        reset_aggregated_jumps!(integrator)
+    end
+    cb = DiscreteCallback(pcondit, paffect!)
+    sol = solve(jprob, SSAStepper(); tstops = [1000.0], callback = cb)
+    @test sol.u[end][1] == 100
 end
-sol = solve(jprob, SSAStepper(), tstops = [1000.0],
-    callback = DiscreteCallback(pcondit, paffect!))
-@test_skip sol.u[end][1] == 100 # TODO: Fix mass-action jumps in JumpProcesses
 
 # observed variable handling
 @variables OBS(t)