add constant_rate field to JumpProblem

Author: isaacsas

src/JumpProcesses.jl

@@ -121,6 +121,13 @@ include("SSA_stepper.jl")
 export SSAStepper
 
 # leaping: 
+"""
+Aggregator to indicate that individual jumps should also be handled via the leaping
+algorithm that is passed to solve.
+"""
+struct LeapOnly <: AbstractAggregatorAlgorithm end
+export LeapOnly
+
 include("simple_regular_solve.jl")
 export SimpleTauLeaping, EnsembleGPUKernel
 

src/problem.jl

@@ -67,8 +67,8 @@ page](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_e
 the DifferentialEquations.jl [docs](https://docs.sciml.ai/JumpProcesses/stable/) for usage
 examples and commonly asked questions.
 """
-mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggregator}, J2,
-    J3, J4, R, K} <: DiffEqBase.AbstractJumpProblem{P, J}
+mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggregator}, J1, 
+        J2, J3, J4, R, K} <: DiffEqBase.AbstractJumpProblem{P, J}
     """The type of problem to couple the jumps to. For a pure jump process use `DiscreteProblem`, to couple to ODEs, `ODEProblem`, etc."""
     prob::P
     """The aggregator algorithm that determines the next jump times and types for `ConstantRateJump`s and `MassActionJump`s. Examples include `Direct`."""
@@ -77,6 +77,8 @@ mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggrega
     discrete_jump_aggregation::J
     """`CallBackSet` with the underlying `ConstantRate` and `VariableRate` jumps."""
     jump_callback::C
+    """The `ConstantRateJump`s."""
+    constant_jumps::J1
     """The `VariableRateJump`s."""
     variable_jumps::J2
     """The `RegularJump`s."""
@@ -88,10 +90,11 @@ mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggrega
     """kwargs to pass on to solve call."""
     kwargs::K
 end
-function JumpProblem(p::P, a::A, dj::J, jc::C, vj::J2, rj::J3, mj::J4,
-        rng::R, kwargs::K) where {P, A, J, C, J2, J3, J4, R, K}
+function JumpProblem(p::P, a::A, dj::J, jc::C, cj::J1, vj::J2, rj::J3, mj::J4,
+        rng::R, kwargs::K) where {P, A, J, C, J1, J2, J3, J4, R, K}
     iip = isinplace_jump(p, rj)
-    JumpProblem{iip, P, A, C, J, J2, J3, J4, R, K}(p, a, dj, jc, vj, rj, mj, rng, kwargs)
+    JumpProblem{iip, P, A, C, J, J1, J2, J3, J4, R, K}(p, a, dj, jc, cj, vj, rj, mj, 
+        rng, kwargs)
 end
 
 ######## remaking ######
@@ -154,8 +157,8 @@ function DiffEqBase.remake(jprob::JumpProblem; kwargs...)
     end
 
     T(newprob, jprob.aggregator, jprob.discrete_jump_aggregation, jprob.jump_callback,
-        jprob.variable_jumps, jprob.regular_jump, jprob.massaction_jump, jprob.rng,
-        jprob.kwargs)
+        jprob.constant_jumps, jprob.variable_jumps, jprob.regular_jump, 
+        jprob.massaction_jump, jprob.rng, jprob.kwargs)
 end
 
 # for updating parameters in JumpProblems to update MassActionJumps
@@ -253,6 +256,7 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS
     end
 
     ndiscjumps = get_num_majumps(maj) + num_crjs(jumps)
+    crjs = jumps.constant_jumps
 
     # separate bounded variable rate jumps *if* the aggregator can use them
     if use_vrj_bounds && supports_variablerates(aggregator) && (num_bndvrjs(jumps) > 0)
@@ -272,41 +276,35 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS
         disc_agg = nothing
         constant_jump_callback = CallbackSet()
     else
-        disc_agg = aggregate(aggregator, u, prob.p, t, end_time, jumps.constant_jumps, maj,
+        disc_agg = aggregate(aggregator, u, prob.p, t, end_time, crjs, maj,
             save_positions, rng; kwargs...)
         constant_jump_callback = DiscreteCallback(disc_agg)
     end
 
     # handle any remaining vrjs
     if length(cvrjs) > 0
         # Handle variable rate jumps based on vr_aggregator
-        new_prob, variable_jump_callback,
-        cont_agg = configure_jump_problem(prob,
-            vr_aggregator, jumps, cvrjs; rng)
+        new_prob, variable_jump_callback = configure_jump_problem(prob, vr_aggregator, 
+            jumps, cvrjs; rng)
     else
         new_prob = prob
         variable_jump_callback = CallbackSet()
-        cont_agg = JumpSet().variable_jumps
+        cvrjs = JumpSet().variable_jumps
     end
 
     jump_cbs = CallbackSet(constant_jump_callback, variable_jump_callback)
     iip = isinplace_jump(prob, jumps.regular_jump)
     solkwargs = make_kwarg(; callback)
 
-    JumpProblem{iip, typeof(new_prob), typeof(aggregator),
-        typeof(jump_cbs), typeof(disc_agg),
-        typeof(cont_agg),
-        typeof(jumps.regular_jump),
+    JumpProblem{iip, typeof(new_prob), typeof(aggregator), typeof(jump_cbs), 
+        typeof(disc_agg), typeof(crjs), typeof(cvrjs), typeof(jumps.regular_jump),
         typeof(maj), typeof(rng), typeof(solkwargs)}(new_prob, aggregator, disc_agg,
-        jump_cbs, cont_agg,
-        jumps.regular_jump, maj, rng,
-        solkwargs)
+        jump_cbs, crjs, cvrjs, jumps.regular_jump, maj, rng, solkwargs)
 end
 
-aggregator(jp::JumpProblem{iip, P, A, C, J}) where {iip, P, A, C, J} = A
+aggregator(jp::JumpProblem{iip, P, A}) where {iip, P, A} = A
 
-@inline function extend_tstops!(tstops,
-        jp::JumpProblem{P, A, C, J, J2}) where {P, A, C, J, J2}
+@inline function extend_tstops!(tstops, jp::JumpProblem) 
     !(jp.jump_callback.discrete_callbacks isa Tuple{}) &&
         push!(tstops, jp.jump_callback.discrete_callbacks[1].condition.next_jump_time)
 end

src/variable_rate.jl

@@ -62,8 +62,7 @@ function configure_jump_problem(prob, vr_aggregator::VR_FRM, jumps, cvrjs;
         rng = DEFAULT_RNG)
     new_prob = extend_problem(prob, cvrjs; rng)
     variable_jump_callback = build_variable_callback(CallbackSet(), 0, cvrjs...; rng)
-    cont_agg = cvrjs
-    return new_prob, variable_jump_callback, cont_agg
+    return new_prob, variable_jump_callback
 end
 
 # extends prob.u0 to an ExtendedJumpArray with Njumps integrated intensity values,
@@ -381,16 +380,14 @@ function configure_jump_problem(prob, ::VR_Direct, jumps, cvrjs; rng = DEFAULT_R
     new_prob = prob
     cache = VR_DirectEventCache(jumps, VR_Direct(), prob, eltype(prob.tspan); rng)
     variable_jump_callback = build_variable_integcallback(cache, cvrjs)
-    cont_agg = cvrjs
-    return new_prob, variable_jump_callback, cont_agg
+    return new_prob, variable_jump_callback
 end
 
 function configure_jump_problem(prob, ::VR_DirectFW, jumps, cvrjs; rng = DEFAULT_RNG)
     new_prob = prob
     cache = VR_DirectEventCache(jumps, VR_DirectFW(), prob, eltype(prob.tspan); rng)
     variable_jump_callback = build_variable_integcallback(cache, cvrjs)
-    cont_agg = cvrjs
-    return new_prob, variable_jump_callback, cont_agg
+    return new_prob, variable_jump_callback
 end
 
 # recursively evaluate the cumulative sum of the rates for type stability