Use single type for mass action kinetics

This avoids type instability in `rate` and `rate_derivative` due to the small type union.
Using a single type is probably better in the long run, but need to keep an eye on `rate_derivative` for higher-order (order >> 2)reactions.

src/model/reaction_system.jl

@@ -2,11 +2,7 @@
 
 abstract type KineticLaw end
 
-struct MassActionOrder0  <: KineticLaw end
-struct MassActionOrder1  <: KineticLaw end
-struct MassActionOrder2A <: KineticLaw end
-struct MassActionOrder2B <: KineticLaw end
-struct MassActionOrderN  <: KineticLaw end
+struct MassAction  <: KineticLaw end
 
 ##### ReactionStruct #####
 
@@ -25,11 +21,7 @@ struct ReactionStruct{MA <: KineticLaw}
     end
 end
 
-is_compatible_law(::MassActionOrder0,  order, num_reactants) = order == 0
-is_compatible_law(::MassActionOrder1,  order, num_reactants) = order == 1
-is_compatible_law(::MassActionOrder2A, order, num_reactants) = order == 2 && num_reactants == 2
-is_compatible_law(::MassActionOrder2B, order, num_reactants) = order == 2 && num_reactants == 1
-is_compatible_law(::MassActionOrderN, order, num_reactants) = true
+is_compatible_law(::MassAction, order, num_reactants) = true
 
 function execute_jump!(x, r::ReactionStruct)
     net_change = r.net_change
@@ -42,65 +34,48 @@ end
 
 ##### rate functions for stochastic mass action kinetics #####
 
-@inline @inbounds function rate(r::ReactionStruct{MassActionOrder0}, x, p)
+@inline @inbounds function rate(r::ReactionStruct{MassAction}, x, p)
     i = r.paramidx
+    total_rate = p[i]   # accumulates terms of the form (x_k - (j-1))
+    prefactor = one(total_rate) # accumulates constants 1 / m_k!
 
-    return p[i]
-end
-
-@inline @inbounds function rate(r::ReactionStruct{MassActionOrder1}, x, p)
-    i = r.paramidx
-    k, _ = r.reactants[1]
-
-    return p[i] * x[k]
-end
-
-@inline @inbounds function rate(r::ReactionStruct{MassActionOrder2A}, x, p)
-    i = r.paramidx
-    k1, _ = r.reactants[1]
-    k2, _ = r.reactants[2]
-
-    return p[i] * x[k1] * x[k2]
-end
+    for (k, m) in r.reactants
+        for j in 1:m
+            total_rate *= (x[k] - (j-1))
+            prefactor *= j
+        end
+    end
 
-@inline @inbounds function rate(r::ReactionStruct{MassActionOrder2B}, x, p)
-    i = r.paramidx
-    k, _ = r.reactants[1]
+    total_rate /= prefactor
 
-    return 0.5 * x[k] * (x[k] - 1) * p[i]
+    return total_rate
 end
 
-@inline @inbounds function rate(r::ReactionStruct{MassActionOrderN}, x, p)
-    i = r.paramidx
-    total_rate = p[i]   # accumulates terms of the form (x_k - (j-1))
+@inline @inbounds function rate_derivative(r::ReactionStruct{MassAction}, x, p, i)
+    total_rate = p[r.paramidx]  # accumulates terms of the form (x_k - (j-1))
     prefactor = one(total_rate) # accumulates constants 1 / m_k!
+    deriv_term = zero(total_rate)
 
     for (k, m) in r.reactants
         for j in 1:m
             total_rate *= (x[k] - (j-1))
             prefactor *= j
+
+            if k == i
+                deriv_term += 1 / (x[k] - (j-1))
+            end
         end
     end
 
-    total_rate /= prefactor
+    total_rate = total_rate / prefactor * deriv_term
 
     return total_rate
 end
 
-##### type union for heterogeneous ReactionStruct arrays #####
-
-ReactionLike = Union{
-ReactionStruct{MassActionOrder0},
-ReactionStruct{MassActionOrder1},
-ReactionStruct{MassActionOrder2A},
-ReactionStruct{MassActionOrder2B},
-ReactionStruct{MassActionOrderN}
-}
-
 ##### ReactionSystem #####
 
 struct ReactionSystem{R,DG<:DependencyGraph}
-    reactions::Vector{ReactionLike}
+    reactions::Vector{ReactionStruct{MassAction}}
     rxn_rates::R
     dep_graph::DG
     spc_graph::DG
@@ -110,7 +85,7 @@ end
 function ReactionSystem(model::Network)
     num_reactions = number_reactions(model)
 
