Merge pull request #444 from isaacsas/speed-up-reaction-jump-sys

Speed up JumpSystem generation from ReactionSystems

Author: ChrisRackauckas

src/systems/jumps/jumpsystem.jl

@@ -45,7 +45,6 @@ end
 
 function JumpSystem(eqs, iv, states, ps; systems = JumpSystem[],
                                           name = gensym(:JumpSystem))
-
     ap = ArrayPartition(MassActionJump[], ConstantRateJump[], VariableRateJump[])
     for eq in eqs
         if eq isa MassActionJump
@@ -62,6 +61,9 @@ function JumpSystem(eqs, iv, states, ps; systems = JumpSystem[],
     JumpSystem{typeof(ap)}(ap, convert(Variable,iv), convert.(Variable, states), convert.(Variable, ps), name, systems)
 end
 
+JumpSystem(eqs::ArrayPartition, iv, states, ps; systems = JumpSystem[], name = gensym(:JumpSystem)) =
+    JumpSystem{typeof(eqs)}(eqs, convert(Variable,iv), convert.(Variable, states), convert.(Variable, ps), name, systems)
+
 
 generate_rate_function(js, rate) = build_function(rate, states(js), parameters(js),
                                         independent_variable(js),
@@ -114,45 +116,64 @@ function assemble_crj_expr(js, crj, statetoid)
     end
 end
 
-function assemble_maj(js, maj::MassActionJump{U,Vector{Pair{V,W}},Vector{Pair{V2,W2}}},
-                      statetoid, subber, invttype) where {U,V,W,V2,W2}
-    sr = maj.scaled_rates
-    if sr isa Operation
-        pval = subber(sr).value
-    elseif sr isa Variable
-        pval = subber(sr()).value
+function numericrate(rate, subber)
+    if rate isa Operation
+        rval = subber(rate).value
+    elseif rate isa Variable
+        rval = subber(rate()).value
     else
-        pval = maj.scaled_rates
+        rval = rate
     end
+    rval
+end
 
+function numericrstoich(mtrs::Vector{Pair{V,W}}, statetoid) where {V,W}
     rs = Vector{Pair{Int,W}}()
-    for (spec,stoich) in maj.reactant_stoch
-        if iszero(spec)
+    for (spec,stoich) in mtrs
+        if !(spec isa Operation) && iszero(spec)
             push!(rs, 0 => stoich)
         else
             push!(rs, statetoid[convert(Variable,spec)] => stoich)
         end
     end
     sort!(rs)
+    rs
+end
 
-    ns = Vector{Pair{Int,W2}}()
-    for (spec,stoich) in maj.net_stoch
-        iszero(spec) && error("Net stoichiometry can not have a species labelled 0.")
+function numericnstoich(mtrs::Vector{Pair{V,W}}, statetoid) where {V,W}
+    ns = Vector{Pair{Int,W}}()
+    for (spec,stoich) in mtrs
+        !(spec isa Operation) && iszero(spec) && error("Net stoichiometry can not have a species labelled 0.")
         push!(ns, statetoid[convert(Variable,spec)] => stoich)
     end
     sort!(ns)
+end
 
-    maj = MassActionJump(convert(invttype, pval), rs, ns, scale_rates = false)
-    return maj
+# assemble a numeric MassActionJump from a MT MassActionJump representing one rx.
+function assemble_maj(maj::MassActionJump, statetoid, subber, invttype)
+    rval = numericrate(maj.scaled_rates, subber)
+    rs   = numericrstoich(maj.reactant_stoch, statetoid)
+    ns   = numericnstoich(maj.net_stoch, statetoid)
+    maj  = MassActionJump(convert(invttype, rval), rs, ns, scale_rates = false)
+    maj
 end
 
+# For MassActionJumps that contain many reactions
+# function assemble_maj(maj::MassActionJump{U,V,W}, statetoid, subber,
+#                       invttype) where {U <: AbstractVector,V,W}
+#     rval = [convert(invttype,numericrate(sr, subber)) for sr in maj.scaled_rates]
+#     rs   = [numericrstoich(rs, statetoid) for rs in maj.reactant_stoch]
+#     ns   = [numericnstoich(ns, statetoid) for ns in maj.net_stoch]
+#     maj  = MassActionJump(rval, rs, ns, scale_rates = false)
+#     maj
+# end
 """
 ```julia
 function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan,
                                     parammap=DiffEqBase.NullParameters; kwargs...)
 ```
 
-Generates a blank DiscreteProblem for a pure jump JumpSystem to utilize as 
+Generates a blank DiscreteProblem for a pure jump JumpSystem to utilize as
 its `prob.prob`. This is used in the case where there are no ODEs
 and no SDEs associated with the system.
 
@@ -167,9 +188,17 @@ dprob = DiscreteProblem(js, u₀map, tspan, parammap)
 """
 function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Tuple,
                                     parammap=DiffEqBase.NullParameters(); kwargs...)
-    u0 = varmap_to_vars(u0map, states(sys))
-    p  = varmap_to_vars(parammap, parameters(sys))
-    # identity function to make syms works
+
+    (u0map isa AbstractVector) || error("For DiscreteProblems u0map must be an AbstractVector.")
+    u0d = Dict( convert(Variable,u[1]) => u[2] for u in u0map)
+    u0 = [u0d[u] for u in states(sys)]
+    if parammap != DiffEqBase.NullParameters()
+        (parammap isa AbstractVector) || error("For DiscreteProblems parammap must be an AbstractVector.")
+        pd  = Dict( convert(Variable,u[1]) => u[2] for u in parammap)
+        p  = [pd[u] for u in parameters(sys)]
+    else
+        p = parammap
+    end
     # EvalFunc because we know that the jump functions are generated via eval
     f  = DiffEqBase.EvalFunc(DiffEqBase.DISCRETE_INPLACE_DEFAULT)
     df = DiscreteFunction(f, syms=Symbol.(states(sys)))
@@ -235,7 +264,7 @@ function DiffEqJump.JumpProblem(js::JumpSystem, prob, aggregator; kwargs...)
     parammap  = map((x,y)->Pair(x(),y), parameters(js), p)
     subber    = substituter(parammap)
 
-    majs = MassActionJump[assemble_maj(js, j, statetoid, subber, invttype) for j in eqs.x[1]]
+    majs = MassActionJump[assemble_maj(j, statetoid, subber, invttype) for j in eqs.x[1]]
     crjs = ConstantRateJump[assemble_crj(js, j, statetoid) for j in eqs.x[2]]
     vrjs = VariableRateJump[assemble_vrj(js, j, statetoid) for j in eqs.x[3]]
     ((prob isa DiscreteProblem) && !isempty(vrjs)) && error("Use continuous problems such as an ODEProblem or a SDEProblem with VariableRateJumps")

src/systems/reaction/reactionsystem.jl

@@ -69,12 +69,12 @@ function Reaction(rate, subs, prods, substoich, prodstoich;
       (isnothing(prodstoich)&&isnothing(substoich)) && error("Both substrate and product stochiometry inputs cannot be nothing.")
       if isnothing(subs)
         subs = Vector{Operation}()
-        (substoich!=nothing) && error("If substrates are nothing, substrate stiocihometries have to be so too.")
+        !isnothing(substoich) && error("If substrates are nothing, substrate stiocihometries have to be so too.")
         substoich = typeof(prodstoich)()
     end
     if isnothing(prods)
         prods = Vector{Operation}()
-        (prodstoich!=nothing) && error("If products are nothing, product stiocihometries have to be so too.")
+        !isnothing(prodstoich) && error("If products are nothing, product stiocihometries have to be so too.")
         prodstoich = typeof(substoich)()
     end
     ns = isnothing(netstoich) ? get_netstoich(subs, prods, substoich, prodstoich) : netstoich
@@ -140,8 +140,9 @@ function ReactionSystem(eqs, iv, species, params; systems = ReactionSystem[],
 
 
     isempty(species) && error("ReactionSystems require at least one species.")
-    paramvars = isempty(params) ? Variable[] : convert.(Variable, params)
-    ReactionSystem(eqs, iv, convert.(Variable,species), paramvars, name, systems)
+    paramvars = map(v -> convert(Variable,v), params)
+    specvars  = map(s -> convert(Variable,s), species)
+    ReactionSystem(eqs, convert(Variable,iv), specvars, paramvars, name, systems)
 end
 
 # Calculate the ODE rate law
@@ -220,7 +221,8 @@ end
 """
 ```julia
 ismassaction(rx, rs; rxvars = get_variables(rx.rate),
-                              haveivdep = any(var -> isequal(rs.iv,convert(Variable,var)), rxvars))
+                              haveivdep = any(var -> isequal(rs.iv,convert(Variable,var)), rxvars),
+                              stateset = Set(states(rs)))
 ```
 
 True if a given reaction is of mass action form, i.e. `rx.rate` does not depend
@@ -232,43 +234,71 @@ explicitly on the independent variable (usually time).
 - `rs`, a [`ReactionSystem`](@ref) containing the reaction.
 - Optional: `rxvars`, `Variable`s which are not in `rxvars` are ignored as possible dependencies.
 - Optional: `haveivdep`, `true` if the [`Reaction`](@ref) `rate` field explicitly depends on the independent variable.
+- Optional: `stateset`, set of states which if the rxvars are within mean rx is non-mass action.
 """
 function ismassaction(rx, rs; rxvars = get_variables(rx.rate),
-                              haveivdep = any(var -> isequal(rs.iv,convert(Variable,var)), rxvars))
+                              haveivdep = any(var -> isequal(rs.iv,convert(Variable,var)), rxvars),
+                              stateset = Set(states(rs)))
     # if no dependencies must be zero order
-    if isempty(rxvars)
-        return true
-    else
-        return !(haveivdep || rx.only_use_rate || any(convert(Variable,rxv) in states(rs) for rxv in rxvars))
+    (length(rxvars)==0) && return true
+    (haveivdep || rx.only_use_rate) && return false
+    @inbounds for i = 1:length(rxvars)
+        (rxvars[i].op in stateset) && return false
     end
+    return true
+end
+
+@inline function makemajump(rx)
+    @unpack rate, substrates, substoich, netstoich = rx
+    havesubstoich = (length(substoich) == 0)
+    reactant_stoch = Vector{Pair{Operation,eltype(substoich)}}(undef, length(substoich))
+    @inbounds for i = 1:length(reactant_stoch)
+        reactant_stoch[i] = var2op(substrates[i].op) => substoich[i]
+    end
+    #push!(rstoich, reactant_stoch)
+    coef           = havesubstoich ? one(eltype(substoich)) : prod(stoich -> factorial(stoich), substoich)
+    rate           = isone(coef) ? rate : rate/coef
+    #push!(rates, rate)
+    net_stoch      = [Pair(var2op(p[1]),p[2]) for p in netstoich]
+    #push!(nstoich, net_stoch)
+    MassActionJump(rate, reactant_stoch, net_stoch, scale_rates=false, useiszero=false)
 end
 
 function assemble_jumps(rs)
-    eqs = Vector{Union{ConstantRateJump, MassActionJump, VariableRateJump}}()
+    meqs = MassActionJump[]; ceqs = ConstantRateJump[]; veqs = VariableRateJump[]
+    stateset = Set(states(rs))
+    #rates = [];  rstoich = []; nstoich = []
+    rxvars = Operation[]
+    ivname = rs.iv.name
 
+    isempty(equations(rs)) && error("Must give at least one reaction before constructing a JumpSystem.")
     for rx in equations(rs)
-        rxvars    = (rx.rate isa Operation) ? get_variables(rx.rate) : Operation[]
-        haveivdep = any(var -> isequal(rs.iv,convert(Variable,var)), rxvars)
-        if ismassaction(rx, rs; rxvars=rxvars, haveivdep=haveivdep)
-            reactant_stoch = isempty(rx.substoich) ? [0 => 1] : [var2op(sub.op) => stoich for (sub,stoich) in zip(rx.substrates,rx.substoich)]
-            coef           = isempty(rx.substoich) ? one(eltype(rx.substoich)) : prod(stoich -> factorial(stoich), rx.substoich)
-            rate           = isone(coef) ? rx.rate : rx.rate/coef
-            net_stoch      = [Pair(var2op(p[1]),p[2]) for p in rx.netstoich]
-            push!(eqs, MassActionJump(rate, reactant_stoch, net_stoch, scale_rates=false))
+        empty!(rxvars)
+        (rx.rate isa Operation) && get_variables!(rxvars, rx.rate)
+        haveivdep = false
+        @inbounds for i = 1:length(rxvars)
+            if rxvars[i].op.name == ivname
+                haveivdep = true
+                break
+            end
+        end
+        if ismassaction(rx, rs; rxvars=rxvars, haveivdep=haveivdep, stateset=stateset)
+            push!(meqs, makemajump(rx))
         else
             rl     = jumpratelaw(rx, rxvars=rxvars)
             affect = Vector{Equation}()
             for (spec,stoich) in rx.netstoich
                 push!(affect, var2op(spec) ~ var2op(spec) + stoich)
             end
             if haveivdep
-                push!(eqs, VariableRateJump(rl,affect))
+                push!(veqs, VariableRateJump(rl,affect))
             else
-                push!(eqs, ConstantRateJump(rl,affect))
+                push!(ceqs, ConstantRateJump(rl,affect))
             end
         end
     end
-    eqs
+    #eqs[1] = MassActionJump(rates, rstoich, nstoich, scale_rates=false, useiszero=false)
+    ArrayPartition(meqs,ceqs,veqs)
 end
 
 """

test/reactionsystem.jl

@@ -84,12 +84,15 @@ G2 = sf.g(u,p,t)
 # test with JumpSystem
 js = convert(JumpSystem, rs)
 
-@test all(map(i -> typeof(js.eqs[i]) <: DiffEqJump.MassActionJump, 1:14))
-@test all(map(i -> typeof(js.eqs[i]) <: DiffEqJump.ConstantRateJump, 15:18))
-@test all(map(i -> typeof(js.eqs[i]) <: DiffEqJump.VariableRateJump, 19:20))
+midxs = 1:14
+cidxs = 15:18
+vidxs = 19:20
+@test all(map(i -> typeof(js.eqs[i]) <: DiffEqJump.MassActionJump, midxs))
+@test all(map(i -> typeof(js.eqs[i]) <: DiffEqJump.ConstantRateJump, cidxs))
+@test all(map(i -> typeof(js.eqs[i]) <: DiffEqJump.VariableRateJump, vidxs))
 
 pars = rand(length(k)); u0 = rand(1:10,4); time = rand();
-jumps = Vector{Union{ConstantRateJump, MassActionJump, VariableRateJump}}(undef,length(js.eqs))
+jumps = Vector{Union{ConstantRateJump, MassActionJump, VariableRateJump}}(undef,length(rxs))
 
 jumps[1] = MassActionJump(pars[1], Vector{Pair{Int,Int}}(), [1 => 1]);
 jumps[2] = MassActionJump(pars[2], [2 => 1], [2 => -1]);
@@ -116,20 +119,20 @@ jumps[20] = VariableRateJump((u,p,t) -> p[20]*t*u[1]*binomial(u[2],2)*u[3], inte
 
 statetoid = Dict(convert(Variable,state) => i for (i,state) in enumerate(states(js)))
 parammap = map((x,y)->Pair(x(),y),parameters(js),pars)
-for i = 1:14
-  maj = MT.assemble_maj(js, js.eqs[i], statetoid, ModelingToolkit.substituter(parammap),eltype(pars))
+for i in midxs
+  maj = MT.assemble_maj(js.eqs[i], statetoid, ModelingToolkit.substituter(parammap),eltype(pars))
   @test abs(jumps[i].scaled_rates - maj.scaled_rates) < 100*eps()
   @test jumps[i].reactant_stoch == maj.reactant_stoch
   @test jumps[i].net_stoch == maj.net_stoch
 end
-for i = 15:18
+for i in cidxs
   crj = MT.assemble_crj(js, js.eqs[i], statetoid)
   @test isapprox(crj.rate(u0,p,time), jumps[i].rate(u0,p,time))
   fake_integrator1 = (u=zeros(4),p=p,t=0); fake_integrator2 = deepcopy(fake_integrator1);
   crj.affect!(fake_integrator1); jumps[i].affect!(fake_integrator2);
   @test fake_integrator1 == fake_integrator2
 end
-for i = 19:20
+for i in vidxs
   crj = MT.assemble_vrj(js, js.eqs[i], statetoid)
   @test isapprox(crj.rate(u0,p,time), jumps[i].rate(u0,p,time))
   fake_integrator1 = (u=zeros(4),p=p,t=0.); fake_integrator2 = deepcopy(fake_integrator1);