Merge pull request #1224 from isaacsas/expand_hills

Expand Catalyst functions when converting

Author: isaacsas

HISTORY.md

@@ -123,7 +123,27 @@
   plot_network(brusselator)
   ```
 - The letter Ø (used in Danish/Norwegian alphabet) is now conisdred the same as ∅ (empty set). It can no longer be used as a species/parameter.
- 
+- When converting a Catalyst `ReactionSystem` to a ModelingToolkit system, for
+  example an `ODESystem`, Catalyst defined functions like `hill(A,B,C,D)` are
+  now replaced with the explicit rational function they represent in the
+  equations of the generated system. For example `mm(X,v,K)` will be replaced
+  with `v*X / (X + K)`. This can be disabled by passing the keyword argument
+  `expand_catalyst_funs = false`. e.g.
+  ```julia
+  using Catalyst
+  rn = @reaction_network begin
+    hill(X,v,K,n), A --> 0
+  end
+  osys = convert(ODESystem, rn)
+  ```
+  generates an ODE system with `D(A) ~ ~ -((v*A(t)*(X^n)) / (K^n + X^n))`, while
+  ```julia
+  osys = convert(ODESystem, rn; expand_catalyst_funs = false)
+  ```
+  generates an ODE system with `D(A) ~ -A(t)*hill(X, v, K, n)`. This keyword
+  argument can also be passed to problems defined over `ReactionSystem`s, i.e.
+  when calling `ODEProblem(rn, u0, tspan, p; expand_catalyst_funs = false)`.
+
 ## Catalyst 14.4.1
 - Support for user-defined functions on the RHS when providing coupled equations 
   for CRNs using the @equations macro. For example, the following now works: 

src/reactionsystem_conversions.jl

@@ -22,9 +22,10 @@ Notes:
     `combinatoric_ratelaw=false` then the ratelaw is `k*S^2`, i.e. the scaling
     factor is ignored.
 """
-function oderatelaw(rx; combinatoric_ratelaw = true)
+function oderatelaw(rx; combinatoric_ratelaw = true, expand_catalyst_funs = true)
     @unpack rate, substrates, substoich, only_use_rate = rx
     rl = rate
+    expand_catalyst_funs && (rl = expand_registered_functions(rl))
 
     # if the stoichiometric coefficients are not integers error if asking to scale rates
     !all(s -> s isa Union{Integer, Symbolic}, substoich) &&
@@ -46,7 +47,8 @@ end
 # including non-species variables.
 drop_dynamics(s) = isconstant(s) || isbc(s) || (!isspecies(s))
 
-function assemble_oderhs(rs, ispcs; combinatoric_ratelaws = true, remove_conserved = false)
+function assemble_oderhs(rs, ispcs; combinatoric_ratelaws = true, 
+        remove_conserved = false, expand_catalyst_funs = true)
     nps = get_networkproperties(rs)
     species_to_idx = Dict(x => i for (i, x) in enumerate(ispcs))
     rhsvec = Any[0 for _ in ispcs]
@@ -57,7 +59,8 @@ function assemble_oderhs(rs, ispcs; combinatoric_ratelaws = true, remove_conserv
     end
 
     for rx in get_rxs(rs)
-        rl = oderatelaw(rx; combinatoric_ratelaw = combinatoric_ratelaws)
+        rl = oderatelaw(rx; combinatoric_ratelaw = combinatoric_ratelaws, 
+            expand_catalyst_funs)
         remove_conserved && (rl = substitute(rl, depspec_submap))
         for (spec, stoich) in rx.netstoich
             # dependent species don't get an ODE, so are skipped
@@ -90,8 +93,9 @@ function assemble_oderhs(rs, ispcs; combinatoric_ratelaws = true, remove_conserv
 end
 
 function assemble_drift(rs, ispcs; combinatoric_ratelaws = true, as_odes = true,
-        include_zero_odes = true, remove_conserved = false)
-    rhsvec = assemble_oderhs(rs, ispcs; combinatoric_ratelaws, remove_conserved)
+        include_zero_odes = true, remove_conserved = false, expand_catalyst_funs = true)
+    rhsvec = assemble_oderhs(rs, ispcs; combinatoric_ratelaws, remove_conserved, 
+        expand_catalyst_funs)
     if as_odes
         D = Differential(get_iv(rs))
         eqs = [Equation(D(x), rhs)
@@ -104,7 +108,7 @@ end
 
 # this doesn't work with constraint equations currently
 function assemble_diffusion(rs, sts, ispcs; combinatoric_ratelaws = true,
-        remove_conserved = false)
+        remove_conserved = false, expand_catalyst_funs = true)
     # as BC species should ultimately get an equation, we include them in the noise matrix
     num_bcsts = count(isbc, get_unknowns(rs))
 
@@ -120,7 +124,9 @@ function assemble_diffusion(rs, sts, ispcs; combinatoric_ratelaws = true,
     end
 
     for (j, rx) in enumerate(get_rxs(rs))
-        rlsqrt = sqrt(abs(oderatelaw(rx; combinatoric_ratelaw = combinatoric_ratelaws)))
+        rl = oderatelaw(rx; combinatoric_ratelaw = combinatoric_ratelaws, 
+            expand_catalyst_funs)
+        rlsqrt = sqrt(abs(rl))
         hasnoisescaling(rx) && (rlsqrt *= getnoisescaling(rx))
         remove_conserved && (rlsqrt = substitute(rlsqrt, depspec_submap))
 
@@ -169,9 +175,12 @@ Notes:
   the ratelaw is `k*S*(S-1)`, i.e. the rate law is not normalized by the scaling
   factor.
 """
-function jumpratelaw(rx; combinatoric_ratelaw = true)
+function jumpratelaw(rx; combinatoric_ratelaw = true, expand_catalyst_funs = true)
     @unpack rate, substrates, substoich, only_use_rate = rx
+    
     rl = rate
+    expand_catalyst_funs && (rl = expand_registered_functions(rl))
+
     if !only_use_rate
         coef = eltype(substoich) <: Number ? one(eltype(substoich)) : 1
         for (i, stoich) in enumerate(substoich)
@@ -313,7 +322,7 @@ function get_depgraph(rs)
     eqeq_dependencies(jdeps, vdeps).fadjlist
 end
 
-function assemble_jumps(rs; combinatoric_ratelaws = true)
+function assemble_jumps(rs; combinatoric_ratelaws = true, expand_catalyst_funs = true)
     meqs = MassActionJump[]
     ceqs = ConstantRateJump[]
     veqs = VariableRateJump[]
@@ -361,7 +370,8 @@ function assemble_jumps(rs; combinatoric_ratelaws = true)
         if (!isvrj) && ismassaction(rx, rs; rxvars, haveivdep = false, unknownset)
             push!(meqs, makemajump(rx; combinatoric_ratelaw = combinatoric_ratelaws))
         else
-            rl = jumpratelaw(rx; combinatoric_ratelaw = combinatoric_ratelaws)
+            rl = jumpratelaw(rx; combinatoric_ratelaw = combinatoric_ratelaws, 
+                expand_catalyst_funs)
             affect = Vector{Equation}()
             for (spec, stoich) in rx.netstoich
                 # don't change species that are constant or BCs
@@ -481,12 +491,15 @@ Keyword args and default values:
 - `remove_conserved=false`, if set to `true` will calculate conservation laws of the
   underlying set of reactions (ignoring constraint equations), and then apply them to reduce
   the number of equations.
+- `expand_catalyst_funs = true`, replaces Catalyst defined functions like `hill(A,B,C,D)`
+  with their rational function representation when converting to another system type. Set to
+  `false`` to disable.
 """
 function Base.convert(::Type{<:ODESystem}, rs::ReactionSystem; name = nameof(rs),
         combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
-        include_zero_odes = true, remove_conserved = false,
-        checks = false, default_u0 = Dict(), default_p = Dict(),
-        defaults = _merge(Dict(default_u0), Dict(default_p)),
+        include_zero_odes = true, remove_conserved = false, checks = false, 
+        default_u0 = Dict(), default_p = Dict(), 
+        defaults = _merge(Dict(default_u0), Dict(default_p)), expand_catalyst_funs = true, 
         kwargs...)
     # Error checks.
     iscomplete(rs) || error(COMPLETENESS_ERROR)
@@ -496,7 +509,7 @@ function Base.convert(::Type{<:ODESystem}, rs::ReactionSystem; name = nameof(rs)
     remove_conserved && conservationlaws(fullrs)
     ists, ispcs = get_indep_sts(fullrs, remove_conserved)
     eqs = assemble_drift(fullrs, ispcs; combinatoric_ratelaws, remove_conserved,
-        include_zero_odes)
+        include_zero_odes, expand_catalyst_funs)
     eqs, us, ps, obs, defs = addconstraints!(eqs, fullrs, ists, ispcs; remove_conserved)
 
     ODESystem(eqs, get_iv(fullrs), us, ps;
@@ -547,13 +560,16 @@ Keyword args and default values:
   conservation laws. See the [FAQ
   entry](https://docs.sciml.ai/Catalyst/stable/faqs/#faq_remake_nonlinprob) for more
   details.
+- `expand_catalyst_funs = true`, replaces Catalyst defined functions like `hill(A,B,C,D)`
+  with their rational function representation when converting to another system type. Set to
+  `false`` to disable.
 """
 function Base.convert(::Type{<:NonlinearSystem}, rs::ReactionSystem; name = nameof(rs),
         combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
         remove_conserved = false, conseqs_remake_warn = true, checks = false,
         default_u0 = Dict(), default_p = Dict(), 
         defaults = _merge(Dict(default_u0), Dict(default_p)),
-        all_differentials_permitted = false, kwargs...)
+        all_differentials_permitted = false, expand_catalyst_funs = true, kwargs...)
     # Error checks.
     iscomplete(rs) || error(COMPLETENESS_ERROR)
     spatial_convert_err(rs::ReactionSystem, NonlinearSystem)