-    reactions = Vector{ReactionLike}(undef, num_reactions)
+    reactions = Vector{ReactionStruct{MassAction}}(undef, num_reactions)
     rxn_rates = zeros(num_reactions)
 
     build_reactions!(reactions, rxn_rates, model)
@@ -130,6 +105,10 @@ end
 
 @inline @inbounds rate(rxn::ReactionSystem, x, j) = rate(rxn.reactions[j], x, rxn.rxn_rates)
 
+@inline @inbounds function rate_derivative(rxn, x, i, j)
+    rate_derivative(rxn.reactions[j], x, rxn.rxn_rates, i)
+end
+
 function netstoichiometry(rxn::ReactionSystem, num_species, num_reactions)
     V = zeros(Int, num_species, num_reactions)
     reactions = rxn.reactions
@@ -164,95 +143,23 @@ function build_reactions!(rxn_set, rxn_rates, model)
                 push!(net_change, (indexmap[s], change))
             end
         end
-        L = get_kinetic_law(rtuples)
+        klaw = get_kinetic_law(rtuples)
 
-        rxn_set[j] = ReactionStruct(L(), rtuples, net_change, j)
+        rxn_set[j] = ReactionStruct(klaw, rtuples, net_change, j)
         rxn_rates[j] = rxn_rate
     end
 
     return rxn_set
 end
 
 function get_kinetic_law(rtuples)
-    num_reactants = length(rtuples)
-    order = num_reactants > 0 ? sum(c for (_, c) in rtuples) : 0
-
-    if order == 0
-        MassActionOrder0
-    elseif order == 1
-        MassActionOrder1
-    elseif order == 2 && num_reactants == 2
-        MassActionOrder2A
-    elseif order == 2 && num_reactants == 1
-        MassActionOrder2B
-    else
-        MassActionOrderN
-    end
+    # num_reactants = length(rtuples)
+    # order = num_reactants > 0 ? sum(c for (_, c) in rtuples) : 0
+    return MassAction()
 end
 
 ##### extras #####
 
 function execute_leap!(state, stoichiometry, number_jumps)
     state += stoichiometry * number_jumps
 end
-
-function rate_derivative(rxn, x, i, j)
-    rate_derivative(rxn.reactions[j], x, rxn.rxn_rates, i)
-end
-
-function rate_derivative(r::ReactionStruct{MassActionOrder0}, x, p, i)
-    return zero(eltype(p))
-end
-
-function rate_derivative(r::ReactionStruct{MassActionOrder1}, x, p, i)
-    idx = r.paramidx
-    k, _ = r.reactants[1]
-
-    return i == k ? p[idx] : zero(eltype(p))
-end
-
-function rate_derivative(r::ReactionStruct{MassActionOrder2A}, x, p, i)
-    idx = r.paramidx
-    k1, _ = r.reactants[1]
-    k2, _ = r.reactants[2]
-
-    if i == k1
-        return p[idx] * x[k2]
-    elseif i == k2
-        return p[idx] * x[k1]
-    else
-        return zero(p[idx])
-    end
-end
-
-function rate_derivative(r::ReactionStruct{MassActionOrder2B}, x, p, i)
-    idx = r.paramidx
-    k, _ = r.reactants[1]
-
-    if i == k
-        return p[idx] * (x[k] - 0.5)
-    else
-        return zero(eltype(p))
-    end
-end
-
-function rate_derivative(r::ReactionStruct{MassActionOrderN}, x, p, i)
-    total_rate = p[r.paramidx]  # accumulates terms of the form (x_k - (j-1))
-    prefactor = one(total_rate) # accumulates constants 1 / m_k!
-    deriv_term = zero(total_rate)
-
-    for (k, m) in r.reactants
-        for j in 1:m
-            total_rate *= (x[k] - (j-1))
-            prefactor *= j
-
-            if k == i
-                deriv_term += 1 / (x[k] - (j-1))
-            end
-        end
-    end
-
-    total_rate = total_rate / prefactor * deriv_term
-
-    return total_rate
-end

test/model/reaction_system.jl

@@ -1,5 +1,5 @@
-import BioSimulator: MassActionOrder0, MassActionOrder1, MassActionOrder2A, MassActionOrder2B, MassActionOrderN
-import BioSimulator: ReactionStruct, ReactionLike, execute_jump!, rate, rate_derivative
+import BioSimulator: MassAction
+import BioSimulator: ReactionStruct, execute_jump!, rate, rate_derivative
 import BioSimulator: ReactionSystem
 
 # define notion of equality for ReactionStruct
@@ -51,18 +51,19 @@ end
   order3a = [
     ([(1, 1), (2, 1), (3, 1)], [(1, -1), (2, -1), (3, -1), (4, 1)], 1)
   ]
