Merge pull request #289 from isaacsas/rx_rate_laws

add ODE ratelaw expressions

Author: ChrisRackauckas

Project.toml

@@ -19,6 +19,7 @@ SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 TreeViews = "a2a6695c-b41b-5b7d-aed9-dbfdeacea5d7"
+UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [compat]

src/ModelingToolkit.jl

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

src/systems/reaction/reactionsystem.jl

@@ -1,9 +1,47 @@
-struct Reaction
+struct Reaction{S <: Variable, T <: Number}
     rate
-    reactants::Vector{Operation}
+    substrates::Vector{Operation}
     products::Vector{Operation}
+    substoich::Vector{T}
+    prodstoich::Vector{T}    
+    netstoich::Vector{Pair{S,T}}
+    only_use_rate::Bool
 end
 
+function Reaction(rate, subs, prods, substoich, prodstoich; 
+                  netstoich=nothing, only_use_rate=false, kwargs...)
+
+    subsv  = isnothing(subs) ? Vector{Operation}() : subs
+    prodsv = isnothing(prods) ? Vector{Operation}() : prods
+    ns = isnothing(netstoich) ? get_netstoich(subsv, prodsv, substoich, prodstoich) : netstoich
+    Reaction(rate, subsv, prodsv, substoich, prodstoich, ns, only_use_rate)
+end
+
+# three argument constructor assumes stoichiometric coefs are one and integers
+function Reaction(rate, subs, prods; kwargs...) 
+
+    sstoich = isnothing(subs) ? Int[] : ones(Int,length(subs))
+    pstoich = isnothing(prods) ? Int[] : ones(Int,length(prods))
+    Reaction(rate, subs, prods, sstoich, pstoich; kwargs...)
+end
+
+# calculates the net stoichiometry of a reaction as a vector of pairs (sub,substoich)
+function get_netstoich(subs, prods, sstoich, pstoich)
+    # stoichiometry as a Dictionary
+    nsdict = Dict{Variable,eltype(sstoich)}(sub.op => -sstoich[i] for (i,sub) in enumerate(subs))
+    for (i,p) in enumerate(prods)
+        coef = pstoich[i]
+        prod = p.op
+        @inbounds nsdict[prod] = haskey(nsdict, prod) ? nsdict[prod] + coef : coef
+    end
+
+    # stoichiometry as a vector
+    ns = [el for el in nsdict if el[2] != zero(el[2])]
+
+    ns
+end
+
+
 struct ReactionSystem <: AbstractSystem
     eqs::Vector{Reaction}
     iv::Variable
@@ -13,59 +51,74 @@ struct ReactionSystem <: AbstractSystem
     systems::Vector{ReactionSystem}
 end
 
-function ReactionSystem(eqs,iv,dvs,ps;
-                        systems = ReactionSystem[],
-                        name=gensym(:ReactionSystem))
-    ReactionSystem(eqs,iv,convert.(Variable,dvs),convert.(Variable,ps),name,systems)
+function ReactionSystem(eqs, iv, species, params; systems = ReactionSystem[],
+                                                  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, only_use_rate = rx    
+    rl = rate
+    if !only_use_rate
+        coef = one(eltype(substoich))
+        for (i,stoich) in enumerate(substoich)
+            coef *= factorial(stoich)        
+            rl   *= isone(stoich) ? substrates[i] : substrates[i]^stoich
+        end
+        (!isone(coef)) && (rl /= coef)
+    end
+    rl
+end
 
-function essemble_drift(rs)
-    D = Differential(rs.iv())
+function assemble_drift(rs)
+    D   = Differential(rs.iv())
     eqs = [D(x(rs.iv())) ~ 0 for x in rs.states]
+    species_to_idx = Dict((x => i for (i,x) in enumerate(rs.states)))
 
-    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))
-        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))
+    for rx in rs.eqs        
+        rl = oderatelaw(rx)
+        for (spec,stoich) in rx.netstoich
+            i = species_to_idx[spec]
+            if iszero(eqs[i].rhs)
+                signedrl = (stoich > zero(stoich)) ? rl : -rl
+                rhs      = isone(abs(stoich)) ? signedrl : stoich * rl                
+            else
+                Δspec = isone(abs(stoich)) ? rl : abs(stoich) * rl            
+                rhs   = (stoich > zero(stoich)) ? (eqs[i].rhs + Δspec) : (eqs[i].rhs - Δspec)
+            end
+            eqs[i] = Equation(eqs[i].lhs, rhs)
         end
     end
     eqs
 end
 
-function essemble_diffusion(rs)
+function assemble_diffusion(rs)
     eqs = Expression[Constant(0) for x in rs.states, y in rs.eqs]
+    species_to_idx = Dict((x => i for (i,x) in enumerate(rs.states)))
 
     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)
-        end
-
-        for product in rx.products
-            i = findfirst(x->x == product.op,rs.states)
-            eqs[i,j] = sqrt(rx.rate) * reaction_expr(rx.reactants)
+        rlsqrt = sqrt(oderatelaw(rx))
+        for (spec,stoich) in rx.netstoich
+            i            = species_to_idx[spec]
+            signedrlsqrt = (stoich > zero(stoich)) ? rlsqrt : -rlsqrt
+            eqs[i,j]     = isone(abs(stoich)) ? signedrlsqrt : stoich * rlsqrt            
         end
     end
     eqs
 end
 
 function Base.convert(::Type{<:ODESystem},rs::ReactionSystem)