@@ -565,7 +581,7 @@ function Base.convert(::Type{<:NonlinearSystem}, rs::ReactionSystem; name = name
     remove_conserved && conservationlaws(fullrs)
     ists, ispcs = get_indep_sts(fullrs, remove_conserved)
     eqs = assemble_drift(fullrs, ispcs; combinatoric_ratelaws, remove_conserved,
-        as_odes = false, include_zero_odes = false)
+        as_odes = false, include_zero_odes = false, expand_catalyst_funs)
     eqs, us, ps, obs, defs, initeqs = addconstraints!(eqs, fullrs, ists, ispcs; 
         remove_conserved, treat_conserved_as_eqs = true)
 
@@ -627,11 +643,16 @@ Notes:
 - `remove_conserved=false`, if set to `true` will calculate conservation laws of the
   underlying set of reactions (ignoring constraint equations), and then apply them to reduce
   the number of equations.
+- `expand_catalyst_funs = true`, replaces Catalyst defined functions like `hill(A,B,C,D)`
+  with their rational function representation when converting to another system type. Set to
+  `false`` to disable.
 """
 function Base.convert(::Type{<:SDESystem}, rs::ReactionSystem;
         name = nameof(rs), combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
         include_zero_odes = true, checks = false, remove_conserved = false,
-        default_u0 = Dict(), default_p = Dict(), defaults = _merge(Dict(default_u0), Dict(default_p)),
+        default_u0 = Dict(), default_p = Dict(), 
+        defaults = _merge(Dict(default_u0), Dict(default_p)),
+        expand_catalyst_funs = true, 
         kwargs...)
     # Error checks.
     iscomplete(rs) || error(COMPLETENESS_ERROR)
@@ -642,9 +663,9 @@ function Base.convert(::Type{<:SDESystem}, rs::ReactionSystem;
     remove_conserved && conservationlaws(flatrs)
     ists, ispcs = get_indep_sts(flatrs, remove_conserved)
     eqs = assemble_drift(flatrs, ispcs; combinatoric_ratelaws, include_zero_odes,
-        remove_conserved)
-    noiseeqs = assemble_diffusion(flatrs, ists, ispcs;
-        combinatoric_ratelaws, remove_conserved)
+        remove_conserved, expand_catalyst_funs)
+    noiseeqs = assemble_diffusion(flatrs, ists, ispcs; combinatoric_ratelaws, 
+        remove_conserved, expand_catalyst_funs)
     eqs, us, ps, obs, defs = addconstraints!(eqs, flatrs, ists, ispcs; remove_conserved)
 
     if any(isbc, get_unknowns(flatrs))
@@ -678,12 +699,14 @@ Notes:
   differential equations.
 - Does not currently support continuous events as these are not supported by
   `ModelingToolkit.JumpSystems`.
+- `expand_catalyst_funs = true`, replaces Catalyst defined functions like `hill(A,B,C,D)`
+  with their rational function representation when converting to another system type. Set to
+  `false`` to disable.
 """
 function Base.convert(::Type{<:JumpSystem}, rs::ReactionSystem; name = nameof(rs),
         combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
-        remove_conserved = nothing, checks = false,
-        default_u0 = Dict(), default_p = Dict(),
-        defaults = _merge(Dict(default_u0), Dict(default_p)),
+        remove_conserved = nothing, checks = false, default_u0 = Dict(), default_p = Dict(),
+        defaults = _merge(Dict(default_u0), Dict(default_p)), expand_catalyst_funs = true, 
         kwargs...)
     iscomplete(rs) || error(COMPLETENESS_ERROR)
     spatial_convert_err(rs::ReactionSystem, JumpSystem)
@@ -696,7 +719,7 @@ function Base.convert(::Type{<:JumpSystem}, rs::ReactionSystem; name = nameof(rs
     (length(MT.continuous_events(flatrs)) > 0) &&
         (@warn "continuous_events will be dropped as they are not currently supported by JumpSystems.")
 
-    eqs = assemble_jumps(flatrs; combinatoric_ratelaws)
+    eqs = assemble_jumps(flatrs; combinatoric_ratelaws, expand_catalyst_funs)
 
     # handle BC species
     sts, ispcs = get_indep_sts(flatrs)
@@ -720,9 +743,9 @@ function DiffEqBase.ODEProblem(rs::ReactionSystem, u0, tspan,
         check_length = false, name = nameof(rs),
         combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
         include_zero_odes = true, remove_conserved = false, checks = false, 
-        structural_simplify = false, kwargs...)
+        expand_catalyst_funs = true, structural_simplify = false, kwargs...)
     osys = convert(ODESystem, rs; name, combinatoric_ratelaws, include_zero_odes, checks,
-        remove_conserved)
+        remove_conserved, expand_catalyst_funs)
 
     # Handles potential differential algebraic equations (which requires `structural_simplify`).
     if structural_simplify
@@ -764,14 +787,19 @@ Keyword args and default values:
   conservation laws. See the [FAQ
   entry](https://docs.sciml.ai/Catalyst/stable/faqs/#faq_remake_nonlinprob) for more
   details.
+- `expand_catalyst_funs = true`, replaces Catalyst defined functions like `hill(A,B,C,D)`
+  with their rational function representation when converting to another system type. Set to
+  `false`` to disable.
 """
 function DiffEqBase.NonlinearProblem(rs::ReactionSystem, u0,
         p = DiffEqBase.NullParameters(), args...;
         name = nameof(rs), combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
-        remove_conserved = false, conseqs_remake_warn = true, checks = false, check_length = false, 
+        remove_conserved = false, conseqs_remake_warn = true, checks = false, 
+        check_length = false, expand_catalyst_funs = true, 
         structural_simplify = false, all_differentials_permitted = false, kwargs...)
     nlsys = convert(NonlinearSystem, rs; name, combinatoric_ratelaws, checks, 
-        all_differentials_permitted, remove_conserved, conseqs_remake_warn)
+        all_differentials_permitted, remove_conserved, conseqs_remake_warn, 
+        expand_catalyst_funs)
     nlsys = structural_simplify ? MT.structural_simplify(nlsys) : complete(nlsys)
     return NonlinearProblem(nlsys, u0, p, args...; check_length,
         kwargs...)
