Merge pull request #30 from isaacsas/massaction-rate-updates

Expanding MassActionJump input options

src/jumps.jl

@@ -41,13 +41,18 @@ struct MassActionJump{T,S} <: AbstractJump
   reactant_stoch::S
   net_stoch::S
 
-  function MassActionJump{T,S}(rates::T, rs::S, ns::S, scale_rates::Bool) where {T,S}
+  function MassActionJump{T,S}(rates::T, rs::S, ns::S, scale_rates::Bool) where {T <: AbstractVector,S}
     sr = copy(rates)
     if scale_rates && !isempty(sr) 
       scalerates!(sr, rs)
     end
     new(sr, rs, ns)
   end
+  function MassActionJump{T,S}(rate::T, rs::S, ns::S, scale_rates::Bool) where {T <: Number,S}
+    sr = scale_rates ? scalerate(rate, rs) : rate
+    new(sr, rs, ns)
+  end
+
 end
 MassActionJump(usr::T, rs::S, ns::S; scale_rates = true ) where {T,S} = MassActionJump{T,S}(usr, rs, ns, scale_rates)
 
@@ -67,9 +72,22 @@ JumpSet() = JumpSet((),(),nothing,nothing)
 JumpSet(jb::Void) = JumpSet()
 
 # For Varargs, use recursion to make it type-stable
-
 JumpSet(jumps::AbstractJump...) = JumpSet(split_jumps((), (), nothing, nothing, jumps...)...)
 
+# handle vector of mass action jumps 
+function JumpSet(vjs, cjs, rj, majv::Vector{T}) where {T <: MassActionJump}
+  if isempty(majv)
+    error("JumpSets do not accept empty mass action jump collections; use \"nothing\" instead.")
+  end
+
+  maj = setup_majump_to_merge(majv[1].scaled_rates, majv[1].reactant_stoch, majv[1].net_stoch)
+  for i = 2:length(majv)
+    massaction_jump_combine(maj, majv[i])
+  end
+
+  JumpSet(vjs, cjs, rj, maj)
+end
+
 @inline split_jumps(vj, cj, rj, maj) = vj, cj, rj, maj
 @inline split_jumps(vj, cj, rj, maj, v::VariableRateJump, args...) = split_jumps((vj..., v), cj, rj, maj, args...)
 @inline split_jumps(vj, cj, rj, maj, c::ConstantRateJump, args...) = split_jumps(vj, (cj..., c), rj, maj, args...)
@@ -85,10 +103,45 @@ regular_jump_combine(rj1::Void,rj2::RegularJump) = rj2
 regular_jump_combine(rj1::Void,rj2::Void) = rj1
 regular_jump_combine(rj1::RegularJump,rj2::RegularJump) = error("Only one regular jump is allowed in a JumpSet")
 
+
+# functionality to merge two mass action jumps together
+
+# if given containers of rates and stoichiometry directly create a jump
+function setup_majump_to_merge(sr::T, rs::AbstractVector{S}, ns::AbstractVector{S}) where {T <: AbstractVector, S <: AbstractArray}
+  MassActionJump(sr, rs, ns)
+end
+
+# if just given the data for one jump (and not in a container) wrap in a vector
+function setup_majump_to_merge(sr::T, rs::S, ns::S) where {T <: Number, S <: AbstractArray}
+  MassActionJump([sr], [rs], [ns])
+end
+
+# when given a collection of reactions to add to maj
+function majump_merge!(maj::MassActionJump{U,V}, sr::U, rs::V, ns::V) where {U <: AbstractVector, V <: AbstractVector}
+  append!(maj.scaled_rates, sr)
+  append!(maj.reactant_stoch, rs)
+  append!(maj.net_stoch, ns)
+  maj
+end
+
+# when given a single jump's worth of data to add to maj
+function majump_merge!(maj::MassActionJump{U,V}, sr::T, rs::S, ns::S) where {U <: AbstractVector, V <: AbstractVector, T <: Number, S <: AbstractArray}
+  push!(maj.scaled_rates, sr)
+  push!(maj.reactant_stoch, rs)
+  push!(maj.net_stoch, ns)
+  maj
+end
+
+# when maj only stores a single jump's worth of data (and not in a collection)
+# create a new jump with the merged data stored in vectors
+function majump_merge!(maj::MassActionJump{T,S}, sr::T, rs::S, ns::S) where {T <: Number, S <: AbstractArray}
+  MassActionJump([maj.scaled_rates, sr], [maj.reactant_stoch, rs], [maj.net_stoch, ns])
+end
+
 massaction_jump_combine(maj1::MassActionJump, maj2::Void) = maj1
 massaction_jump_combine(maj1::Void, maj2::MassActionJump) = maj2
 massaction_jump_combine(maj1::Void, maj2::Void) = maj1
-massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) = error("Only one mass action jump type is allowed in a JumpSet")
+massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) = majump_merge!(maj1, maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch)
 
 
 ##### helper methods for unpacking rates and affects! from constant jumps #####

src/massaction_rates.jl

@@ -1,40 +1,48 @@
 ###############################################################################
-# Stochiometry for a given reaction is a vector of pairs mapping species id to 
+# Stochiometry for a given reaction is a vector of pairs mapping species id to
 # stochiometric coefficient.
 ###############################################################################
 
-@inbounds @fastmath function evalrxrate(speciesvec::AbstractVector{T}, rateconst,
-                                        stochmat::AbstractVector{Pair{S,V}})::typeof(rateconst) where {T,S,V}
+@fastmath function evalrxrate(speciesvec::AbstractVector{T}, rateconst,
+                              stochmat::AbstractVector{Pair{S,V}})::typeof(rateconst) where {T,S,V}
     val = one(T)
 
-    for specstoch in stochmat
+    @inbounds for specstoch in stochmat
         specpop = speciesvec[specstoch[1]]
         val    *= specpop
-        for k = 2:specstoch[2]
+        @inbounds for k = 2:specstoch[2]
             specpop -= one(specpop)
             val     *= specpop
         end
     end
 
-    return rateconst * val
+     rateconst * val
 end
 
-@inline @inbounds @fastmath function executerx!(speciesvec::AbstractVector{T},
-                                                net_stoch::AbstractVector{Pair{S,V}}) where {T,S,V}
-    for specstoch in net_stoch
+@inline @fastmath function executerx!(speciesvec::AbstractVector{T},
+                                      net_stoch::AbstractVector{Pair{S,V}}) where {T,S,V}
+    @inbounds for specstoch in net_stoch
         speciesvec[specstoch[1]] += specstoch[2]
     end
     nothing
 end
 
 
-@inbounds function scalerates!(unscaled_rates, stochmat::Vector{Vector{Pair{S,T}}}) where {S,T}
-    for i in eachindex(unscaled_rates)
+function scalerates!(unscaled_rates::Vector{U}, stochmat::Vector{Vector{Pair{S,T}}}) where {U,S,T}
+    @inbounds for i in eachindex(unscaled_rates)
         coef = one(T)
-        for specstoch in stochmat[i]            
+        @inbounds for specstoch in stochmat[i]
             coef *= factorial(specstoch[2])
         end
         unscaled_rates[i] /= coef
     end
     nothing
 end
+
+function scalerate(unscaled_rate::U, stochmat::Vector{Pair{S,T}}) where {U <: Number, S, T}
+    coef = one(T)
+    @inbounds for specstoch in stochmat
+        coef *= factorial(specstoch[2])
+    end
+    unscaled_rate /= coef
+end

test/bimolerx_test.jl

@@ -11,7 +11,7 @@ dotestmean   = true
 doprintmeans = false
 
 # SSAs to test
-SSAalgs = (Direct(),FRM()) #,DirectFW(), FRM(), FRMFW())
+SSAalgs = (Direct(),DirectFW(), FRM(), FRMFW())
 
 Nsims        = 32000
 tf           = .01

test/linearreaction_test.jl

@@ -1,4 +1,5 @@
 # calculates the mean from N stochastic A->B reactions at different rates