-    eqs = essemble_drift(rs)
+    eqs = assemble_drift(rs)
     ODESystem(eqs,rs.iv,rs.states,rs.ps,name=rs.name,
               systems=convert.(ODESystem,rs.systems))
 end
 
 function Base.convert(::Type{<:SDESystem},rs::ReactionSystem)
-    eqs = essemble_drift(rs)
-    noiseeqs = essemble_diffusion(rs)
+    eqs = assemble_drift(rs)
+    noiseeqs = assemble_diffusion(rs)
     SDESystem(eqs,noiseeqs,rs.iv,rs.states,rs.ps,
               name=rs.name,systems=convert.(SDESystem,rs.systems))
 end

test/reactionsystem.jl

@@ -1,10 +1,68 @@
-using ModelingToolkit
-
-@parameters t a b c
-@variables x(t) y(t) z(t)
-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)
+using ModelingToolkit, LinearAlgebra, Test
+
+@parameters t k[1:13]
+@variables A(t) B(t) C(t) D(t)
+rxs = [Reaction(k[1], nothing, [A]),            # 0 -> A
+       Reaction(k[2], [B], nothing),            # B -> 0
+       Reaction(k[3],[A],[C]),                  # A -> C
+       Reaction(k[4], [C], [A,B]),              # C -> A + B        
+       Reaction(k[5], [C], [A], [1], [2]),      # C -> A + A
+       Reaction(k[6], [A,B], [C]),              # A + B -> C
+       Reaction(k[7], [B], [A], [2], [1]),      # 2B -> A
+       Reaction(k[8], [A,B], [A,C]),            # A + B -> A + C
+       Reaction(k[9], [A,B], [C,D]),            # A + B -> C + D
+       Reaction(k[10], [A], [C,D], [2], [1,1]), # 2A -> C + D
+       Reaction(k[11], [A], [A,B], [2], [1,1]), # 2A -> A + B
+       Reaction(k[12], [A,B,C], [C,D], [1,3,4], [2, 3]), # A+3B+4C -> 2C + 3D
+       Reaction(k[13]*A/(2+A), [A], nothing; only_use_rate=true)  # A -> 0 with custom rate 
+       ]  
+rs = ReactionSystem(rxs,t,[A,B,C,D],k)
+odesys = convert(ODESystem,rs)
+sdesys = convert(SDESystem,rs)
+
+# hard coded ODE rhs
+function oderhs(u,k,t)
+       A = u[1]; B = u[2]; C = u[3]; D = u[4];
+       du = zeros(eltype(u),4)
+       du[1] = k[1] - k[3]*A + k[4]*C + 2*k[5]*C - k[6]*A*B + k[7]*B^2/2 - k[9]*A*B - k[10]*A^2 - k[11]*A^2/2 - k[12]*A*B^3*C^4/144 - k[13]*A/(2+A)
+       du[2] = -k[2]*B + k[4]*C - k[6]*A*B - k[7]*B^2 - k[8]*A*B - k[9]*A*B + k[11]*A^2/2 - 3*k[12]*A*B^3*C^4/144
+       du[3] = k[3]*A - k[4]*C - k[5]*C  + k[6]*A*B + k[8]*A*B + k[9]*A*B + k[10]*A^2/2 - 2*k[12]*A*B^3*C^4/144
+       du[4] = k[9]*A*B + k[10]*A^2/2 + 3*k[12]*A*B^3*C^4/144
+       du
+end
+
+# sde noise coefs
+function sdenoise(u,k,t)
+       A = u[1]; B = u[2]; C = u[3]; D = u[4];
+       G = zeros(eltype(u),length(k),length(u))
+       z = zero(eltype(u))
+
+       G = [sqrt(k[1]) z  z z;
+            z -sqrt(k[2]*B) z z;
+            -sqrt(k[3]*A) z sqrt(k[3]*A) z;
+            sqrt(k[4]*C) sqrt(k[4]*C) -sqrt(k[4]*C) z;
+            2*sqrt(k[5]*C) z -sqrt(k[5]*C) z;
+            -sqrt(k[6]*A*B) -sqrt(k[6]*A*B) sqrt(k[6]*A*B) z;
+            sqrt(k[7]*B^2/2) -2*sqrt(k[7]*B^2/2) z z;
+            z -sqrt(k[8]*A*B) sqrt(k[8]*A*B) z;
+            -sqrt(k[9]*A*B) -sqrt(k[9]*A*B) sqrt(k[9]*A*B) sqrt(k[9]*A*B);
+            -2*sqrt(k[10]*A^2/2) z sqrt(k[10]*A^2/2) sqrt(k[10]*A^2/2);
+            -sqrt(k[11]*A^2/2) sqrt(k[11]*A^2/2) z z;
+            -sqrt(k[12]*A*B^3*C^4/144) -3*sqrt(k[12]*A*B^3*C^4/144) -2*sqrt(k[12]*A*B^3*C^4/144) 3*sqrt(k[12]*A*B^3*C^4/144);
+            -sqrt(k[13]*A/(2+A)) z z z]'
+       
+       return G
+end
+
+# test by evaluating drift and diffusion terms
+p  = rand(length(k))
+u  = rand(length(k))
+t  = 0.
+du = oderhs(u,p,t)
+G  = sdenoise(u,p,t)
+sdesys = convert(SDESystem,rs)
+sf = SDEFunction{false}(sdesys, states(rs), parameters(rs))
+du2 = sf.f(u,p,t)
+@test norm(du-du2) < 100*eps()
+G2 = sf.g(u,p,t)
+@test norm(G-G2) < 100*eps()