add ODE ratelaw expressions

Author: isaacsas

Project.toml

@@ -13,6 +13,7 @@ Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 NaNMath = "77ba4419-2d1f-58cd-9bb1-8ffee604a2e3"
+Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"

src/ModelingToolkit.jl

@@ -4,6 +4,7 @@ using DiffEqBase, Distributed
 using StaticArrays, LinearAlgebra, SparseArrays
 using Latexify, Unitful
 using MacroTools
+using Parameters: @unpack
 
 using Base.Threads
 import MacroTools: splitdef, combinedef, postwalk, striplines

src/systems/reaction/reactionsystem.jl

@@ -1,71 +1,97 @@
-struct Reaction
+struct Reaction{T <: Number}
     rate
-    reactants::Vector{Operation}
+    substrates::Vector{Operation}
     products::Vector{Operation}
+    substoich::Vector{T}
+    prodstoich::Vector{T}
 end
 
 struct ReactionSystem <: AbstractSystem
     eqs::Vector{Reaction}
     iv::Variable
-    states::Vector{Variable}
-    ps::Vector{Variable}
+    species::Vector{Variable}
+    params::Vector{Variable}
     name::Symbol
     systems::Vector{ReactionSystem}
 end
 
-function ReactionSystem(eqs,iv,dvs,ps;
+function ReactionSystem(eqs, iv, species, params;
                         systems = ReactionSystem[],
-                        name=gensym(:ReactionSystem))
-    ReactionSystem(eqs,iv,convert.(Variable,dvs),convert.(Variable,ps),name,systems)
+                        name = gensym(:ReactionSystem))
+
+    ReactionSystem(eqs, iv, convert.(Variable,species), 
+                   convert.(Variable,params), name, systems)
 end
 
-# TODO: Make it do the combinatorics stuff
-reaction_expr(reactants) = *(reactants...)
+# Calculate the ODE rate law
+function oderatelaw(rx)
+    @unpack rate, substrates, substoich = rx    
+    rl   = rate
+    coef = one(eltype(substoich))
+    for (i,stoich) in enumerate(substoich)
+        coef *= factorial(stoich)        
+        rl   *= (stoich != one(stoich)) ? substrates[i]^stoich : substrates[i]
+    end
+    (coef != one(coef)) && (rl /= coef)
+
+    rl
+end
 
 function essemble_drift(rs)
-    D = Differential(rs.iv())
-    eqs = [D(x(rs.iv())) ~ 0 for x in rs.states]
+    D   = Differential(rs.iv())
+    eqs = [D(x(rs.iv())) ~ 0 for x in rs.species]
+    species_to_idx = Dict(enumerate(rs.species))
 
-    for rx in rs.eqs
-        for reactant in rx.reactants
-            i = findfirst(x->x == reactant.op,rs.states)
-            eqs[i] = Equation(eqs[i].lhs,eqs[i].rhs - rx.rate * reaction_expr(rx.reactants))
+    # note, this should really use the net stoichiometry to avoid 
+    # adding and substracting the same term for A + X -> B + X
+    for rx in rs.eqs        
+        rl     = oderatelaw(rx.rate, rx.substrates, rx.substoich)
+        stoich = rx.substoich
+        for (sidx,substrate) in enumerate(rx.substrates)
+            i      = species_to_idx[substrate.op]
+            eqs[i] = Equation(eqs[i].lhs, eqs[i].rhs - stoich[sidx]*rl)
         end
 
-        for product in rx.products
-            i = findfirst(x->x == product.op,rs.states)
-            eqs[i] = Equation(eqs[i].lhs,eqs[i].rhs + rx.rate * reaction_expr(rx.reactants))
+        stoich = rx.prodstoich
+        for (pidx,product) in enumerate(rx.products)
+            i      = species_to_idx[product.op]
+            eqs[i] = Equation(eqs[i].lhs,eqs[i].rhs + stoich[pidx]*rl)
         end
     end
+
     eqs
 end
 
 function essemble_diffusion(rs)
-    eqs = Expression[Constant(0) for x in rs.states, y in rs.eqs]
+    eqs = Expression[Constant(0) for x in rs.species, y in rs.eqs]
+    species_to_idx = Dict(enumerate(rs.species))
 
     for (j,rx) in enumerate(rs.eqs)
-        for reactant in rx.reactants
-            i = findfirst(x->x == reactant.op,rs.states)
-            eqs[i,j] = -sqrt(rx.rate) * reaction_expr(rx.reactants)
+        rlsqrt = sqrt(oderatelaw(rx.rate, rx.substrates, rx.substoich))
+        stoich = rx.substoich
+        for (sidx,substrate) in enumerate(rx.substrates)
+            i        = species_to_idx[substrate.op]
+            eqs[i,j] = -stoich[sidx] * rlsqrt
         end
 
-        for product in rx.products
-            i = findfirst(x->x == product.op,rs.states)
-            eqs[i,j] = sqrt(rx.rate) * reaction_expr(rx.reactants)
+        for (pidx,product) in enumerate(rx.products)
+            i        = species_to_idx[product.op]
+            Δp       = stoich[pidx] * rlsqrt            
+            eqs[i,j] = (eqs[i,j]==Constant(0)) ? Δp : (eqs[i,j] + Δp)
         end
     end
     eqs
 end
 
 function Base.convert(::Type{<:ODESystem},rs::ReactionSystem)
     eqs = essemble_drift(rs)
-    ODESystem(eqs,rs.iv,rs.states,rs.ps,name=rs.name,
+    ODESystem(eqs,rs.iv,rs.species,rs.params,name=rs.name,
               systems=convert.(ODESystem,rs.systems))
 end
 
 function Base.convert(::Type{<:SDESystem},rs::ReactionSystem)
     eqs = essemble_drift(rs)
     noiseeqs = essemble_diffusion(rs)
-    SDESystem(eqs,noiseeqs,rs.iv,rs.states,rs.ps,
+    SDESystem(eqs,noiseeqs,rs.iv,rs.species,rs.params,
               name=rs.name,systems=convert.(SDESystem,rs.systems))
 end

test/reactionsystem.jl

@@ -6,5 +6,5 @@ rxs = [Reaction(a,[x],[y]),
        Reaction(b,[y],[x]),
        Reaction(c,[x,y],[x])]
 rs = ReactionSystem(rxs,t,[x,y],[a,b])
-sys = convert(ODESystem,rs)
-sys = convert(SDESystem,rs)
+odesys = convert(ODESystem,rs)
+sdesys = convert(SDESystem,rs)