Merge pull request #519 from isaacsas/cleanup_leaping_and_support_nonregjumps

Cleanup exports and add constant rate jumps to JumpProblem

Author: isaacsas

src/JumpProcesses.jl

@@ -56,73 +56,94 @@ const USE_RSSA_THRESHOLD = 100
 const USE_SORTINGDIRECT_THRESHOLD = 200
 
 include("jumps.jl")
+export ConstantRateJump, VariableRateJump, RegularJump, MassActionJump, JumpSet
+
 include("massaction_rates.jl")
+
+# constant rate aggregators (i.e. SSAs)
 include("aggregators/aggregators.jl")
+export get_num_majumps, needs_depgraph, needs_vartojumps_map, reset_aggregated_jumps!
+
 include("aggregators/ssajump.jl")
+
 include("aggregators/direct.jl")
+export Direct, DirectFW
+
 include("aggregators/frm.jl")
+export FRM, FRMFW
+
 include("aggregators/sortingdirect.jl")
+export SortingDirect
+
 include("aggregators/nrm.jl")
+export NRM
+
 include("aggregators/bracketing.jl")
+export BracketData
+
 include("aggregators/rssa.jl")
+export RSSA
+
 include("aggregators/prioritytable.jl")
+
 include("aggregators/directcr.jl")
+export DirectCR
+
 include("aggregators/rssacr.jl")
+export RSSACR
+
 include("aggregators/rdirect.jl")
-include("aggregators/coevolve.jl")
-include("aggregators/ccnrm.jl")
+export RDirect
 
-# spatial:
-include("spatial/spatial_massaction_jump.jl")
-include("spatial/topology.jl")
-include("spatial/hop_rates.jl")
-include("spatial/reaction_rates.jl")
-include("spatial/flatten.jl")
-include("spatial/utils.jl")
-include("spatial/bracketing.jl")
+include("aggregators/coevolve.jl")
+export Coevolve
 
-include("spatial/nsm.jl")
-include("spatial/directcrdirect.jl")
+include("aggregators/ccnrm.jl")
+export CCNRM
 
 include("aggregators/aggregated_api.jl")
 
+# variable rate aggregators (i.e. SSAs)
 include("extended_jump_array.jl")
-include("variable_rate.jl")
-include("problem.jl")
-include("solve.jl")
-include("coupled_array.jl")
-include("coupling.jl")
-include("SSA_stepper.jl")
-include("simple_regular_solve.jl")
+export ExtendedJumpArray
 
-export ConstantRateJump, VariableRateJump, RegularJump,
-       MassActionJump, JumpSet
+include("variable_rate.jl")
+export VariableRateAggregator, VR_FRM, VR_Direct, VR_DirectFW
 
-export JumpProblem
+# core problem and timestepping
+include("problem.jl")
+export JumpProblem, SplitCoupledJumpProblem
 
-export SplitCoupledJumpProblem
+include("solve.jl")
+export init, solve, solve!
 
-export Direct, DirectFW, SortingDirect, DirectCR
-export BracketData, RSSA
-export FRM, FRMFW, NRM, CCNRM
-export RSSACR, RDirect
-export Coevolve
+include("SSA_stepper.jl")
+export SSAStepper
 
-export get_num_majumps, needs_depgraph, needs_vartojumps_map
+# leaping: 
+include("simple_regular_solve.jl")
+export SimpleTauLeaping, EnsembleGPUKernel
 
-export init, solve, solve!
+# spatial:
+include("spatial/spatial_massaction_jump.jl")
+export SpatialMassActionJump
 
-export reset_aggregated_jumps!
+include("spatial/topology.jl")
+export CartesianGrid, CartesianGridRej, outdegree, num_sites, neighbors
 
-export ExtendedJumpArray
+include("spatial/hop_rates.jl")
+include("spatial/reaction_rates.jl")
+include("spatial/flatten.jl")
+include("spatial/utils.jl")
+include("spatial/bracketing.jl")
+include("spatial/nsm.jl")
+export NSM
 
-# Export VariableRateAggregator types
-export VariableRateAggregator, VR_FRM, VR_Direct, VR_DirectFW
+include("spatial/directcrdirect.jl")
+export DirectCRDirect
 
-# spatial structs and functions
-export CartesianGrid, CartesianGridRej
-export SpatialMassActionJump
-export outdegree, num_sites, neighbors
-export NSM, DirectCRDirect
+# coupling
+include("coupled_array.jl")
+include("coupling.jl")
 
 end # module

src/SSA_stepper.jl

@@ -381,7 +381,6 @@ function DiffEqBase.terminate!(integrator::SSAIntegrator, retcode = ReturnCode.T
     nothing
 end
 
-export SSAStepper
 
 function SciMLBase.isdenseplot(sol::ODESolution{
         T, N, uType, uType2, DType, tType, rateType, discType, P,

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/simple_regular_solve.jl

@@ -63,4 +63,3 @@ function EnsembleGPUKernel()
     EnsembleGPUKernel(nothing, 0.0)
 end
 
-export SimpleTauLeaping, EnsembleGPUKernel

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