@@ -781,9 +809,10 @@ end
 function DiffEqBase.SDEProblem(rs::ReactionSystem, u0, tspan,
         p = DiffEqBase.NullParameters(), args...;
         name = nameof(rs), combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
-        include_zero_odes = true, checks = false, check_length = false, remove_conserved = false,
-        structural_simplify = false, kwargs...)
-    sde_sys = convert(SDESystem, rs; name, combinatoric_ratelaws,
+        include_zero_odes = true, checks = false, check_length = false, 
+        remove_conserved = false, structural_simplify = false, 
+        expand_catalyst_funs = true, kwargs...)
+    sde_sys = convert(SDESystem, rs; name, combinatoric_ratelaws, expand_catalyst_funs, 
         include_zero_odes, checks, remove_conserved)
 
     # Handles potential differential algebraic equations (which requires `structural_simplify`).
@@ -848,8 +877,9 @@ plot(sol, idxs = :A)
 """
 function JumpInputs(rs::ReactionSystem, u0, tspan, p = DiffEqBase.NullParameters();
         name = nameof(rs), combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
-        checks = false, kwargs...)
-    jsys = complete(convert(JumpSystem, rs; name, combinatoric_ratelaws, checks))
+        expand_catalyst_funs = true, checks = false, kwargs...)
+    jsys = complete(convert(JumpSystem, rs; name, combinatoric_ratelaws, 
+        expand_catalyst_funs, checks))
     if MT.has_variableratejumps(jsys)
         prob = ODEProblem(jsys, u0, tspan, p; kwargs...)
     else
@@ -873,11 +903,11 @@ end
 
 # DiscreteProblem from AbstractReactionNetwork
 function DiffEqBase.DiscreteProblem(rs::ReactionSystem, u0, tspan::Tuple,
-        p = DiffEqBase.NullParameters(), args...;
-        name = nameof(rs),
-        combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
-        checks = false, kwargs...)
-    jsys = convert(JumpSystem, rs; name, combinatoric_ratelaws, checks)
+        p = DiffEqBase.NullParameters(), args...; name = nameof(rs), 
+        combinatoric_ratelaws = get_combinatoric_ratelaws(rs), checks = false, 
+        expand_catalyst_funs = true, kwargs...)
+    jsys = convert(JumpSystem, rs; name, combinatoric_ratelaws, checks, 
+        expand_catalyst_funs)
     jsys = complete(jsys)
     return DiscreteProblem(jsys, u0, tspan, p, args...; kwargs...)
 end
@@ -886,8 +916,9 @@ end
 function JumpProcesses.JumpProblem(rs::ReactionSystem, prob,
         aggregator = JumpProcesses.NullAggregator(); name = nameof(rs),
         combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
-        checks = false, kwargs...)
-    jsys = convert(JumpSystem, rs; name, combinatoric_ratelaws, checks)
+        expand_catalyst_funs = true, checks = false, kwargs...)
+    jsys = convert(JumpSystem, rs; name, combinatoric_ratelaws, 
+        expand_catalyst_funs, checks)
     jsys = complete(jsys)
     return JumpProblem(jsys, prob, aggregator; kwargs...)
 end
@@ -905,9 +936,9 @@ function DiffEqBase.SteadyStateProblem(rs::ReactionSystem, u0,
         check_length = false, name = nameof(rs),
         combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
         remove_conserved = false, include_zero_odes = true, checks = false, 
-        structural_simplify = false, kwargs...)
+        expand_catalyst_funs = true, structural_simplify = false, kwargs...)
     osys = convert(ODESystem, rs; name, combinatoric_ratelaws, include_zero_odes, checks,
-        remove_conserved)
+        remove_conserved, expand_catalyst_funs)
 
     # Handles potential differential algebraic equations (which requires `structural_simplify`).
     if structural_simplify