fix conversion of ReactionSystems to NonlinearSystems

Author: isaacsas

src/systems/nonlinear/nonlinearsystem.jl

@@ -75,7 +75,7 @@ jacobian_sparsity(sys::NonlinearSystem) =
 
 """
 ```julia
-function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem,u0map,tspan,
+function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem,u0map,
                                           parammap=DiffEqBase.NullParameters();
                                           jac = false, sparse=false,
                                           checkbounds = false,
@@ -104,7 +104,7 @@ end
 
 """
 ```julia
-function DiffEqBase.NonlinearProblemExpr{iip}(sys::NonlinearSystem,u0map,tspan,
+function DiffEqBase.NonlinearProblemExpr{iip}(sys::NonlinearSystem,u0map,
                                           parammap=DiffEqBase.NullParameters();
                                           jac = false, sparse=false,
                                           checkbounds = false,
@@ -118,7 +118,7 @@ numerical enhancements.
 """
 struct NonlinearProblemExpr{iip} end
 
-function NonlinearProblemExpr{iip}(sys::NonlinearSystem,u0map,tspan,
+function NonlinearProblemExpr{iip}(sys::NonlinearSystem,u0map,
                                           parammap=DiffEqBase.NullParameters();
                                           jac = false, sparse=false,
                                           checkbounds = false,

src/systems/reaction/reactionsystem.jl

@@ -188,25 +188,34 @@ function oderatelaw(rx; combinatoric_ratelaw=true)
     rl
 end
 
-function assemble_drift(rs; combinatoric_ratelaws=true)
-    D   = Differential(rs.iv)
-    eqs = [D(x) ~ 0 for x in rs.states]
+function assemble_oderhs(rs; combinatoric_ratelaws=true)
     species_to_idx = Dict((x => i for (i,x) in enumerate(rs.states)))
-
+    rhsvec         = [Num(0) for i in eachindex(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)
-                signedrl = (stoich > zero(stoich)) ? rl : -rl
-                rhs      = isone(abs(stoich)) ? signedrl : stoich * rl
+            if _iszero(rhsvec[i])
+                signedrl  = (stoich > zero(stoich)) ? rl : -rl
+                rhsvec[i] = 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)
+                Δspec     = isone(abs(stoich)) ? rl : abs(stoich) * rl
+                rhsvec[i] = (stoich > zero(stoich)) ? (rhsvec[i] + Δspec) : (rhsvec[i] - Δspec)
             end
-            eqs[i] = Equation(eqs[i].lhs, rhs)
         end
     end
+
+    rhsvec
+end
+
+function assemble_drift(rs; combinatoric_ratelaws=true, as_odes=true)
+    rhsvec = assemble_oderhs(rs; combinatoric_ratelaws=combinatoric_ratelaws)
+    if as_odes
+        D   = Differential(rs.iv)
+        eqs = [Equation(D(x),rhs) for (x,rhs) in zip(rs.states,rhsvec)]
+    else
+        eqs = [Equation(Num(0),rhs) for rhs in rhsvec]
+    end
     eqs
 end
 
@@ -383,6 +392,24 @@ function Base.convert(::Type{<:ODESystem}, rs::ReactionSystem; combinatoric_rate
               systems=convert.(ODESystem,rs.systems))
 end
 
+"""
+```julia
+Base.convert(::Type{<:NonlinearSystem},rs::ReactionSystem)
+```
+
+Convert a [`ReactionSystem`](@ref) to an [`NonlinearSystem`](@ref).
+
+Notes:
+- `combinatoric_ratelaws=true` uses factorial scaling factors in calculating the rate
+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,rs.states,rs.ps,name=rs.name,systems=convert.(NonlinearSystem,rs.systems))
+end
+
 """
 ```julia
 Base.convert(::Type{<:SDESystem},rs::ReactionSystem)
