Coevolve aggregator for VariableRateJump

docs/make.jl

@@ -1,7 +1,9 @@
 using Documenter, JumpProcesses
 
-cp("./docs/Manifest.toml", "./docs/src/assets/Manifest.toml", force = true)
-cp("./docs/Project.toml", "./docs/src/assets/Project.toml", force = true)
+docpath = Base.source_dir()
+assetpath = joinpath(docpath, "src", "assets")
+cp(joinpath(docpath, "Manifest.toml"), joinpath(assetpath, "Manifest.toml"), force = true)
+cp(joinpath(docpath, "Project.toml"), joinpath(assetpath, "Project.toml"), force = true)
 
 include("pages.jl")
 

docs/src/index.md

@@ -79,30 +79,32 @@ versioninfo() # hide
 ```
 ```@example
 using Pkg # hide
-Pkg.status(;mode = PKGMODE_MANIFEST) # hide
+Pkg.status(; mode = PKGMODE_MANIFEST) # hide
 ```
 ```@raw html
 </details>
 ```
 ```@raw html
-You can also download the 
+You can also download the
 <a href="
 ```
 ```@eval
 using TOML
-version = TOML.parse(read("../../Project.toml",String))["version"]
-name = TOML.parse(read("../../Project.toml",String))["name"]
-link = "https://github.com/SciML/"*name*".jl/tree/gh-pages/v"*version*"/assets/Manifest.toml"
+projtoml = joinpath("..", "..", "Project.toml")
+version = TOML.parse(read(projtoml, String))["version"]
+name = TOML.parse(read(projtoml, String))["name"]
+link = "https://github.com/SciML/" * name * ".jl/tree/gh-pages/v" * version * "/assets/Manifest.toml"
 ```
 ```@raw html
 ">manifest</a> file and the
 <a href="
 ```
 ```@eval
 using TOML
-version = TOML.parse(read("../../Project.toml",String))["version"]
-name = TOML.parse(read("../../Project.toml",String))["name"]
-link = "https://github.com/SciML/"*name*".jl/tree/gh-pages/v"*version*"/assets/Project.toml"
+projtoml = joinpath("..", "..", "Project.toml")
+version = TOML.parse(read(projtoml, String))["version"]
+name = TOML.parse(read(projtoml, String))["name"]
+link = "https://github.com/SciML/" * name * ".jl/tree/gh-pages/v" * version * "/assets/Project.toml"
 ```
 ```@raw html
 ">project</a> file.

src/SSA_stepper.jl

