jump system stuff

Co-authored-by: "Yingbo Ma" <[email protected]>

Author: shashi

src/systems/jumps/jumpsystem.jl

@@ -32,12 +32,12 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractSystem
     """
     eqs::U
     """The independent variable, usually time."""
-    iv::Variable
+    iv::Any
     """The dependent variables, representing the state of the system."""
-    states::Vector{Variable}
+    states::Vector
     """The parameters of the system."""
-    ps::Vector{Variable}
-    pins::Vector{Variable}
+    ps::Vector
+    pins::Vector
     observed::Vector{Equation}
     """The name of the system."""
     name::Symbol
@@ -46,7 +46,7 @@ struct JumpSystem{U <: ArrayPartition} <: AbstractSystem
 end
 
 function JumpSystem(eqs, iv, states, ps;
-                    pins = Variable[],
+                    pins = [],
                     observed = Equation[],
                     systems = JumpSystem[],
                     name = gensym(:JumpSystem))
@@ -64,7 +64,7 @@ function JumpSystem(eqs, iv, states, ps;
         end
     end
 
-    JumpSystem{typeof(ap)}(ap, convert(Variable,iv), convert.(Variable, states), convert.(Variable, ps), pins, observed, name, systems)
+    JumpSystem{typeof(ap)}(ap, value(iv), value.(states), value.(ps), pins, observed, name, systems)
 end
 
 generate_rate_function(js, rate) = build_function(rate, states(js), parameters(js),

src/systems/reaction/reactionsystem.jl

@@ -42,13 +42,13 @@ Notes:
 - The three-argument form assumes all reactant and product stoichiometric coefficients
   are one.
 """
-struct Reaction{S <: Variable, T <: Number}
+struct Reaction{S, T <: Number}
     """The rate function (excluding mass action terms)."""
     rate
     """Reaction substrates."""
-    substrates::Vector{Operation}
+    substrates::Vector
     """Reaction products."""
-    products::Vector{Operation}
+    products::Vector
     """The stoichiometric coefficients of the reactants."""
     substoich::Vector{T}
     """The stoichiometric coefficients of the products."""
@@ -65,20 +65,22 @@ end
 function Reaction(rate, subs, prods, substoich, prodstoich;
                   netstoich=nothing, only_use_rate=false, kwargs...)
 
-      (isnothing(prods)&&isnothing(subs)) && error("A reaction requires a non-nothing substrate or product vector.")
-      (isnothing(prodstoich)&&isnothing(substoich)) && error("Both substrate and product stochiometry inputs cannot be nothing.")
-      if isnothing(subs)
-        subs = Vector{Operation}()
+    (isnothing(prods)&&isnothing(subs)) && error("A reaction requires a non-nothing substrate or product vector.")
+    (isnothing(prodstoich)&&isnothing(substoich)) && error("Both substrate and product stochiometry inputs cannot be nothing.")
+    if isnothing(subs)
+        subs = Vector{Term}()
         !isnothing(substoich) && error("If substrates are nothing, substrate stiocihometries have to be so too.")
         substoich = typeof(prodstoich)()
     end
     if isnothing(prods)
-        prods = Vector{Operation}()
+        prods = Vector{Term}()
         !isnothing(prodstoich) && error("If products are nothing, product stiocihometries have to be so too.")
         prodstoich = typeof(substoich)()
     end
+    subs = value.(subs)
+    prods = value.(prods)
     ns = isnothing(netstoich) ? get_netstoich(subs, prods, substoich, prodstoich) : netstoich
-    Reaction(rate, subs, prods, substoich, prodstoich, ns, only_use_rate)
+    Reaction(value(rate), subs, prods, substoich, prodstoich, ns, only_use_rate)
 end
 
 