@@ -450,55 +477,29 @@ function Base.convert(::Type{<:JumpSystem},rs::ReactionSystem; combinatoric_rate
 end
 
 
-"""
-```julia
-Base.convert(::Type{<:NonlinearSystem},rs::ReactionSystem)
-```
-
-Convert a [`ReactionSystem`](@ref) to an [`NonlinearSystem`](@ref).
-
-Notes:
-- `combinatoric_ratelaws=true` uses factorial scaling factors in calculating the rate
-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)
-    states_swaps = value.(rs.states)
-    eqs = map(eq -> 0 ~ make_sub!(eq,states_swaps),getproperty.(assemble_drift(rs; combinatoric_ratelaws=combinatoric_ratelaws),:rhs))
-    NonlinearSystem(eqs,rs.states,rs.ps,name=rs.name,
-              systems=convert.(NonlinearSystem,rs.systems))
-end
-
-# Used for Base.convert(::Type{<:NonlinearSystem},rs::ReactionSystem) only, should likely be removed.
-function make_sub!(eq,states_swaps)
-	for (i,arg) in enumerate(eq.args)
-		if any(isequal.(states_swaps,arg))
-			eq.args[i] = var2op(arg.op)
-		else
-			make_sub!(arg,states_swaps)
-		end
-	end
-	return eq
-end
-
 ### Converts a reaction system to ODE or SDE problems ###
 
 
 # ODEProblem from AbstractReactionNetwork
-function DiffEqBase.ODEProblem(rs::ReactionSystem, u0::Union{AbstractArray, Number}, tspan, p=DiffEqBase.NullParameters(), args...; kwargs...)
+function DiffEqBase.ODEProblem(rs::ReactionSystem, u0, tspan, p=DiffEqBase.NullParameters(), args...; kwargs...)
     return ODEProblem(convert(ODESystem,rs),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...)
+end
+
+
 # SDEProblem from AbstractReactionNetwork
-function DiffEqBase.SDEProblem(rs::ReactionSystem, u0::Union{AbstractArray, Number}, tspan, p=DiffEqBase.NullParameters(), args...; noise_scaling=nothing, kwargs...)
+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(rs.states), length(rs.eqs))
     return SDEProblem(sde_sys,u0,tspan,p,args...; noise_rate_prototype=p_matrix,kwargs...)
 end
 
 # DiscreteProblem from AbstractReactionNetwork
-function DiffEqBase.DiscreteProblem(rs::ReactionSystem, u0::Union{AbstractArray, Number}, tspan::Tuple, p=DiffEqBase.NullParameters(), args...; kwargs...)
+function DiffEqBase.DiscreteProblem(rs::ReactionSystem, u0, tspan::Tuple, p=DiffEqBase.NullParameters(), args...; kwargs...)
     return DiscreteProblem(convert(JumpSystem,rs), u0,tspan,p, args...; kwargs...)
 end
 
@@ -508,7 +509,7 @@ function DiffEqJump.JumpProblem(rs::ReactionSystem, prob, aggregator, args...; k
 end
 
 # SteadyStateProblem from AbstractReactionNetwork
-function DiffEqBase.SteadyStateProblem(rs::ReactionSystem, u0::Union{AbstractArray, Number}, p=DiffEqBase.NullParameters(), args...; kwargs...)
+function DiffEqBase.SteadyStateProblem(rs::ReactionSystem, u0, p=DiffEqBase.NullParameters(), args...; kwargs...)
     return SteadyStateProblem(ODEFunction(convert(ODESystem,rs)),u0,p, args...; kwargs...)
 end
 

test/reactionsystem.jl

@@ -81,6 +81,13 @@ du2 = sf.f(u,p,t)
 G2 = sf.g(u,p,t)
 @test norm(G-G2) < 100*eps()
 
+# test conversion to NonlinearSystem
+ns = convert(NonlinearSystem,rs)
+fnl = eval(generate_function(ns)[2])
+dunl = similar(du)
+fnl(dunl,u,p)
+@test norm(du-dunl) < 100*eps()
+
 # tests the noise_scaling argument.
 p  = rand(length(k)+1)
 u  = rand(length(k))