+# this really tests different ways of constructing the jump problems
 using DiffEqBase, DiffEqJump
 using Base.Test
 
@@ -13,7 +14,7 @@ tf          = .1
 baserate    = .1
 A0          = 100
 exactmean   = (t,ratevec) -> A0 * exp(-sum(ratevec) * t)
-SSAalgs     = [FRM(), FRMFW(), Direct(), DirectFW()]
+SSAalgs     = [Direct()]#, DirectFW(), FRM(), FRMFW()]
 
 rates = ones(Float64, Nrxs) * baserate;
 cumsum!(rates, rates)    
@@ -72,7 +73,7 @@ function A_to_B_vec(N, method)
     jump_prob
 end
 
-# uses a mass action jump to represent all reactions
+# uses a single mass action jump to represent all reactions
 function A_to_B_ma(N, method)
     reactstoch = Vector{Vector{Pair{Int,Int}}}();
     netstoch   = Vector{Vector{Pair{Int,Int}}}();
@@ -89,8 +90,8 @@ function A_to_B_ma(N, method)
     jump_prob
 end
 
-# uses a mass action jump to represent half the reactions and a vector of constant jumps for the other half
-# stores them in a JumpSet
+# uses one mass action jump to represent half the reactions and a vector 
+# of constant jumps for the other half. Stores them in a JumpSet
 function A_to_B_hybrid(N, method)
     # half reactions are treated as mass action and half as constant jumps
     switchidx = (N//2).num
@@ -122,8 +123,8 @@ function A_to_B_hybrid(N, method)
     jump_prob
 end
 
-# uses a mass action jump to represent half the reactions and a vector of constant jumps for the other half
-# passes them to JumpProblem as a splatted tuple
+# uses a mass action jump to represent half the reactions and a vector
+# of constant jumps for the other half. Passes them to JumpProblem as a splatted tuple
 function A_to_B_hybrid_nojset(N, method)
     # half reactions are treated as mass action and half as constant jumps
     switchidx = (N//2).num
@@ -155,8 +156,131 @@ function A_to_B_hybrid_nojset(N, method)
     jump_prob
 end
 
-#jump_prob_gens = [A_to_B_tuple, A_to_B_vec, A_to_B_ma, A_to_B_hybrid, A_to_B_hybrid_nojset]
-jump_prob_gens = [A_to_B_tuple, A_to_B_ma, A_to_B_hybrid]
+
+# uses a vector of mass action jumps of vectors to represent half the reactions and a vector 
+# of constant jumps for the other half. Passes them to JumpProblem as a JumpSet
+function A_to_B_hybrid_vecs(N, method)
+    # half reactions are treated as mass action and half as constant jumps
+    switchidx = (N//2).num
+
+    # mass action reactions
+    majumps = Vector{MassActionJump}()
+    for i in 1:switchidx
+        push!(majumps, MassActionJump([rates[i]], [[1 => 1]], [[1 => -1, 2=>1]] ))
+    end
+
+     # jump reactions
+     jumpvec = Vector{ConstantRateJump}()
+     for i in (switchidx+1):N
+         ratefunc = (u,p,t) -> rates[i] * u[1]
+         affect!  = function (integrator)
+             integrator.u[1] -= 1
+             integrator.u[2] += 1
+         end
+         push!(jumpvec, ConstantRateJump(ratefunc, affect!))
+     end
+    jset      = JumpSet((), jumpvec, nothing, majumps)
+    prob      = DiscreteProblem([A0,0], (0.0,tf))
+    jump_prob = JumpProblem(prob, method, jset; save_positions=(false,false))
+
+    jump_prob
+end
+
+# uses a vector of scalar mass action jumps to represent half the reactions and a vector 
+# of constant jumps for the other half. Passes them to JumpProblem as a JumpSet
+function A_to_B_hybrid_vecs_scalars(N, method)
+    # half reactions are treated as mass action and half as constant jumps
+    switchidx = (N//2).num
+
+    # mass action reactions
+    majumps = Vector{MassActionJump}()
+    for i in 1:switchidx
+        push!(majumps, MassActionJump(rates[i], [1 => 1], [1 => -1, 2=>1] ))
+    end
+
+     # jump reactions
+     jumpvec = Vector{ConstantRateJump}()
+     for i in (switchidx+1):N
+         ratefunc = (u,p,t) -> rates[i] * u[1]
+         affect!  = function (integrator)
+             integrator.u[1] -= 1
+             integrator.u[2] += 1
+         end
+         push!(jumpvec, ConstantRateJump(ratefunc, affect!))
+     end
+    jset      = JumpSet((), jumpvec, nothing, majumps)
+    prob      = DiscreteProblem([A0,0], (0.0,tf))
+    jump_prob = JumpProblem(prob, method, jset; save_positions=(false,false))
+
+    jump_prob
+end
+
+
+# uses a vector of scalar mass action jumps to represent half the reactions and a vector 
+# of constant jumps for the other half. Passes them to JumpProblem as a single splatted tuple.
+function A_to_B_hybrid_tups_scalars(N, method)
+    # half reactions are treated as mass action and half as constant jumps
+    switchidx = (N//2).num
+
+    # mass action reactions
+    majumpsv = Vector{MassActionJump}()
+    for i in 1:switchidx
+        push!(majumpsv, MassActionJump(rates[i], [1 => 1], [1 => -1, 2=>1] ))
+    end
+
+    # jump reactions
+    jumpvec = Vector{ConstantRateJump}()
+    for i in (switchidx+1):N
+        ratefunc = (u,p,t) -> rates[i] * u[1]
+        affect!  = function (integrator)
+            integrator.u[1] -= 1
+            integrator.u[2] += 1
+        end
+        push!(jumpvec, ConstantRateJump(ratefunc, affect!))
+    end
+
+    jumps     = ((maj for maj in majumpsv)..., (jump for jump in jumpvec)...)
+    prob      = DiscreteProblem([A0,0], (0.0,tf))
+    jump_prob = JumpProblem(prob, method, jumps...; save_positions=(false,false))
+
+    jump_prob
+end
+
+
+# uses a mass action jump to represent half the reactions and a tuple
+# of constant jumps for the other half. Passes them to JumpProblem as a JumpSet.
+function A_to_B_hybrid_tups(N, method)
+    # half reactions are treated as mass action and half as constant jumps
+    switchidx = (N//2).num
+
+    # mass action reactions
+    majumps = Vector{MassActionJump}()
+    for i in 1:switchidx
+        push!(majumps, MassActionJump([rates[i]], [[1 => 1]], [[1 => -1, 2=>1]] ))
+    end
+
+     # jump reactions
+     jumpvec = Vector{ConstantRateJump}()
+     for i in (switchidx+1):N
+        ratefunc = (u,p,t) -> rates[i] * u[1]
+        affect!  = function (integrator)
+            integrator.u[1] -= 1
+            integrator.u[2] += 1
+        end
+        push!(jumpvec, ConstantRateJump(ratefunc, affect!))
+     end
+    jumps    = ((jump for jump in jumpvec)...)
+    jset      = JumpSet((), jumps, nothing, majumps)
+    prob      = DiscreteProblem([A0,0], (0.0,tf))
+    jump_prob = JumpProblem(prob, method, jset; save_positions=(false,false))
+
+    jump_prob
+end
+
+
+ jump_prob_gens = [A_to_B_tuple, A_to_B_vec, A_to_B_ma, A_to_B_hybrid, A_to_B_hybrid_nojset,
+                   A_to_B_hybrid_vecs, A_to_B_hybrid_vecs_scalars, A_to_B_hybrid_tups,A_to_B_hybrid_tups_scalars]
+#jump_prob_gens = [A_to_B_tuple, A_to_B_ma, A_to_B_hybrid, A_to_B_hybrid_vecs, A_to_B_hybrid_vecs_scalars,A_to_B_hybrid_tups_scalars]
 
 for method in SSAalgs
     for jump_prob_gen in jump_prob_gens