@@ -93,11 +95,10 @@ 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))
+    nsdict = Dict{Any, eltype(sstoich)}(sub => -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
+        @inbounds nsdict[p] = haskey(nsdict, p) ? nsdict[p] + coef : coef
     end
 
     # stoichiometry as a vector
@@ -124,12 +125,12 @@ struct ReactionSystem <: AbstractSystem
     """The reactions defining the system."""
     eqs::Vector{Reaction}
     """Independent variable (usually time)."""
-    iv::Variable
+    iv::Any
     """Dependent (state) variables representing amount of each species."""
-    states::Vector{Variable}
+    states::Vector
     """Parameter variables."""
-    ps::Vector{Variable}
-    pins::Vector{Variable}
+    ps::Vector
+    pins::Vector
     observed::Vector{Equation}
     """The name of the system"""
     name::Symbol
@@ -138,15 +139,13 @@ struct ReactionSystem <: AbstractSystem
 end
 
 function ReactionSystem(eqs, iv, species, params;
-                        pins = Variable[],
-                        observed = Operation[],
+                        pins = [],
+                        observed = [],
                         systems = ReactionSystem[],
                         name = gensym(:ReactionSystem))
 
     isempty(species) && error("ReactionSystems require at least one species.")
-    paramvars = map(v -> convert(Variable,v), params)
-    specvars  = map(s -> convert(Variable,s), species)
-    ReactionSystem(eqs, convert(Variable,iv), specvars, paramvars,
+    ReactionSystem(eqs, value(iv), value.(species), value.(params),
                    pins, observed, name, systems)
 end
 
@@ -187,15 +186,15 @@ function oderatelaw(rx; combinatoric_ratelaw=true)
 end
 
 function assemble_drift(rs; combinatoric_ratelaws=true)
-    D   = Differential(rs.iv())
-    eqs = [D(x(rs.iv())) ~ 0 for x in rs.states]
+    D   = Differential(rs.iv)
+    eqs = [D(x) ~ 0 for x in rs.states]
     species_to_idx = Dict((x => i for (i,x) in enumerate(rs.states)))
 
     for rx in rs.eqs
         rl = oderatelaw(rx; combinatoric_ratelaw=combinatoric_ratelaws)
         for (spec,stoich) in rx.netstoich
             i = species_to_idx[spec]
-            if iszero(eqs[i].rhs)
+            if _iszero(eqs[i].rhs)
                 signedrl = (stoich > zero(stoich)) ? rl : -rl
                 rhs      = isone(abs(stoich)) ? signedrl : stoich * rl
             else
@@ -209,7 +208,7 @@ function assemble_drift(rs; combinatoric_ratelaws=true)
 end
 
 function assemble_diffusion(rs, noise_scaling; combinatoric_ratelaws=true)
-    eqs = Expression[Constant(0) for x in rs.states, y in rs.eqs]
+    eqs = fill(Num(0), length(rs.states), length(rs.eqs))
     species_to_idx = Dict((x => i for (i,x) in enumerate(rs.states)))
 
     for (j,rx) in enumerate(rs.eqs)
@@ -281,7 +280,7 @@ 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,var), rxvars),
                               stateset = Set(states(rs)))
 ```
 
@@ -297,7 +296,7 @@ explicitly on the independent variable (usually time).
 - 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(isequal(rs.iv), rxvars),
                               stateset = Set(states(rs)))
     # if no dependencies must be zero order
     (length(rxvars)==0) && return true
@@ -311,7 +310,7 @@ end
 @inline function makemajump(rx; combinatoric_ratelaw=true)
     @unpack rate, substrates, substoich, netstoich = rx
     zeroorder = (length(substoich) == 0)
-    reactant_stoch = Vector{Pair{Operation,eltype(substoich)}}(undef, length(substoich))
+    reactant_stoch = Vector{Pair{Term,eltype(substoich)}}(undef, length(substoich))
     @inbounds for i = 1:length(reactant_stoch)
         reactant_stoch[i] = var2op(substrates[i].op) => substoich[i]
     end
@@ -321,27 +320,22 @@ end
     #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)
+
+    #XXX:              vv--- this sucks
+    MassActionJump(Num(rate), reactant_stoch, net_stoch, scale_rates=false, useiszero=false)
 end
 
 function assemble_jumps(rs; combinatoric_ratelaws=true)
     meqs = MassActionJump[]; ceqs = ConstantRateJump[]; veqs = VariableRateJump[]
     stateset = Set(states(rs))
     #rates = [];  rstoich = []; nstoich = []
-    rxvars = Operation[]
-    ivname = rs.iv.name
+    rxvars = []
 
     isempty(equations(rs)) && error("Must give at least one reaction before constructing a JumpSystem.")
     for rx in equations(rs)
         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
+        (rx.rate isa Term) && get_variables!(rxvars, rx.rate)
+        haveivdep = any(isequal(rs.iv), rxvars)
         if ismassaction(rx, rs; rxvars=rxvars, haveivdep=haveivdep, stateset=stateset)
             push!(meqs, makemajump(rx, combinatoric_ratelaw=combinatoric_ratelaws))
         else
@@ -379,6 +373,9 @@ function Base.convert(::Type{<:ODESystem}, rs::ReactionSystem; combinatoric_rate
               systems=convert.(ODESystem,rs.systems))
 end
 
+makeparam(s::Sym) = Sym{Parameter{Number}}(s.name)
+makeparam(s::Sym{<:Parameter}) = s
+
 """
 ```julia
 Base.convert(::Type{<:SDESystem},rs::ReactionSystem)
