Fix ReactionSystems tests

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

Author: shashi

src/build_function.jl

@@ -209,7 +209,7 @@ Special Keyword Argumnets:
   - `DaggerForm()`: Multithreading and multiprocessing using Julia's Dagger.jl
     for dynamic scheduling and load balancing.
 - `conv`: The conversion function of the Operation to Expr. By default this uses
-  the `toexpr` function utilized in `convert(Expr,x)`.
+  the `toexpr` function.
 - `checkbounds`: For whether to enable bounds checking inside of the generated
   function. Defaults to false, meaning that `@inbounds` is applied.
 - `linenumbers`: Determines whether the generated function expression retains

src/systems/jumps/jumpsystem.jl

@@ -67,16 +67,22 @@ function JumpSystem(eqs, iv, states, ps;
     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),
-                                        independent_variable(js),
-                                        expression=Val{true})
+function generate_rate_function(js, rate)
+    build_function(rate, states(js), parameters(js),
+                   independent_variable(js),
+                   conv = states_to_sym(states(js)),
+                   expression=Val{true})
+end
 
-generate_affect_function(js, affect, outputidxs) = build_function(affect, states(js),
-                                                      parameters(js),
-                                                      independent_variable(js),
-                                                      expression=Val{true},
-                                                      headerfun=add_integrator_header,
-                                                      outputidxs=outputidxs)[2]
+function generate_affect_function(js, affect, outputidxs)
+    build_function(affect, states(js),
+                   parameters(js),
+                   conv = states_to_sym(states(js)),
+                   independent_variable(js),
+                   expression=Val{true},
+                   headerfun=add_integrator_header,
+                   outputidxs=outputidxs)[2]
+end
 
 function assemble_vrj(js, vrj, statetoid)
     rate   = eval(generate_rate_function(js, vrj.rate))
@@ -121,7 +127,7 @@ end
 function numericrstoich(mtrs::Vector{Pair{V,W}}, statetoid) where {V,W}
     rs = Vector{Pair{Int,W}}()
     for (spec,stoich) in mtrs
-        if !(spec isa Term) && iszero(spec)
+        if !(spec isa Term) && _iszero(spec)
             push!(rs, 0 => stoich)
         else
             push!(rs, statetoid[value(spec)] => stoich)
@@ -134,8 +140,8 @@ end
 function numericnstoich(mtrs::Vector{Pair{V,W}}, statetoid) where {V,W}
     ns = Vector{Pair{Int,W}}()
     for (spec,stoich) in mtrs
-        !(spec isa Term) && iszero(spec) && error("Net stoichiometry can not have a species labelled 0.")
-        push!(ns, statetoid[value(spec)] => stoich)
+        !(spec isa Term) && _iszero(spec) && error("Net stoichiometry can not have a species labelled 0.")
+        push!(ns, statetoid[spec] => stoich)
     end
     sort!(ns)
 end
@@ -278,13 +284,13 @@ end
 
 ### Functions to determine which states a jump depends on
 function get_variables!(dep, jump::Union{ConstantRateJump,VariableRateJump}, variables) 
-    (jump.rate isa Operation) && get_variables!(dep, jump.rate, variables)
+    (jump.rate isa Symbolic) && get_variables!(dep, jump.rate, variables)
     dep
 end
 
 function get_variables!(dep, jump::MassActionJump, variables)
     sr = jump.scaled_rates
-    (sr isa Term) && get_variables!(dep, sr, variables)
+    (sr isa Symbolic) && get_variables!(dep, sr, variables)
     for varasop in jump.reactant_stoch
         any(isequal(varasop[1]), variables) && push!(dep, varasop[1])
     end

src/systems/reaction/reactionsystem.jl

@@ -349,7 +349,7 @@ function assemble_jumps(rs; combinatoric_ratelaws=true)
             rl     = jumpratelaw(rx, rxvars=rxvars, combinatoric_ratelaw=combinatoric_ratelaws)
             affect = Vector{Equation}()
             for (spec,stoich) in rx.netstoich
-                push!(affect, var2op(spec) ~ var2op(spec) + stoich)
+                push!(affect, spec ~ spec + stoich)
             end
             if haveivdep
                 push!(veqs, VariableRateJump(rl,affect))
@@ -522,16 +522,16 @@ end
 
 # determine which species a reaction depends on
 function get_variables!(deps::Set, rx::Reaction, variables)
