Merge pull request #920 from isaacsas/composing-reactionsystems

Composing `ReactionSystem`s

Author: isaacsas

.gitignore

@@ -2,3 +2,5 @@
 *.jl.*.cov
 *.jl.mem
 Manifest.toml
+.vscode
+.vscode/*

src/systems/reaction/reactionsystem.jl

@@ -92,6 +92,14 @@ function Reaction(rate, subs, prods; kwargs...)
     Reaction(rate, subs, prods, sstoich, pstoich; kwargs...)
 end
 
+function namespace_equation(rx::Reaction, name, iv)
+    Reaction(namespace_expr(rx.rate, name, iv), 
+             namespace_expr(rx.substrates, name, iv),
+             namespace_expr(rx.products, name, iv),
+             rx.substoich, rx.prodstoich,            
+             [namespace_expr(n[1],name,iv) => n[2] for n in rx.netstoich], rx.only_use_rate)
+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
@@ -134,7 +142,7 @@ struct ReactionSystem <: AbstractSystem
     """The name of the system"""
     name::Symbol
     """systems: The internal systems"""
-    systems::Vector{ReactionSystem}
+    systems::Vector
 
     function ReactionSystem(eqs, iv, states, ps, observed, name, systems)
         new(eqs, value(iv), value.(states), value.(ps), observed, name, systems)
@@ -143,13 +151,31 @@ end
 
 function ReactionSystem(eqs, iv, species, params;
                         observed = [],
-                        systems = ReactionSystem[],
+                        systems = [],
                         name = gensym(:ReactionSystem))
 
-    isempty(species) && error("ReactionSystems require at least one species.")
+    #isempty(species) && error("ReactionSystems require at least one species.")
     ReactionSystem(eqs, iv, species, params, observed, name, systems)
 end
 
+function ReactionSystem(iv; kwargs...)
+    ReactionSystem(Reaction[], iv, [], []; kwargs...)
+end
+
+function equations(sys::ModelingToolkit.ReactionSystem)
+    eqs = get_eqs(sys)
+    systems = get_systems(sys)
+    if isempty(systems)
+        return eqs
+    else
+        eqs = [eqs;
+               reduce(vcat,
+                      namespace_equations.(get_systems(sys));
+                      init=[])]
+        return eqs
+    end
+end
+
 """
     oderatelaw(rx; combinatoric_ratelaw=true)
 
@@ -187,11 +213,11 @@ function oderatelaw(rx; combinatoric_ratelaw=true)
 end
 
 function assemble_oderhs(rs; combinatoric_ratelaws=true)
-    sts = states(rs)
+    sts = get_states(rs)
     species_to_idx = Dict((x => i for (i,x) in enumerate(sts)))
     rhsvec         = Any[0 for i in eachindex(sts)]
 
-    for rx in equations(rs)
+    for rx in get_eqs(rs)
         rl = oderatelaw(rx; combinatoric_ratelaw=combinatoric_ratelaws)
         for (spec,stoich) in rx.netstoich
             i = species_to_idx[spec]
@@ -212,16 +238,16 @@ function assemble_drift(rs; combinatoric_ratelaws=true, as_odes=true)
     rhsvec = assemble_oderhs(rs; combinatoric_ratelaws=combinatoric_ratelaws)
     if as_odes
         D   = Differential(get_iv(rs))
-        eqs = [Equation(D(x),rhs) for (x,rhs) in zip(states(rs),rhsvec)]
+        eqs = [Equation(D(x),rhs) for (x,rhs) in zip(get_states(rs),rhsvec) if (!_iszero(rhs))]
     else
-        eqs = [Equation(0,rhs) for rhs in rhsvec]
+        eqs = [Equation(0,rhs) for rhs in rhsvec if (!_iszero(rhs))]
     end
     eqs
 end
 
 function assemble_diffusion(rs, noise_scaling; combinatoric_ratelaws=true)
-    sts = states(rs)
-    eqs  = Matrix{Any}(undef, length(sts), length(equations(rs)))
+    sts  = get_states(rs)
+    eqs  = Matrix{Any}(undef, length(sts), length(get_eqs(rs)))
     eqs .= 0
     species_to_idx = Dict((x => i for (i,x) in enumerate(sts)))
 
@@ -302,7 +328,7 @@ explicitly on the independent variable (usually time).
 """
 function ismassaction(rx, rs; rxvars = get_variables(rx.rate),
                               haveivdep = any(var -> isequal(get_iv(rs),var), rxvars),
-                              stateset = Set(states(rs)))
+                              stateset = Set(get_states(rs)))
     # if no dependencies must be zero order
     (length(rxvars)==0) && return true
     haveivdep && return false
@@ -331,7 +357,7 @@ end
 
 function assemble_jumps(rs; combinatoric_ratelaws=true)
     meqs = MassActionJump[]; ceqs = ConstantRateJump[]; veqs = VariableRateJump[]
-    stateset = Set(states(rs))
+    stateset = Set(get_states(rs))
     #rates = [];  rstoich = []; nstoich = []
     rxvars = []
     ivname = nameof(get_iv(rs))
@@ -378,10 +404,11 @@ law, i.e. for `2S -> 0` at rate `k` the ratelaw would be `k*S^2/2!`. If
 `combinatoric_ratelaws=false` then the ratelaw is `k*S^2`, i.e. the scaling factor is
 ignored.
 """
-function Base.convert(::Type{<:ODESystem}, rs::ReactionSystem; combinatoric_ratelaws=true)
-    eqs = assemble_drift(rs; combinatoric_ratelaws=combinatoric_ratelaws)
-    ODESystem(eqs,get_iv(rs),states(rs),get_ps(rs),name=nameof(rs),
-              systems=convert.(ODESystem,get_systems(rs)))
+function Base.convert(::Type{<:ODESystem}, rs::ReactionSystem; 
+                      name=nameof(rs), combinatoric_ratelaws=true, kwargs...)
+    eqs     = assemble_drift(rs; combinatoric_ratelaws=combinatoric_ratelaws)
+    systems = map(sys -> (sys isa ODESystem) ? sys : convert(ODESystem, sys), get_systems(rs))
+    ODESystem(eqs, get_iv(rs), get_states(rs), get_ps(rs), name=name, systems=systems)
 end
 
 """
@@ -397,9 +424,11 @@ law, i.e. for `2S -> 0` at rate `k` the ratelaw would be `k*S^2/2!`. If
 `combinatoric_ratelaws=false` then the ratelaw is `k*S^2`, i.e. the scaling factor is
 ignored.
 """
-function Base.convert(::Type{<:NonlinearSystem},rs::ReactionSystem; combinatoric_ratelaws=true)
-    eqs = assemble_drift(rs; combinatoric_ratelaws=combinatoric_ratelaws, as_odes=false)
-    NonlinearSystem(eqs,states(rs),get_ps(rs),name=nameof(rs),systems=convert.(NonlinearSystem,get_systems(rs)))
+function Base.convert(::Type{<:NonlinearSystem},rs::ReactionSystem;
+                      name=nameof(rs), combinatoric_ratelaws=true, kwargs...)
+    eqs     = assemble_drift(rs; combinatoric_ratelaws=combinatoric_ratelaws, as_odes=false)
+    systems = convert.(NonlinearSystem, get_systems(rs))
+    NonlinearSystem(eqs, get_states(rs), get_ps(rs), name=name, systems=systems)
 end
 
 """
@@ -423,7 +452,8 @@ Finally, a `Vector{Operation}` can be provided (the length must be equal to the
 Here the noise for each reaction is scaled by the corresponding parameter in the input vector.
 This input may contain repeat parameters.
 """
-function Base.convert(::Type{<:SDESystem},rs::ReactionSystem, combinatoric_ratelaws=true; noise_scaling=nothing)
+function Base.convert(::Type{<:SDESystem}, rs::ReactionSystem; 
+                      noise_scaling=nothing, name=nameof(rs), combinatoric_ratelaws=true, kwargs...)
 
     if noise_scaling isa Vector
         (length(noise_scaling)!=length(equations(rs))) &&
@@ -434,19 +464,14 @@ function Base.convert(::Type{<:SDESystem},rs::ReactionSystem, combinatoric_ratel
         noise_scaling = fill(value(noise_scaling),length(equations(rs)))
     end
 
-    eqs = assemble_drift(rs; combinatoric_ratelaws=combinatoric_ratelaws)
-
+    eqs      = assemble_drift(rs; combinatoric_ratelaws=combinatoric_ratelaws)
     noiseeqs = assemble_diffusion(rs,noise_scaling;
                                   combinatoric_ratelaws=combinatoric_ratelaws)
-
-    SDESystem(eqs,
-              noiseeqs,
-              get_iv(rs),
-              states(rs),
-              (noise_scaling===nothing) ?
-                    get_ps(rs) :
-                    union(get_ps(rs),toparam.(noise_scaling)),
-                    name=nameof(rs),systems=convert.(SDESystem,get_systems(rs)))
+    systems  = convert.(SDESystem, get_systems(rs))
+    SDESystem(eqs, noiseeqs, get_iv(rs), get_states(rs),
+              (noise_scaling===nothing) ? get_ps(rs) : union(get_ps(rs), toparam.(noise_scaling)),
+              name=name, 
+              systems=systems)
 end
 
 """
@@ -462,10 +487,11 @@ Notes:
   the ratelaw is `k*S*(S-1)`, i.e. the rate law is not normalized by the scaling
   factor.
 """
-function Base.convert(::Type{<:JumpSystem},rs::ReactionSystem; combinatoric_ratelaws=true)
-    eqs = assemble_jumps(rs; combinatoric_ratelaws=combinatoric_ratelaws)
-    JumpSystem(eqs,get_iv(rs),states(rs),get_ps(rs),name=nameof(rs),
-              systems=convert.(JumpSystem,get_systems(rs)))
+function Base.convert(::Type{<:JumpSystem},rs::ReactionSystem; 
+                      name=nameof(rs), combinatoric_ratelaws=true, kwargs...)
+    eqs     = assemble_jumps(rs; combinatoric_ratelaws=combinatoric_ratelaws)
+    systems = convert.(JumpSystem, get_systems(rs))
+    JumpSystem(eqs, get_iv(rs), get_states(rs), get_ps(rs), name=name, systems=systems)
 end
 
 
@@ -474,35 +500,35 @@ end
 
 # ODEProblem from AbstractReactionNetwork
 function DiffEqBase.ODEProblem(rs::ReactionSystem, u0, tspan, p=DiffEqBase.NullParameters(), args...; kwargs...)
-    return ODEProblem(convert(ODESystem,rs),u0,tspan,p, args...; kwargs...)
+    return ODEProblem(convert(ODESystem,rs; kwargs...),u0,tspan,p, args...; kwargs...)
 end
 
 # NonlinearProblem from AbstractReactionNetwork
 function DiffEqBase.NonlinearProblem(rs::ReactionSystem, u0, p=DiffEqBase.NullParameters(), args...; kwargs...)
-    return NonlinearProblem(convert(NonlinearSystem,rs), u0, p, args...; kwargs...)
+    return NonlinearProblem(convert(NonlinearSystem,rs; kwargs...), u0, p, args...; kwargs...)
 end
 
 
 # SDEProblem from AbstractReactionNetwork
 function DiffEqBase.SDEProblem(rs::ReactionSystem, u0, tspan, p=DiffEqBase.NullParameters(), args...; noise_scaling=nothing, kwargs...)
-    sde_sys = convert(SDESystem,rs,noise_scaling=noise_scaling)
-    p_matrix = zeros(length(states(rs)), length(equations(rs)))
+    sde_sys  = convert(SDESystem,rs;noise_scaling=noise_scaling, kwargs...)
+    p_matrix = zeros(length(get_states(rs)), length(get_eqs(rs)))
     return SDEProblem(sde_sys,u0,tspan,p,args...; noise_rate_prototype=p_matrix,kwargs...)
 end
 
 # DiscreteProblem from AbstractReactionNetwork
 function DiffEqBase.DiscreteProblem(rs::ReactionSystem, u0, tspan::Tuple, p=DiffEqBase.NullParameters(), args...; kwargs...)
-    return DiscreteProblem(convert(JumpSystem,rs), u0,tspan,p, args...; kwargs...)
+    return DiscreteProblem(convert(JumpSystem,rs; kwargs...), u0,tspan,p, args...; kwargs...)
 end
 
 # JumpProblem from AbstractReactionNetwork
 function DiffEqJump.JumpProblem(rs::ReactionSystem, prob, aggregator, args...; kwargs...)
-    return JumpProblem(convert(JumpSystem,rs), prob, aggregator, args...; kwargs...)
+    return JumpProblem(convert(JumpSystem,rs; kwargs...), prob, aggregator, args...; kwargs...)
 end
 
 # SteadyStateProblem from AbstractReactionNetwork
 function DiffEqBase.SteadyStateProblem(rs::ReactionSystem, u0, p=DiffEqBase.NullParameters(), args...; kwargs...)
-    return SteadyStateProblem(ODEFunction(convert(ODESystem,rs)),u0,p, args...; kwargs...)
+    return SteadyStateProblem(ODEFunction(convert(ODESystem,rs; kwargs...)),u0,p, args...; kwargs...)
 end
 
 # determine which species a reaction depends on

test/reactionsystem_components.jl

@@ -0,0 +1,86 @@
+using ModelingToolkit, LinearAlgebra, OrdinaryDiffEq, Test
+MT = ModelingToolkit
+
+# Repressilator model
+@parameters t α₀ α K n δ β μ 
+@variables m(t) P(t) R(t)
+rxs = [Reaction(α₀, nothing, [m]),
+       Reaction(α / (1 + (R/K)^n), nothing, [m]),
+       Reaction(δ, [m], nothing),
+       Reaction(β, [m], [m,P]),
+       Reaction(μ, [P], nothing)    
+]
+
+specs = [m,P,R]
+pars  = [α₀,α,K,n,δ,β,μ]
+@named rs = ReactionSystem(rxs, t, specs, pars)
+
+# using ODESystem components
+@named os₁ = convert(ODESystem, rs)
+@named os₂ = convert(ODESystem, rs)
+@named os₃ = convert(ODESystem, rs)
+connections = [os₁.R ~ os₃.P, 
+               os₂.R ~ os₁.P,
+               os₃.R ~ os₂.P]
+@named connected = ODESystem(connections, t, [], [], systems=[os₁,os₂,os₃])
+oderepressilator = structural_simplify(connected)
+
+pvals = [os₁.α₀ => 5e-4, 
+         os₁.α => .5, 
+         os₁.K => 40.0, 
+         os₁.n => 2, 
+         os₁.δ => (log(2)/120), 
+         os₁.β => (20*log(2)/120),
+         os₁.μ => (log(2)/600),
+         os₂.α₀ => 5e-4, 
+         os₂.α => .5, 
+         os₂.K => 40.0, 
+         os₂.n => 2, 
+         os₂.δ => (log(2)/120), 
+         os₂.β => (20*log(2)/120),
+         os₂.μ => (log(2)/600),
+         os₃.α₀ => 5e-4, 
+         os₃.α => .5, 
+         os₃.K => 40.0, 
+         os₃.n => 2, 
+         os₃.δ => (log(2)/120), 
+         os₃.β => (20*log(2)/120),
+         os₃.μ => (log(2)/600)]
+u₀    = [os₁.m => 0.0, os₁.P => 20.0, os₂.m => 0.0, os₂.P => 0.0, os₃.m => 0.0, os₃.P => 0.0]
+tspan = (0.0, 100000.0)
+oprob = ODEProblem(oderepressilator, u₀, tspan, pvals)
+sol = solve(oprob, Tsit5())
+
+# hardcoded network
+function repress!(f, y, p, t)
+    α = p.α; α₀ = p.α₀; β = p.β; δ = p.δ; μ = p.μ; K = p.K; n = p.n
+    f[1] = α / (1 + (y[6] / K)^n) - δ * y[1] + α₀
+    f[2] = α / (1 + (y[4] / K)^n) - δ * y[2] + α₀
+    f[3] = α / (1 + (y[5] / K)^n) - δ * y[3] + α₀
+    f[4] = β * y[1] - μ * y[4]
+    f[5] = β * y[2] - μ * y[5]    
+    f[6] = β * y[3] - μ * y[6]    
+    nothing
+end
+ps = (α₀=5e-4, α=.5, K=40.0, n=2, δ=(log(2)/120), β=(20*log(2)/120), μ=(log(2)/600))
+u0 = [0.0,0.0,0.0,20.0,0.0,0.0]
+oprob2 = ODEProblem(repress!, u0, tspan, ps)
+sol2 = solve(oprob2, Tsit5())
+tvs = 0:1:tspan[end]
+
+indexof(sym,syms) = findfirst(isequal(sym),syms)
+i = indexof(os₁.P, states(oderepressilator))
+@test all(isapprox(u[1],u[2],atol=1e-4) for u in zip(sol(tvs, idxs=2), sol2(tvs, idxs=4)))
+
+# using ReactionSystem components
+
+# @named rs₁ = ReactionSystem(rxs, t, specs, pars)
+# @named rs₂ = ReactionSystem(rxs, t, specs, pars)
+# @named rs₃ = ReactionSystem(rxs, t, specs, pars)
+# connections = [rs₁.R ~ rs₃.P, 
+#                rs₂.R ~ rs₁.P,
+#                rs₃.R ~ rs₂.P]
+# @named csys = ODESystem(connections, t, [], [])
+# @named repressilator = ReactionSystem(t; systems=[csys,rs₁,rs₂,rs₃])
+# @named oderepressilator2 = convert(ODESystem, repressilator)
+# sys2 = structural_simplify(oderepressilator2)  # FAILS currently

test/runtests.jl

@@ -15,6 +15,7 @@ using SafeTestsets, Test
 @safetestset "NonlinearSystem Test" begin include("nonlinearsystem.jl") end
 @safetestset "OptimizationSystem Test" begin include("optimizationsystem.jl") end
 @safetestset "ReactionSystem Test" begin include("reactionsystem.jl") end
+@safetestset "ReactionSystem Test" begin include("reactionsystem_components.jl") end
 @safetestset "JumpSystem Test" begin include("jumpsystem.jl") end
 @safetestset "ControlSystem Test" begin include("controlsystem.jl") end
 @safetestset "Domain Test" begin include("domains.jl") end