@@ -400,12 +397,30 @@ 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::Union{Vector{Operation},Operation,Nothing})
-    (typeof(noise_scaling) <: Vector{Operation}) && (length(noise_scaling)!=length(rs.eqs)) && error("The number of elements in 'noise_scaling' must be equal to the number of reactions in the reaction system.")
-    (typeof(noise_scaling) <: Operation) && (noise_scaling = fill(noise_scaling,length(rs.eqs)))
+function Base.convert(::Type{<:SDESystem},rs::ReactionSystem, combinatoric_ratelaws=true; noise_scaling=nothing)
+
+    if noise_scaling isa Vector
+        (length(noise_scaling)!=length(rs.eqs)) &&
+        error("The number of elements in 'noise_scaling' must be equal " *
+              "to the number of reactions in the reaction system.")
+        noise_scaling = value.(noise_scaling)
+    elseif !isnothing(noise_scaling)
+        noise_scaling = fill(value(noise_scaling),length(rs.eqs))
+    end
+
     eqs = assemble_drift(rs; combinatoric_ratelaws=combinatoric_ratelaws)
-    noiseeqs = assemble_diffusion(rs,noise_scaling; combinatoric_ratelaws=combinatoric_ratelaws)
-    SDESystem(eqs,noiseeqs,rs.iv,rs.states,(noise_scaling===nothing) ? rs.ps : union(rs.ps,Variable{ModelingToolkit.Parameter{Number}}.(noise_scaling)),name=rs.name,systems=convert.(SDESystem,rs.systems))
+
+    noiseeqs = assemble_diffusion(rs,noise_scaling;
+                                  combinatoric_ratelaws=combinatoric_ratelaws)
+
+    SDESystem(eqs,
+              noiseeqs,
+              rs.iv,
+              rs.states,
+              (noise_scaling===nothing) ?
+                    rs.ps :
+                    union(rs.ps,makeparam.(noise_scaling)),
+                    name=rs.name,systems=convert.(SDESystem,rs.systems))
 end
 
 """
@@ -499,7 +514,7 @@ function DiffEqBase.SteadyStateProblem(rs::ReactionSystem, u0::Union{AbstractArr
 end
 
 # determine which species a reaction depends on
-function get_variables!(deps::Set{Operation}, rx::Reaction, variables)
+function get_variables!(deps::Set, rx::Reaction, variables)
     (rx.rate isa Operation) && get_variables!(deps, rx.rate, variables)
     for s in rx.substrates
         push!(deps, s)

test/reactionsystem.jl

@@ -149,7 +149,7 @@ jumps[18] = ConstantRateJump((u,p,t) -> p[18]*u[2], integrator -> (integrator.u[
 jumps[19] = VariableRateJump((u,p,t) -> p[19]*u[1]*t, integrator -> (integrator.u[1] -= 1; integrator.u[2] += 1))
 jumps[20] = VariableRateJump((u,p,t) -> p[20]*t*u[1]*binomial(u[2],2)*u[3], integrator -> (integrator.u[2] -= 2; integrator.u[3] -= 1; integrator.u[4] += 2))
 
-statetoid = Dict(convert(Variable,state) => i for (i,state) in enumerate(states(js)))
+statetoid = Dict(state => i for (i,state) in enumerate(states(js)))
 parammap = map((x,y)->Pair(x(),y),parameters(js),pars)
 for i in midxs
   maj = MT.assemble_maj(js.eqs[i], statetoid, ModelingToolkit.substituter(parammap),eltype(pars))