diff --git a/.github/workflows/FormatCheck.yml b/.github/workflows/FormatCheck.yml index f80d0b18..329ae654 100644 --- a/.github/workflows/FormatCheck.yml +++ b/.github/workflows/FormatCheck.yml @@ -27,7 +27,7 @@ jobs: # # julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="0.13.0"))' run: | - julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter"))' + julia -e 'using Pkg; Pkg.add(PackageSpec(name="JuliaFormatter", version="1.0.32"))' julia -e 'using JuliaFormatter; format(".", verbose=true)' - name: Format check run: | diff --git a/docs/src/jump_solve.md b/docs/src/jump_solve.md index 68a5dddc..41ac4412 100644 --- a/docs/src/jump_solve.md +++ b/docs/src/jump_solve.md @@ -49,7 +49,19 @@ algorithms are optimized for pure jump problems. - `SSAStepper`: a stepping integrator for `JumpProblem`s defined over `DiscreteProblem`s involving `ConstantRateJump`s, `MassActionJump`s, and/or bounded `VariableRateJump`s . Supports handling of `DiscreteCallback`s and - saving controls like `saveat`. + saving controls like `saveat`. Note that DifferentialEquations.jl treats + jumps as similar to callbacks, and hence `SSAStepper` only implements a + subset of ODE/SDE solver saving controls. In particular, `save_everystep` is + not supported as saving jumps each step is controlled via the + `save_positions` argument to [`JumpProblem`](@ref)s. Note, in contrast to + when [`JumpProblem`](@ref)s are coupled with ODE and SDE timesteppers, with + [`SSAStepper`](@ref), setting `save_positions = (false, true)`, + `save_positions = (true, false)` or `save_positions = (true, true)` are + equivalent and save only after the jump has occurred (as opposed to saving + the state both before and after a jump). This is because the underlying + [`SSAStepper`](@ref) generated-solution uses piecewise constant + interpolation, and can therefore exactly evaluate the sampled solution + path at any time when only saving the post-jump state for each jump. ## RegularJump Compatible Methods diff --git a/docs/src/jump_types.md b/docs/src/jump_types.md index 721414cd..69115150 100644 --- a/docs/src/jump_types.md +++ b/docs/src/jump_types.md @@ -212,15 +212,6 @@ must also specify - `rateinterval(u, p, t)`, a function which computes a time interval `t` to `t + rateinterval(u, p, t)` given state `u` and parameters `p` over which the `urate` bound will hold (and `lrate` bound if provided, see below). -Note that it is ok if the `urate` bound would be violated within the -`rateinterval` due to a change in `u` arising from another `ConstantRateJump`, -`MassActionJump` or *bounded* `VariableRateJump` being executed, as the chosen -aggregator will then handle recalculating the rate bound and interval. *However, -if the bound could be violated within the time interval due to a change in `u` -arising from continuous dynamics such as a coupled ODE, SDE, or a general -`VariableRateJump`, bounds should not be given.* This ensures the jump is -classified as a general `VariableRateJump` and properly handled. - For increased performance, one can also specify a lower bound that should be valid over the same `rateinterval` @@ -232,14 +223,32 @@ valid over the same `rateinterval` Note that - It is currently only possible to simulate `VariableRateJump`s with - `SSAStepper` when using systems with only bounded `VariableRateJump`s and the + `SSAStepper` when using systems with bounded `VariableRateJump`s and the `Coevolve` aggregator. - - When choosing a different aggregator than `Coevolve`, `SSAStepper` cannot - currently be used, and the `JumpProblem` must be coupled to a continuous - problem type such as an `ODEProblem` to handle time-stepping. The continuous - time-stepper treats *all* `VariableRateJump`s as `ContinuousCallback`s, using - the `rate(u, p, t)` function to construct the `condition` function that - triggers a callback. + - Any `JumpProblem` with `VariableRateJump` that *does not use* the + `Coevolve` aggregator must be coupled to a continuous problem type such as + an `ODEProblem` to handle time-stepping. Continuous time-stepper will ignore + the provided aggregator and treat *all* `VariableRateJump`s as + `ContinuousCallback`s, using the `rate(u, p, t)` function to construct the + `condition` function that triggers a callback. + - When using `Coevolve` with a `JumpProblem` coupled to a continuous problem + such as an `ODEProblem`, the aggregator will handle the jumps in same way + that it does with `SSAStepper`. However, *ensure that given `t` the bounds + will hold for the duration of `rateinterval(t)` for the full coupled system's + dynamics or the algorithm will not give correct samples*. Numerical and + analytical solutions are generally not guaranteed to satisfy the same bounds, + especially in large complicated models. Consider adding some slack on the + bounds and approach complex models with care. In most simple cases the bounds + should be close enough. For debugging purposes one might want to add safety + checks in the bound functions. + - In some circumstances with complex model of many variables it can be + difficult to determine good a priori bounds on the ODE variables. For some + discussion on the bound setting problem see [1]. + +[1] V. Lemaire, M. Thieullen and N. Thomas, Exact Simulation of the Jump +Times of a Class of Piecewise Deterministic Markov Processes, Journal of +Scientific Computing, 75 (3), 1776-1807 (2018). +doi:10.1007/s10915-017-0607-4. #### Defining a Regular Jump @@ -331,10 +340,12 @@ aggregator requires various types of dependency graphs, see the next section): - *`NRM`*: The Gibson-Bruck Next Reaction Method [7]. For some reaction network structures, this may offer better performance than `Direct` (for example, large, linear chains of reactions). - - *`Coevolve`*: An adaptation of the COEVOLVE algorithm of Farajtabar et al [8]. - Currently the only aggregator that also supports *bounded* - `VariableRateJump`s. Essentially reduces to `NRM` in handling - `ConstantRateJump`s and `MassActionJump`s. + - *`Coevolve`*: An improvement of the COEVOLVE algorithm of Farajtabar et al + [8]. Currently the only aggregator that also supports *bounded* + `VariableRateJump`s. As opposed + to COEVOLVE, this method syncs the thinning procedure with the stepper + which allows it to handle dependencies on continuous dynamics. Essentially + reduces to `NRM` in handling `ConstantRateJump`s and `MassActionJump`s. To pass the aggregator, pass the instantiation of the type. For example: @@ -345,7 +356,7 @@ JumpProblem(prob, Direct(), jump1, jump2) will build a problem where the jumps are simulated using Gillespie's Direct SSA method. -[1] Daniel T. Gillespie, A general method for numerically simulating the stochastic +[1] D. T. Gillespie, A general method for numerically simulating the stochastic time evolution of coupled chemical reactions, Journal of Computational Physics, 22 (4), 403–434 (1976). doi:10.1016/0021-9991(76)90041-3. @@ -431,10 +442,12 @@ For representing and aggregating jumps jumps. - Use `VariableRateJump`s for any remaining jumps with variable rate between jumps. If possible, construct a bounded [`VariableRateJump`](@ref) as - described above and in the doc string. The tighter and easier to compute the - bounds are, the faster the resulting simulation will be. Use the `Coevolve` - aggregator to ensure such jumps are handled via the more efficient aggregator - interface. + described above and in the doc string. The tighter and easier to compute + the bounds are, the faster the resulting simulation will be. Use the + `Coevolve` aggregator to ensure such jumps are handled via the more + efficient aggregator interface. `Coevolve` handles continuous steppers so + can be coupled with a continuous problem type such as an `ODEProblem` as + long as the bounds are satisfied given changes in `u` over `rateinterval`. For systems with only `ConstantRateJump`s and `MassActionJump`s, diff --git a/docs/src/tutorials/discrete_stochastic_example.md b/docs/src/tutorials/discrete_stochastic_example.md index 3e0cff37..30ce297a 100644 --- a/docs/src/tutorials/discrete_stochastic_example.md +++ b/docs/src/tutorials/discrete_stochastic_example.md @@ -596,10 +596,10 @@ jump4 = ConstantRateJump(rate4, affect4!) With the jumps defined, we can build a [`DiscreteProblem`](https://docs.sciml.ai/DiffEqDocs/stable/types/discrete_types/). Bounded `VariableRateJump`s over a `DiscreteProblem` can currently only be -simulated with the `Coevolve` aggregator. The aggregator requires a dependency -graph to indicate when a given jump occurs which other jumps in the system -should have their rate recalculated (i.e., their rate depends on states modified -by one occurrence of the first jump). This ensures that rates, rate bounds, and +simulated with `Coevolve`. The aggregator requires a dependency graph to +indicate when a given jump occurs and which other jumps in the system should +have their rate recalculated (i.e., their rate depends on states modified by +one occurrence of the first jump). This ensures that rates, rate bounds, and rate intervals are recalculated when invalidated due to changes in `u`. For the current example, both processes mutually affect each other, so we have @@ -628,11 +628,13 @@ We see that the time-dependent infection rate leads to a lower peak of the infection throughout the population. Note that bounded `VariableRateJump`s over `DiscreteProblem`s can be quite -general, but it is not possible to handle rates that change according to an -ODE/SDE modified variable. A rate such as `p[2]*u[1]*u[4]` when `u[4]` is the -solution of a continuous problem such as an ODE or SDE can only be handled using -a general `VariableRateJump` within a continuous integrator as discussed -[below](@ref VariableRateJumpSect). +general. However, when handling rates that change according to an ODE/SDE +modified variable we will need a continuous integrator. One example is +discussed [below](@ref VariableRateJumpSect) in which we have a new reaction +added to the model with rate `p[2]*u[1]*u[4]` where `u[4]` is the solution of +an ODE. In such models, you will also need to be more careful in setting +rate bounds as they must be valid for the full coupled system's +dynamics. ## [Reducing Memory Use: Controlling Saving Behavior](@id save_positions_docs) @@ -701,7 +703,7 @@ plot(sol; label = ["S(t)" "I(t)" "R(t)"]) ``` Note that we can combine `MassActionJump`s, `ConstantRateJump`s and bounded -`VariableRateJump`s using the `Coevolve` aggregator. +`VariableRateJump`s using the `Coevolve` aggregators. ## Adding Jumps to a Differential Equation @@ -713,7 +715,7 @@ only acts on some new 4th component: ```@example tut2 using OrdinaryDiffEq function f(du, u, p, t) - du[4] = u[2] * u[3] / 100000 - u[1] * u[4] / 100000 + du[4] = u[2] * u[3] / 1e5 - u[1] * u[4] / 1e5 nothing end u₀ = [999.0, 10.0, 0.0, 100.0] @@ -733,13 +735,13 @@ sol = solve(jump_prob, Tsit5()) plot(sol; label = ["S(t)" "I(t)" "R(t)" "u₄(t)"]) ``` -Note, when using `ConstantRateJump`s, `MassActionJump`s, and bounded -`VariableRateJump`s, the ODE derivative function `f(du, u, p, t)` should not -modify any states in `du` that the corresponding jump rate functions depend on. -However, the opposite where jumps modify the ODE variables is allowed. If one -needs to change a component of `u` in the ODE for which a rate function is -dependent, then one must use a general `VariableRateJump` as described in the -next section. +Note that in general, the ODE derivative `f(du, u, p, t)` could modify any +element in `du` which the jump rate functions depend on. In this section, +`f(du, u, p, t)` does not modify the jump rates so it is safe to couple them with any +type of jump and use any type of aggregator. However, the implementation does +not enforce this requirement, so one must be careful. Alternatively, when `f(du, u, p, t)` *does* modify variables that affect the jump rate, we have to +implement another strategy as described in the next [next Section](@ref +VariableRateJumpSect]. ## [Adding a general VariableRateJump that Depends on a Continuously Evolving Variable](@id VariableRateJumpSect) @@ -758,10 +760,10 @@ jump5 = VariableRateJump(rate5, affect5!) ``` Notice, since `rate5` depends on a variable that evolves continuously, and hence -is not constant between jumps, *we must use a general `VariableRateJump` without -upper/lower bounds*. +is not constant between jumps, *we must either use a general `VariableRateJump` without +upper/lower bounds or a bounded `VariableRateJump`*. -Solving the equation is exactly the same: +In the general case, solving the equation is exactly the same: ```@example tut2 u₀ = [999.0, 10.0, 0.0, 1.0] @@ -774,6 +776,72 @@ plot(sol; label = ["S(t)" "I(t)" "R(t)" "u₄(t)"]) *Note that general `VariableRateJump`s require using a continuous problem, like an ODE/SDE/DDE/DAE problem, and using floating point initial conditions.* +Alternatively, the case of bounded `VariableRateJump` requires some maths. +First, we need to obtain the upper bounds of `rate5` at time `t` given `u`. +Note that `rate5` evolves according to `u[4]` which is a separable first order +differential equation of the form ``x' = b - a x`` with general solution: + +```math +x(t) = - \frac{e^{-a t - c_1 a}}{a} + \frac{b}{a} +``` + +This is bounded by ``b / a`` which is too high for our purposes since it would +lead to a high rate of rejection during sampling. However, since the function +is increasing we can compute the upper bound given an interval ``\Delta t`` +as following: + +```math +\bar{x}(s) = x(t) \, e^{-a (t + \Delta t)} + \frac{b}{a} (1 - e^{- a (t + \Delta t)}) \text{ , } \forall s \in [t, t + \Delta t] +``` + +However, when ``a = 0`` the differential equation becomes ``x' = b`` whose solution is ``x(t) = b t``. In which case, we obtain a different upper bound given by: + +```math +\bar{x}(s) = x(t) + b * (t + \Delta t) \text{ , } \forall s \in [t, t + \Delta t] +``` + +These expressions allow us to write the upper-bound and the rate interval in +Julia. In this example we use analytical boundaries for *illustration +purposes*. However, in some circumstances with complex model of many variables +it can be difficult to determine good a priori bounds on the ODE variables. +Moreover, numerical and analytical solutions are generally not guaranteed to +strictly satisfy the same bounds. In most cases the bounds should be close +enough. Thus, consider adding some slack on the bounds and approach complex +models with care. + +```@example tut2 +function urate2(u, p, t) + if u[1] > 0 + 1e-2 * max(u[4], + (u[4] * exp(-1 * u[1] / 1e5) + + (u[2] * u[3] / u[1]) * (1 - exp(-1 * u[1] / 1e5)))) + else + 1e-2 * (u[4] + 1 * u[2] * u[3] / 1e5) + end +end +rateinterval2(u, p, t) = 1 +``` + +We can then formulate the jump problem. The only aggregator that supports +bounded `VariableRateJump`s is `Coevolve`. We formulate and solve the +jump problem with this aggregator. `Coevolve` can be formulated as either +a discrete or continuous problem. In this case, we must formulate the problem +as continuous as it depends on a continuous variable. + +```@example tut2 +jump6 = VariableRateJump(rate5, affect5!; urate = urate2, rateinterval = rateinterval2) +dep_graph2 = [[1, 2, 3], [1, 2, 3], [1, 2, 3]] +jump_prob = JumpProblem(prob, Coevolve(), jump, jump2, jump6; dep_graph = dep_graph2) +sol = solve(jump_prob, Tsit5()) +plot(sol; label = ["S(t)" "I(t)" "R(t)" "u₄(t)"]) +``` + +We obtain the same solution as with `Direct`, but `Coevolve` runs faster +because it doesn't need to compute the derivative of `rate5`. Each aggregator +faces a different trade-off, so the the choice of best aggregator will depend +on the problem at hand. `Coevolve` requires a good understanding of the +equations involved, passing a wrong boundary can result in silent bugs. + Lastly, we are not restricted to ODEs. For example, we can solve the same jump problem except with multiplicative noise on `u[4]` by using an `SDEProblem` instead: diff --git a/docs/src/tutorials/jump_diffusion.md b/docs/src/tutorials/jump_diffusion.md index ead2d0f0..40848279 100644 --- a/docs/src/tutorials/jump_diffusion.md +++ b/docs/src/tutorials/jump_diffusion.md @@ -147,11 +147,6 @@ plot(sol) In this way we have solved a mixed jump-ODE, i.e., a piecewise deterministic Markov process. -Note that in this case, the rates of the `VariableRateJump`s depend on a -variable that is driven by an `ODEProblem`, and thus they would not satisfy the -conditions to be represented as bounded `VariableRateJump`s (and hence cannot -be simulated with the `Coevolve` aggregator). - ## Jump Diffusion Now we will finally solve the jump diffusion problem. The steps are the same diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl index 7b1915f7..d0476921 100644 --- a/src/JumpProcesses.jl +++ b/src/JumpProcesses.jl @@ -22,6 +22,8 @@ abstract type AbstractJump end abstract type AbstractMassActionJump <: AbstractJump end abstract type AbstractAggregatorAlgorithm end abstract type AbstractJumpAggregator end +abstract type AbstractSSAIntegrator{Alg, IIP, U, T} <: + DiffEqBase.DEIntegrator{Alg, IIP, U, T} end import Base.Threads @static if VERSION < v"1.3" diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index 1b9b5eeb..50a91128 100644 --- a/src/SSA_stepper.jl +++ b/src/SSA_stepper.jl @@ -8,7 +8,17 @@ Highly efficient integrator for pure jump problems that involve only `ConstantRa - Only works with `JumpProblem`s defined from `DiscreteProblem`s. - Only works with collections of `ConstantRateJump`s, `MassActionJump`s, and `VariableRateJump`s with rate bounds. -- Only supports `DiscreteCallback`s for events. +- Only supports `DiscreteCallback`s for events, which are executed for every step + taken by `SSAStepper`. `Coevolve` may take a number of time steps larger + than the number of jumps when simulating `VariableRateJump`s. All the other + aggregators take a number of steps equal to the number of jumps. +- Only supports a limited subset of the output controls from the common solver + and `DiscreteCallback`. In particular, the options `save_positions = (false, + true)`, `save_positions = (true, false)` or `save_positions = (true, true)` are + equivalent and will save every jump after it has occurred. Alternatively, the + option `save_everystep` from the common solver is silently ignored without any + effect on saving behaviour. Finally, `saveat` behaves the same way as in the + common solver. ## Examples SIR model: @@ -51,7 +61,7 @@ Solution objects for pure jump problems solved via `SSAStepper`. $(FIELDS) """ mutable struct SSAIntegrator{F, uType, tType, tdirType, P, S, CB, SA, OPT, TS} <: - DiffEqBase.DEIntegrator{SSAStepper, Nothing, uType, tType} + AbstractSSAIntegrator{SSAStepper, Nothing, uType, tType} """The underlying `prob.f` function. Not currently used.""" f::F """The current solution values.""" @@ -258,10 +268,13 @@ function DiffEqBase.step!(integrator::SSAIntegrator) # FP error means the new time may equal the old if the next jump time is # sufficiently small, hence we add this check to execute jumps until # this is no longer true. + integrator.u_modified = true while integrator.t == integrator.tstop doaffect && integrator.cb.affect!(integrator) end + jump_modified_u = integrator.u_modified + if !(typeof(integrator.opts.callback.discrete_callbacks) <: Tuple{}) discrete_modified, saved_in_cb = DiffEqBase.apply_discrete_callback!(integrator, integrator.opts.callback.discrete_callbacks...) @@ -269,7 +282,7 @@ function DiffEqBase.step!(integrator::SSAIntegrator) saved_in_cb = false end - !saved_in_cb && savevalues!(integrator) + !saved_in_cb && jump_modified_u && savevalues!(integrator) nothing end diff --git a/src/aggregators/aggregators.jl b/src/aggregators/aggregators.jl index 0d268e08..c1553d03 100644 --- a/src/aggregators/aggregators.jl +++ b/src/aggregators/aggregators.jl @@ -120,10 +120,12 @@ systems with many species and many channels, Journal of Physical Chemistry A, struct NRM <: AbstractAggregatorAlgorithm end """ -An adaptation of the COEVOLVE algorithm for simulating any compound jump process -that evolves through time. This method handles variable intensity rates with -user-defined bounds and inter-dependent processes. It reduces to NRM when rates -are constant. +An improvement of the COEVOLVE algorithm for simulating any compound jump +process that evolves through time. This method handles variable intensity +rates with user-defined bounds and inter-dependent processes. It reduces to +NRM when rates are constant. As opposed to COEVOLVE, this method syncs the +thinning procedure with the stepper which allows it to handle dependencies on +continuous dynamics. M. Farajtabar, Y. Wang, M. Gomez-Rodriguez, S. Li, H. Zha, and L. Song, COEVOLVE: a joint point process model for information diffusion and network diff --git a/src/aggregators/coevolve.jl b/src/aggregators/coevolve.jl index eb5a0e75..795de02e 100644 --- a/src/aggregators/coevolve.jl +++ b/src/aggregators/coevolve.jl @@ -20,12 +20,14 @@ mutable struct CoevolveJumpAggregation{T, S, F1, F2, RNG, GR, PQ} <: urates::F1 # vector of rate upper bound functions rateintervals::F1 # vector of interval length functions haslratevec::Vector{Bool} # vector of whether an lrate was provided for this vrj + cur_lrates::Vector{T} # the last computed lower rate for each rate end function CoevolveJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::Nothing, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, rng::RNG; u::U, dep_graph = nothing, lrates, urates, - rateintervals, haslratevec) where {T, S, F1, F2, RNG, U} + rateintervals, haslratevec, + cur_lrates::Vector{T}) where {T, S, F1, F2, RNG, U} if dep_graph === nothing if (get_num_majumps(maj) == 0) || !isempty(urates) error("To use Coevolve a dependency graph between jumps must be supplied.") @@ -48,8 +50,49 @@ function CoevolveJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::Not pq = MutableBinaryMinHeap{T}() affecttype = F2 <: Tuple ? F2 : Any CoevolveJumpAggregation{T, S, F1, affecttype, RNG, typeof(dg), - typeof(pq)}(nj, nj, njt, et, crs, sr, maj, rs, affs!, sps, rng, - dg, pq, lrates, urates, rateintervals, haslratevec) + typeof(pq)}(nj, nj, njt, et, crs, sr, maj, + rs, affs!, sps, rng, dg, pq, + lrates, urates, rateintervals, + haslratevec, cur_lrates) +end + +# display +function num_constant_rate_jumps(aggregator::CoevolveJumpAggregation) + length(aggregator.urates) +end + +# condition for jump to occur +function (p::CoevolveJumpAggregation)(u, t, integrator) + p.next_jump_time == t && + accept_next_jump!(p, integrator, integrator.u, integrator.p, integrator.t) +end + +# executing jump at the next jump time +function (p::CoevolveJumpAggregation)(integrator::I) where {I <: + AbstractSSAIntegrator} + if !accept_next_jump!(p, integrator, integrator.u, integrator.p, integrator.t) + return nothing + end + affects! = p.affects! + if affects! isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}} + execute_jumps!(p, integrator, integrator.u, integrator.p, integrator.t, affects!) + else + error("Error, invalid affects! type. Expected a vector of function wrappers and got $(typeof(affects!))") + end + generate_jumps!(p, integrator, integrator.u, integrator.p, integrator.t) + register_next_jump_time!(integrator, p, integrator.t) + nothing +end + +function (p::CoevolveJumpAggregation{T, S, F1, F2})(integrator::AbstractSSAIntegrator) where + {T, S, F1, F2 <: Union{Tuple, Nothing}} + if !accept_next_jump!(p, integrator, integrator.u, integrator.p, integrator.t) + return nothing + end + execute_jumps!(p, integrator, integrator.u, integrator.p, integrator.t, p.affects!) + generate_jumps!(p, integrator, integrator.u, integrator.p, integrator.t) + register_next_jump_time!(integrator, p, integrator.t) + nothing end # creating the JumpAggregation structure (tuple-based variable jumps) @@ -92,12 +135,14 @@ function aggregate(aggregator::Coevolve, u, p, t, end_time, constant_jumps, num_jumps = get_num_majumps(ma_jumps) + nrjs cur_rates = Vector{typeof(t)}(undef, num_jumps) + cur_lrates = zeros(typeof(t), nvrjs) sum_rate = nothing next_jump = 0 next_jump_time = typemax(t) CoevolveJumpAggregation(next_jump, next_jump_time, end_time, cur_rates, sum_rate, ma_jumps, rates, affects!, save_positions, rng; - u, dep_graph, lrates, urates, rateintervals, haslratevec) + u, dep_graph, lrates, urates, rateintervals, haslratevec, + cur_lrates) end # set up a new simulation and calculate the first jump / jump time @@ -109,12 +154,12 @@ function initialize!(p::CoevolveJumpAggregation, integrator, u, params, t) end # execute one jump, changing the system state -function execute_jumps!(p::CoevolveJumpAggregation, integrator, u, params, t, affects!) +function execute_jumps!(p::CoevolveJumpAggregation, integrator, u, params, t, + affects!) # execute jump - u = update_state!(p, integrator, u, affects!) - + update_state!(p, integrator, u, affects!) # update current jump rates and times - update_dependent_rates!(p, u, params, t) + update_dependent_rates!(p, integrator.u, integrator.p, t) nothing end @@ -125,13 +170,67 @@ function generate_jumps!(p::CoevolveJumpAggregation, integrator, u, params, t) end ######################## SSA specific helper routines ######################## +function accept_next_jump!(p::CoevolveJumpAggregation, integrator, u, params, t) + @unpack next_jump, ma_jumps = p + + num_majumps = get_num_majumps(ma_jumps) + + (next_jump <= num_majumps) && return true + + @unpack cur_rates, rates, rng, urates, cur_lrates = p + num_cjumps = length(urates) - length(rates) + uidx = next_jump - num_majumps + lidx = uidx - num_cjumps + + # lidx <= 0 indicates that we have a constant jump which we always accept + (lidx <= 0) && return true + + @inbounds urate = cur_rates[next_jump] + @inbounds lrate = cur_lrates[lidx] + + s = -one(t) + + # the first condition applies when the candidate time t is rejected because + # it was longer than the rateinterval when it was proposed + if lrate == typemax(t) + urate = get_urate(p, uidx, u, params, t) + s = urate == zero(t) ? typemax(t) : randexp(rng) / urate + elseif lrate < urate + # when the lower and upper bound are the same, then v < 1 = lrate / urate = urate / urate + v = rand(rng) * urate + # first inequality is less expensive and short-circuits the evaluation + if (v > lrate) + rate = get_rate(p, lidx, u, params, t) + if v > rate + urate = get_urate(p, uidx, u, params, t) + s = urate == zero(t) ? typemax(t) : randexp(rng) / urate + end + end + end + + # candidate t was accepted + (s < 0) && return true + + # candidate t was reject, compute a new candidate time + t = next_candidate_time!(p, u, params, t, s, lidx) + update!(p.pq, next_jump, t) + @inbounds cur_rates[next_jump] = urate + + p.prev_jump = next_jump + generate_jumps!(p, integrator, integrator.u, integrator.p, integrator.t) + register_next_jump_time!(integrator, p, integrator.t) + u_modified!(integrator, false) + + return false +end + function update_dependent_rates!(p::CoevolveJumpAggregation, u, params, t) @inbounds deps = p.dep_gr[p.next_jump] - @unpack cur_rates, end_time, pq = p + @unpack cur_rates, pq = p for (ix, i) in enumerate(deps) - ti, last_urate_i = next_time(p, u, params, t, i, end_time) + ti, urate_i = next_time(p, u, params, t, i) update!(pq, i, ti) - @inbounds cur_rates[i] = last_urate_i + @inbounds cur_rates[i] = urate_i end nothing end @@ -156,61 +255,56 @@ end @inbounds return p.rates[lidx](u, params, t) end -function next_time(p::CoevolveJumpAggregation{T}, u, params, t, i, tstop::T) where {T} - @unpack rng, haslratevec = p - num_majumps = get_num_majumps(p.ma_jumps) - num_cjumps = length(p.urates) - length(p.rates) +function next_time(p::CoevolveJumpAggregation, u, params, t, i) + @unpack next_jump, cur_rates, ma_jumps, rates, rng, pq, urates = p + num_majumps = get_num_majumps(ma_jumps) + num_cjumps = length(urates) - length(rates) uidx = i - num_majumps lidx = uidx - num_cjumps urate = uidx > 0 ? get_urate(p, uidx, u, params, t) : get_ma_urate(p, i, u, params, t) - last_urate = p.cur_rates[i] - if i != p.next_jump && last_urate > zero(t) - s = urate == zero(t) ? typemax(t) : last_urate / urate * (p.pq[i] - t) - else - s = urate == zero(t) ? typemax(t) : randexp(rng) / urate - end - _t = t + s - if lidx > 0 - while t < tstop - rateinterval = get_rateinterval(p, lidx, u, params, t) - if s > rateinterval - t = t + rateinterval - urate = get_urate(p, uidx, u, params, t) - s = urate == zero(t) ? typemax(t) : randexp(rng) / urate - _t = t + s - continue - end - (_t >= tstop) && break - - lrate = haslratevec[lidx] ? get_lrate(p, lidx, u, params, t) : zero(t) - if lrate < urate - # when the lower and upper bound are the same, then v < 1 = lrate / urate = urate / urate - v = rand(rng) * urate - # first inequality is less expensive and short-circuits the evaluation - if (v > lrate) && (v > get_rate(p, lidx, u, params, _t)) - t = _t - urate = get_urate(p, uidx, u, params, t) - s = urate == zero(t) ? typemax(t) : randexp(rng) / urate - _t = t + s - continue - end - elseif lrate > urate - error("The lower bound should be lower than the upper bound rate for t = $(t) and i = $(i), but lower bound = $(lrate) > upper bound = $(urate)") - end - break + # we can only re-use the rng in the case of constant rates because the rng + # used to compute the next candidate time has not been accepted or rejected + if i != next_jump && lidx <= 0 + last_urate = cur_rates[i] + if last_urate > zero(t) + s = urate == zero(t) ? typemax(t) : last_urate / urate * (pq[i] - t) + return next_candidate_time!(p, u, params, t, s, lidx), urate end end - return _t, urate + s = urate == zero(t) ? typemax(t) : randexp(rng) / urate + return next_candidate_time!(p, u, params, t, s, lidx), urate +end + +function next_candidate_time!(p::CoevolveJumpAggregation, u, params, t, s, lidx) + if lidx <= 0 + return t + s + end + @unpack end_time, haslratevec, cur_lrates = p + rateinterval = get_rateinterval(p, lidx, u, params, t) + if s > rateinterval + t = t + rateinterval + # we set the lrate to typemax(t) to indicate rejection due to candidate being larger than rateinterval + @inbounds cur_lrates[lidx] = typemax(t) + return t + end + t = t + s + if t < end_time + lrate = haslratevec[lidx] ? get_lrate(p, lidx, u, params, t) : zero(t) + @inbounds cur_lrates[lidx] = lrate + else + # no need to compute the lower bound when time is past the end time + @inbounds cur_lrates[lidx] = typemax(t) + end + return t end -# reevaulate all rates, recalculate all jump times, and reinit the priority queue +# re-evaluates all rates, recalculate all jump times, and reinit the priority queue function fill_rates_and_get_times!(p::CoevolveJumpAggregation, u, params, t) - @unpack end_time = p num_jumps = get_num_majumps(p.ma_jumps) + length(p.urates) p.cur_rates = zeros(typeof(t), num_jumps) jump_times = Vector{typeof(t)}(undef, num_jumps) @inbounds for i in 1:num_jumps - jump_times[i], p.cur_rates[i] = next_time(p, u, params, t, i, end_time) + jump_times[i], p.cur_rates[i] = next_time(p, u, params, t, i) end p.pq = MutableBinaryMinHeap(jump_times) nothing diff --git a/src/aggregators/ssajump.jl b/src/aggregators/ssajump.jl index a5904e86..07b4b031 100644 --- a/src/aggregators/ssajump.jl +++ b/src/aggregators/ssajump.jl @@ -17,9 +17,7 @@ An aggregator interface for SSA-like algorithms. ### Optional fields: - - `dep_gr` # dependency graph, dep_gr[i] = indices of reactions that should - - # be updated when rx i occurs. + - `dep_gr` # dependency graph, dep_gr[i] = indices of reactions that should be updated when rx i occurs. """ abstract type AbstractSSAJumpAggregator{T, S, F1, F2, RNG} <: AbstractJumpAggregator end @@ -36,7 +34,7 @@ end # execute_jumps! # generate_jumps! -@inline function makewrapper(::Type{T}, aff) where T +@inline function makewrapper(::Type{T}, aff) where {T} # rewrap existing wrappers if aff isa FunctionWrappers.FunctionWrapper T(aff.obj[]) @@ -50,7 +48,7 @@ end @inline function concretize_affects!(p::AbstractSSAJumpAggregator, ::I) where {I <: DiffEqBase.DEIntegrator} if (p.affects! isa Vector) && - !(p.affects! isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}}) + !(p.affects! isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}}) AffectWrapper = FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}} p.affects! = AffectWrapper[makewrapper(AffectWrapper, aff) for aff in p.affects!] end diff --git a/src/solve.jl b/src/solve.jl index a880a408..abaf0f05 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -22,12 +22,14 @@ function DiffEqBase.__init(_jump_prob::DiffEqBase.AbstractJumpProblem{P}, # DDEProblems do not have a recompile_flag argument if jump_prob.prob isa DiffEqBase.AbstractDDEProblem + # callback comes after jump consistent with SSAStepper integrator = init(jump_prob.prob, alg, timeseries, ts, ks; - callback = CallbackSet(callback, jump_prob.jump_callback), + callback = CallbackSet(jump_prob.jump_callback, callback), kwargs...) else + # callback comes after jump consistent with SSAStepper integrator = init(jump_prob.prob, alg, timeseries, ts, ks, recompile; - callback = CallbackSet(callback, jump_prob.jump_callback), + callback = CallbackSet(jump_prob.jump_callback, callback), kwargs...) end end diff --git a/test/functionwrappers.jl b/test/functionwrappers.jl index 369ebf33..8ea9aaf3 100644 --- a/test/functionwrappers.jl +++ b/test/functionwrappers.jl @@ -3,15 +3,16 @@ rng = StableRNG(12345) # https://github.com/SciML/JumpProcesses.jl/issues/324 let - rate(u, p, t; debug=true) = 5.0 + rate(u, p, t; debug = true) = 5.0 function affect!(integrator) integrator.u[1] += 1 nothing end - jump = VariableRateJump(rate, affect!; urate=(u,p,t)->10.0, rateinterval=(u,p,t)->0.1) + jump = VariableRateJump(rate, affect!; urate = (u, p, t) -> 10.0, + rateinterval = (u, p, t) -> 0.1) prob = DiscreteProblem([0.0], (0.0, 2.0), [1.0]) - jprob = JumpProblem(prob, Coevolve(), jump; dep_graph=[[1]], rng) + jprob = JumpProblem(prob, Coevolve(), jump; dep_graph = [[1]], rng) agg = jprob.discrete_jump_aggregation @test typeof(agg.affects!) <: Vector{Any} diff --git a/test/hawkes_test.jl b/test/hawkes_test.jl index 0de428e3..46b504e6 100644 --- a/test/hawkes_test.jl +++ b/test/hawkes_test.jl @@ -60,7 +60,8 @@ function hawkes_jump(u, g, h; uselrate = true) return [hawkes_jump(i, g, h; uselrate) for i in 1:length(u)] end -function hawkes_problem(p, agg::Coevolve; u = [0.0], tspan = (0.0, 50.0), +function hawkes_problem(p, agg::Coevolve; u = [0.0], + tspan = (0.0, 50.0), save_positions = (false, true), g = [[1]], h = [[]], uselrate = true) dprob = DiscreteProblem(u, tspan, p) @@ -133,10 +134,10 @@ for (i, alg) in enumerate(algs) end # test stepping Coevolve with continuous integrator and bounded jumps -let +let alg = Coevolve() oprob = ODEProblem(f!, u0, tspan, p) jumps = hawkes_jump(u0, g, h) - jprob = JumpProblem(oprob, Coevolve(), jumps...; dep_graph = g, rng) + jprob = JumpProblem(oprob, alg, jumps...; dep_graph = g, rng) @test ((jprob.variable_jumps === nothing) || isempty(jprob.variable_jumps)) sols = Vector{ODESolution}(undef, Nsims) for n in 1:Nsims @@ -149,10 +150,10 @@ let end # test disabling bounded jumps and using continuous integrator -let +let alg = Coevolve() oprob = ODEProblem(f!, u0, tspan, p) jumps = hawkes_jump(u0, g, h) - jprob = JumpProblem(oprob, Coevolve(), jumps...; dep_graph = g, rng, + jprob = JumpProblem(oprob, alg, jumps...; dep_graph = g, rng, use_vrj_bounds = false) @test length(jprob.variable_jumps) == 1 sols = Vector{ODESolution}(undef, Nsims) diff --git a/test/runtests.jl b/test/runtests.jl index c78d86a1..3b928ed2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -1,3 +1,4 @@ +#! format: off using JumpProcesses, DiffEqBase, SafeTestsets @time begin @@ -20,6 +21,7 @@ using JumpProcesses, DiffEqBase, SafeTestsets @time @safetestset "Composition-Rejection Table Tests" begin include("table_test.jl") end @time @safetestset "Extinction test" begin include("extinction_test.jl") end @time @safetestset "Saveat Regression test" begin include("saveat_regression.jl") end + @time @safetestset "Save_positions test" begin include("save_positions.jl") end @time @safetestset "Ensemble Uniqueness test" begin include("ensemble_uniqueness.jl") end @time @safetestset "Thread Safety test" begin include("thread_safety.jl") end @time @safetestset "A + B <--> C" begin include("reversible_binding.jl") end diff --git a/test/save_positions.jl b/test/save_positions.jl new file mode 100644 index 00000000..cc08394a --- /dev/null +++ b/test/save_positions.jl @@ -0,0 +1,28 @@ +using JumpProcesses, OrdinaryDiffEq, Test +using StableRNGs +rng = StableRNG(12345) + +# test that we only save when a jump occurs +for alg in (Coevolve(),) + u0 = [0] + tspan = (0.0, 30.0) + + dprob = DiscreteProblem(u0, tspan) + # set the rate to 0, so that no jump ever occurs; but urate is positive so + # Coevolve will consider many candidates before the end of the simmulation. + # None of these points should be saved. + jump = VariableRateJump((u, p, t) -> 0, (integrator) -> integrator.u[1] += 1; + urate = (u, p, t) -> 1.0, rateinterval = (u, p, t) -> 5.0) + jumpproblem = JumpProblem(dprob, alg, jump; dep_graph = [[1]], + save_positions = (false, true)) + sol = solve(jumpproblem, SSAStepper()) + @test sol.t == [0.0, 30.0] + + oprob = ODEProblem((du, u, p, t) -> 0, u0, tspan) + jump = VariableRateJump((u, p, t) -> 0, (integrator) -> integrator.u[1] += 1; + urate = (u, p, t) -> 1.0, rateinterval = (u, p, t) -> 5.0) + jumpproblem = JumpProblem(oprob, alg, jump; dep_graph = [[1]], + save_positions = (false, true)) + sol = solve(jumpproblem, Tsit5(); save_everystep = false) + @test sol.t == [0.0, 30.0] +end diff --git a/test/spatial/run_spatial_tests.jl b/test/spatial/run_spatial_tests.jl index 75f43161..d143e3a3 100644 --- a/test/spatial/run_spatial_tests.jl +++ b/test/spatial/run_spatial_tests.jl @@ -1,3 +1,5 @@ +#! format: off + using JumpProcesses, DiffEqBase, SafeTestsets @time begin diff --git a/test/variable_rate.jl b/test/variable_rate.jl index 4080c3ac..bf854e67 100644 --- a/test/variable_rate.jl +++ b/test/variable_rate.jl @@ -179,15 +179,47 @@ let constant_rate_jump = ConstantRateJump(cs_rate1, affect!) jumpset_ = JumpSet((), (constant_rate_jump,), nothing, mass_action_jump_) - u0 = [0] - tspan = (0.0, 30.0) - dprob_ = DiscreteProblem(u0, tspan) - alg = Coevolve() - @test_throws ErrorException JumpProblem(dprob_, alg, jumpset_, - save_positions = (false, false)) - - vrj = VariableRateJump(cs_rate1, affect!; urate = ((u, p, t) -> 1.0), - rateinterval = ((u, p, t) -> 1.0)) - @test_throws ErrorException JumpProblem(dprob_, alg, mass_action_jump_, vrj; - save_positions = (false, false)) + for alg in (Coevolve(),) + u0 = [0] + tspan = (0.0, 30.0) + dprob_ = DiscreteProblem(u0, tspan) + @test_throws ErrorException JumpProblem(dprob_, alg, jumpset_, + save_positions = (false, false)) + + vrj = VariableRateJump(cs_rate1, affect!; urate = ((u, p, t) -> 1.0), + rateinterval = ((u, p, t) -> 1.0)) + @test_throws ErrorException JumpProblem(dprob_, alg, mass_action_jump_, vrj; + save_positions = (false, false)) + end +end + +# Test that rate, urate and lrate do not get called past tstop +# https://github.com/SciML/JumpProcesses.jl/issues/330 +let + test_rate(u, p, t) = 0.1 + test_affect!(integrator) = (integrator.u[1] += 1) + function test_lrate(u, p, t) + if t > 1.0 + error("test_urate does not handle t > 1.0") + else + return 0.05 + end + end + function test_urate(u, p, t) + if t > 1.0 + error("test_urate does not handle t > 1.0") + else + return 0.2 + end + end + + test_jump = VariableRateJump(test_rate, test_affect!; urate = test_urate, + rateinterval = (u, p, t) -> 1.0) + + dprob = DiscreteProblem([0], (0.0, 1.0), nothing) + jprob = JumpProblem(dprob, Coevolve(), test_jump; dep_graph = [[1]]) + + @test_nowarn for i in 1:50 + solve(jprob, SSAStepper()) + end end