+
   # 3 * A + B + 2 * C --> D
   order3b = [
     ([(1, 3), (2, 1), (3, 2)], [(1, -3), (2, -1), (3, -2), (4, 1)], 1)
   ]
 
   @testset "ReactionStruct" begin
-    rxn0  = [ReactionStruct(MassActionOrder0(),  r, v, i) for (r, v, i) in order0]
-    rxn1  = [ReactionStruct(MassActionOrder1(),  r, v, i) for (r, v, i) in order1]
-    rxn2a = [ReactionStruct(MassActionOrder2A(), r, v, i) for (r, v, i) in order2a]
-    rxn2b = [ReactionStruct(MassActionOrder2B(), r, v, i) for (r, v, i) in order2b]
-    rxn3a = [ReactionStruct(MassActionOrderN(), r, v, i) for (r, v, i) in order3a]
-    rxn3b = [ReactionStruct(MassActionOrderN(), r, v, i) for (r, v, i) in order3b]
+    rxn0  = [ReactionStruct(MassAction(), r, v, i) for (r, v, i) in order0]
+    rxn1  = [ReactionStruct(MassAction(), r, v, i) for (r, v, i) in order1]
+    rxn2a = [ReactionStruct(MassAction(), r, v, i) for (r, v, i) in order2a]
+    rxn2b = [ReactionStruct(MassAction(), r, v, i) for (r, v, i) in order2b]
+    rxn3a = [ReactionStruct(MassAction(), r, v, i) for (r, v, i) in order3a]
+    rxn3b = [ReactionStruct(MassAction(), r, v, i) for (r, v, i) in order3b]
 
     @testset "fire_reaction!" begin
       result   = copy(x)
@@ -159,54 +160,54 @@ end
     m = ReactionSystem(model)
 
     @testset "MassActionOrder0" begin
-      expected = ReactionStruct(MassActionOrder0(), order0[1][1], order0[1][2], 1)
+      expected = ReactionStruct(MassAction(), order0[1][1], order0[1][2], 1)
 
       @test m.reactions[1] == expected
       @test m.reactions[1].paramidx == expected.paramidx
 
-      expected = ReactionStruct(MassActionOrder0(), order0[2][1], order0[2][2], 2)
+      expected = ReactionStruct(MassAction(), order0[2][1], order0[2][2], 2)
       @test m.reactions[2] == expected
       @test m.reactions[2].paramidx == expected.paramidx
     end
 
     @testset "MassActionOrder1" begin
-      expected = ReactionStruct(MassActionOrder1(), order1[1][1], order1[1][2], 3)
+      expected = ReactionStruct(MassAction(), order1[1][1], order1[1][2], 3)
 
       @test m.reactions[3] == expected
       @test m.reactions[3].paramidx == expected.paramidx
 
-      expected = ReactionStruct(MassActionOrder1(), order1[2][1], order1[2][2], 4)
+      expected = ReactionStruct(MassAction(), order1[2][1], order1[2][2], 4)
 
       @test m.reactions[4] == expected
       @test m.reactions[4].paramidx == expected.paramidx
 
-      expected = ReactionStruct(MassActionOrder1(), order1[3][1], order1[3][2], 5)
+      expected = ReactionStruct(MassAction(), order1[3][1], order1[3][2], 5)
 
       @test m.reactions[5] == expected
       @test m.reactions[5].paramidx == expected.paramidx
     end
 
     @testset "MassActionOrder2A" begin
-      expected = ReactionStruct(MassActionOrder2A(), order2a[1][1], order2a[1][2], 6)
+      expected = ReactionStruct(MassAction(), order2a[1][1], order2a[1][2], 6)
 
       @test m.reactions[6] == expected
       @test m.reactions[6].paramidx == expected.paramidx
     end
 
     @testset "MassActionOrder2B" begin
-      expected = ReactionStruct(MassActionOrder2B(), order2b[1][1], order2b[1][2], 7)
+      expected = ReactionStruct(MassAction(), order2b[1][1], order2b[1][2], 7)
 
       @test m.reactions[7] == expected
       @test m.reactions[7].paramidx == expected.paramidx
     end
 
     @testset "MassActionOrderN" begin
-      expected = ReactionStruct(MassActionOrderN(), order3a[1][1], order3a[1][2], 8)
+      expected = ReactionStruct(MassAction(), order3a[1][1], order3a[1][2], 8)
 
       @test m.reactions[8] == expected
       @test m.reactions[8].paramidx == expected.paramidx
 
-      expected = ReactionStruct(MassActionOrderN(), order3b[1][1], order3b[1][2], 9)
+      expected = ReactionStruct(MassAction(), order3b[1][1], order3b[1][2], 9)
       @test m.reactions[9] == expected
       @test m.reactions[9].paramidx == expected.paramidx
     end