-    (rx.rate isa Term) && get_variables!(deps, rx.rate, variables)
+    (rx.rate isa Symbolic) && get_variables!(deps, rx.rate, variables)
     for s in rx.substrates
         push!(deps, s)
     end
     @show deps
 end
 
 # determine which species a reaction modifies
-function modified_states!(mstates, rx::Reaction, sts)
+function modified_states!(mstates, rx::Reaction, sts::Set)
     for (species,stoich) in rx.netstoich
-        (species in sts) && push!(mstates, species())
+        (species in sts) && push!(mstates, species)
     end
 end

src/utils.jl

@@ -65,15 +65,16 @@ Returns the variables in the expression
 """
 get_variables(e::Num, varlist=nothing) = get_variables(value(e), varlist)
 get_variables!(vars, e, varlist=nothing) = vars
-get_variables!(vars, e::Sym, varlist=nothing) = push!(vars, e)
 
 is_singleton(e::Term) = e.op isa Sym
 is_singleton(e::Sym) = true
 is_singleton(e) = false
 
-function get_variables!(vars, e::Term, varlist=nothing)
+get_variables!(vars, e::Number, varlist=nothing) = vars
+
+function get_variables!(vars, e::Symbolic, varlist=nothing)
     if is_singleton(e)
-        if isnothing(varlist) || e in varlist
+        if isnothing(varlist) || any(isequal(e), varlist)
             push!(vars, e)
         end
     else
@@ -86,9 +87,10 @@ function get_variables!(vars, e::Equation, varlist=nothing)
   get_variables!(vars, e.rhs, varlist)
 end
 
+get_variables(e, varlist=nothing) = get_variables!([], e, varlist)
+
 modified_states!(mstates, e::Equation, statelist=nothing) = get_variables!(mstates, e.lhs, statelist)
 
-get_variables(e, varlist=nothing) = get_variables!([], e, varlist)
 
 # variable substitution
 # Piracy but mild
@@ -122,3 +124,22 @@ macro showarr(x)
 end
 
 @deprecate substitute_expr!(expr,s) substitute(expr,s)
+
+function states_to_sym(states)
+    function _states_to_sym(O)
+        if O isa Equation
+            Expr(:(=), _states_to_sym(O.lhs), _states_to_sym(O.rhs))
+        elseif O isa Term
+            if isa(O.op, Sym)
+                any(isequal(O), states) && return O.op.name  # dependent variables
+                return build_expr(:call, Any[O.op.name; _states_to_sym.(O.args)])
+            else
+                return build_expr(:call, Any[Symbol(O.op); _states_to_sym.(O.args)])
+            end
+        elseif O isa Num
+            return _states_to_sym(value(O))
+        else
+            return toexpr(O)
+        end
+    end
+end

test/reactionsystem.jl

@@ -161,7 +161,8 @@ for i in cidxs
   crj = MT.assemble_crj(js, js.eqs[i], statetoid)
   @test isapprox(crj.rate(u0,p,time), jumps[i].rate(u0,p,time))
   fake_integrator1 = (u=zeros(4),p=p,t=0); fake_integrator2 = deepcopy(fake_integrator1);
-  crj.affect!(fake_integrator1); jumps[i].affect!(fake_integrator2);
+  crj.affect!(fake_integrator1);
+  jumps[i].affect!(fake_integrator2);
   @test fake_integrator1 == fake_integrator2
 end
 for i in vidxs
@@ -190,12 +191,12 @@ rxs = [Reaction(k1*S, [S,I], [I], [2,3], [2]),
 rs = ReactionSystem(rxs, t, [S,I,R], [k1,k2])
 @test isequal(ModelingToolkit.oderatelaw(rs.eqs[1]), k1*S*S^2*I^3/(factorial(2)*factorial(3)))
 @test_skip isequal(ModelingToolkit.jumpratelaw(rs.eqs[1]), k1*S*binomial(S,2)*binomial(I,3))
-dep = Set{Operation}()
-ModelingToolkit.get_variables!(dep, rxs[2], states(rs))
+dep = Set()
+ModelingToolkit.get_variables!(dep, rxs[2], Set(states(rs)))
 dep2  = Set([R,I])
 @test dep == dep2
-dep = Set{Operation}()
-ModelingToolkit.modified_states!(dep, rxs[2], states(rs))
+dep = Set()
+ModelingToolkit.modified_states!(dep, rxs[2], Set(states(rs)))
 @test dep == Set([R,I])
 
 isequal2(a,b) = isequal(simplify(a), simplify(b))