@@ -1,12 +1,13 @@
 """
 $(TYPEDEF)
 
-Highly efficient integrator for pure jump problems that involve only
-`ConstantRateJump`s and/or `MassActionJump`s.
+Highly efficient integrator for pure jump problems that involve only `ConstantRateJump`s,
+`MassActionJump`s, and/or `VariableRateJump`s *with rate bounds*.
 
 ## Notes
-- Only works with `JumProblem`s defined from `DiscreteProblem`s.
-- Only works with collections of `ConstantRateJump`s, `VariableRateJump`s and `MassActionJump`s.
+- 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.
 
 ## Examples

src/aggregators/aggregators.jl

@@ -178,6 +178,10 @@ needs_vartojumps_map(aggregator::AbstractAggregatorAlgorithm) = false
 needs_vartojumps_map(aggregator::RSSA) = true
 needs_vartojumps_map(aggregator::RSSACR) = true
 
+# true if aggregator supports variable rates
+supports_variablerates(aggregator::AbstractAggregatorAlgorithm) = false
+supports_variablerates(aggregator::Coevolve) = true
+
 is_spatial(aggregator::AbstractAggregatorAlgorithm) = false
 is_spatial(aggregator::NSM) = true
 is_spatial(aggregator::DirectCRDirect) = true

src/aggregators/coevolve.jl

@@ -7,9 +7,9 @@ mutable struct CoevolveJumpAggregation{T, S, F1, F2, RNG, GR, PQ} <:
     prev_jump::Int                    # the previous jump that was executed
     next_jump_time::T                 # the time of the next jump
     end_time::T                       # the time to stop a simulation
-    cur_rates::Nothing                # not used
+    cur_rates::Vector{T}              # the last computed upper bound for each rate
     sum_rate::Nothing                 # not used
-    ma_jumps::S                       # not used
+    ma_jumps::S                       # MassActionJumps
     rates::F1                         # vector of rate functions
     affects!::F2                      # vector of affect functions for VariableRateJumps
     save_positions::Tuple{Bool, Bool} # tuple for whether to save the jumps before and/or after event
@@ -18,77 +18,76 @@ mutable struct CoevolveJumpAggregation{T, S, F1, F2, RNG, GR, PQ} <:
     pq::PQ                            # priority queue of next time
     lrates::F1                        # vector of rate lower bound functions
     urates::F1                        # vector of rate upper bound functions
-    Ls::F1                            # vector of interval length functions
+    rateintervals::F1                 # vector of interval length functions
 end
 
-function CoevolveJumpAggregation(nj::Int, njt::T, et::T, crs::Nothing, sr::Nothing,
+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, Ls) where {T, S, F1, F2, RNG, U}
+                                 rng::RNG; u::U, dep_graph = nothing, lrates, urates,
+                                 rateintervals) where {T, S, F1, F2, RNG, U}
     if dep_graph === nothing
         if (get_num_majumps(maj) == 0) || !isempty(rs)
-            error("To use VariableRateJumps with the Queue Method algorithm a dependency graph between jumps must be supplied.")
+            error("To use Coevolve a dependency graph between jumps must be supplied.")
         else
             dg = make_dependency_graph(length(u), maj)
         end
     else
         # using a Set to ensure that edges are not duplicate
-        dg = [Set{Int}(append!([], jumps, [var]))
-              for (var, jumps) in enumerate(dep_graph)]
-        dg = [sort!(collect(i)) for i in dg]
+        dgsets = [Set{Int}(append!(Int[], jumps, [var]))
+                  for (var, jumps) in enumerate(dep_graph)]
+        dg = [sort!(collect(i)) for i in dgsets]
     end
 
-    num_jumps = get_num_majumps(maj) + length(rs)
+    num_jumps = get_num_majumps(maj) + length(urates)
 
     if length(dg) != num_jumps
         error("Number of nodes in the dependency graph must be the same as the number of jumps.")
     end
 
     pq = MutableBinaryMinHeap{T}()
     CoevolveJumpAggregation{T, S, F1, F2, RNG, typeof(dg),
-                            typeof(pq)
-                            }(nj, nj, njt,
-                              et,
-                              crs, sr, maj,
-                              rs,
-                              affs!, sps,
-                              rng,
-                              dg, pq,
-                              lrates, urates, Ls)
+                            typeof(pq)}(nj, nj, njt, et, crs, sr, maj, rs, affs!, sps, rng,
+                                        dg, pq, lrates, urates, rateintervals)
 end
 
 # creating the JumpAggregation structure (tuple-based variable jumps)
-function aggregate(aggregator::Coevolve, u, p, t, end_time, variable_jumps,
-                   ma_jumps, save_positions, rng;
-                   dep_graph = nothing,
-                   kwargs...)
+function aggregate(aggregator::Coevolve, u, p, t, end_time, constant_jumps,
+                   ma_jumps, save_positions, rng; dep_graph = nothing,
+                   variable_jumps = nothing, kwargs...)
     AffectWrapper = FunctionWrappers.FunctionWrapper{Nothing, Tuple{Any}}
     RateWrapper = FunctionWrappers.FunctionWrapper{typeof(t),
                                                    Tuple{typeof(u), typeof(p), typeof(t)}}
+    affects! = Vector{AffectWrapper}()
+    rates = Vector{RateWrapper}()
+    lrates = Vector{RateWrapper}()
+    urates = Vector{RateWrapper}()
+    rateintervals = Vector{RateWrapper}()
+
+    if (constant_jumps !== nothing) && !isempty(constant_jumps)
+        append!(affects!,
+                [AffectWrapper((integrator) -> (j.affect!(integrator); nothing))
+                 for j in constant_jumps])
+        append!(urates, [RateWrapper(j.rate) for j in constant_jumps])
+    end
+
     if (variable_jumps !== nothing) && !isempty(variable_jumps)
-        affects! = [AffectWrapper((integrator) -> (c.affect!(integrator); nothing))
-                    for c in variable_jumps]
-        rates = [RateWrapper(c.rate) for c in variable_jumps]
-        lrates = Any[RateWrapper(c.lrate) for c in variable_jumps]
-        urates = Any[RateWrapper(c.urate) for c in variable_jumps]
-        Ls = Any[RateWrapper(c.L) for c in variable_jumps]
-    else
-        affects! = Vector{AffectWrapper}()
-        rates = Vector{RateWrapper}()
-        lrates = Vector{RateWrapper}()
-        urates = Vector{RateWrapper}()
-        Ls = Vector{RateWrapper}()
+        append!(affects!,
+                [AffectWrapper((integrator) -> (j.affect!(integrator); nothing))
+                 for j in variable_jumps])
+        append!(rates, [RateWrapper(j.rate) for j in variable_jumps])
+        append!(lrates, [RateWrapper(j.lrate) for j in variable_jumps])
+        append!(urates, [RateWrapper(j.urate) for j in variable_jumps])
+        append!(rateintervals, [RateWrapper(j.rateinterval) for j in variable_jumps])
     end
-    cur_rates = nothing
+
+    num_jumps = get_num_majumps(ma_jumps) + length(urates)
+    cur_rates = Vector{typeof(t)}(undef, num_jumps)
     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 = u,
-                            dep_graph = dep_graph,
-                            lrates = lrates, urates = urates, Ls = Ls)
+                            u, dep_graph, lrates, urates, rateintervals)
 end
 
 # set up a new simulation and calculate the first jump / jump time
@@ -102,7 +101,7 @@ end
 # execute one jump, changing the system state
 function execute_jumps!(p::CoevolveJumpAggregation, integrator, u, params, t)
     # execute jump
-    u = update_state!(p, integrator, u, params, t)
+    u = update_state!(p, integrator, u)
     # update current jump rates and times
     update_dependent_rates!(p, u, params, t)
     nothing
@@ -114,92 +113,98 @@ function generate_jumps!(p::CoevolveJumpAggregation, integrator, u, params, t)
     nothing
 end
 
-@inline function update_state!(p::CoevolveJumpAggregation, integrator, u, params, t)
-    @unpack ma_jumps, next_jump = p
-    num_majumps = get_num_majumps(ma_jumps)
-    if next_jump <= num_majumps
-        if u isa SVector
-            integrator.u = executerx(u, next_jump, ma_jumps)
-        else
-            @inbounds executerx!(u, next_jump, ma_jumps)
-        end
-    else
-        idx = next_jump - num_majumps
-        @inbounds p.affects![idx](integrator)
-    end
-    p.prev_jump = next_jump
-    return integrator.u
-end
-
 ######################## SSA specific helper routines ########################
 function update_dependent_rates!(p::CoevolveJumpAggregation, u, params, t)
     @inbounds deps = p.dep_gr[p.next_jump]
-    @unpack end_time, pq = p
+    @unpack cur_rates, end_time, pq = p
     for (ix, i) in enumerate(deps)
-        ti = next_time(p, u, params, t, i, end_time)
+        ti, last_urate_i = next_time(p, u, params, t, i, end_time)
         update!(pq, i, ti)
+        cur_rates[i] = last_urate_i
     end
     nothing
 end
 
-function get_rates(p::CoevolveJumpAggregation, i, u)
-    ma_jumps = p.ma_jumps
-    num_majumps = get_num_majumps(ma_jumps)
-    if i <= num_majumps
-        _rate = evalrxrate(u, i, ma_jumps)
-        rate(u, p, t) = _rate
-        lrate = rate
-        urate = rate
-        L(u, p, t) = typemax(t)
-    else
-        idx = i - num_majumps
-        rate, lrate, urate, L = p.rates[idx], p.lrates[idx], p.urates[idx], p.Ls[idx]
-    end
-    return rate, lrate, urate, L
+@inline function get_ma_urate(p::CoevolveJumpAggregation, i, u, params, t)
+    return evalrxrate(u, i, p.ma_jumps)
+end
+
+@inline function get_urate(p::CoevolveJumpAggregation, uidx, u, params, t)
+    @inbounds return p.urates[uidx](u, params, t)
+end
+
+@inline function get_rateinterval(p::CoevolveJumpAggregation, lidx, u, params, t)
+    @inbounds return p.rateintervals[lidx](u, params, t)
+end
+
+@inline function get_lrate(p::CoevolveJumpAggregation, lidx, u, params, t)
+    @inbounds return p.lrates[lidx](u, params, t)
+end
+
+@inline function get_rate(p::CoevolveJumpAggregation, lidx, u, params, t)
+    @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 = p
-    rate, lrate, urate, L = get_rates(p, i, u)
-    while t < tstop
-        _urate = urate(u, params, t)
-        _L = L(u, params, t)
-        s = _urate == zero(t) ? typemax(t) : randexp(rng) / _urate
-        _t = t + s
-        if s > _L
-            t = t + _L
-            continue
-        end
-        if _t >= tstop
-            break
-        end
-        _lrate = lrate(u, params, t)
-        if _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)")
-        elseif _lrate < _urate
-            # when the lower and upper bound are the same, then v < 1 = _lrate / _urate = _rate / _urate
-            v = rand(rng)
-            # first inequality is less expensive and short-circuits the evaluation
-            if (v > _lrate / _urate)
-                _rate = rate(u, params, _t)
-                if (v > _rate / _urate)
-                    t = _t
-                    continue
+    num_majumps = get_num_majumps(p.ma_jumps)
+    num_cjumps = length(p.urates) - length(p.rates)
+    uidx = i - num_majumps
+    lidx = i - num_majumps - 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
+            if _t >= tstop
+                break
+            end
+            lrate = p.urates[uidx] === p.lrates[lidx] ? urate :
+                    get_lrate(p, lidx, u, params, t)
+            if 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)")
+            elseif lrate < urate
+                # when the lower and upper bound are the same, then v < 1 = lrate / urate = urate / urate
+                v = rand(rng)
+                # first inequality is less expensive and short-circuits the evaluation
+                if (v > lrate / urate)
+                    if (v > get_rate(p, lidx, u, params, _t) / urate)
+                        t = _t
+                        urate = get_urate(p, uidx, u, params, t)
+                        s = urate == zero(t) ? typemax(t) : randexp(rng) / urate
+                        _t = t + s
+                        continue
+                    end
                 end
             end
+            break
         end
-        return _t
     end
-    return typemax(t)
+    return _t, urate
 end
 
 # reevaulate 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.rates)
-    jump_times = Vector{eltype(t)}(undef, num_jumps)
+    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] = next_time(p, u, params, t, i, end_time)
+        jump_times[i], p.cur_rates[i] = next_time(p, u, params, t, i, end_time)
     end
     p.pq = MutableBinaryMinHeap(jump_times)
     nothing