Merge pull request #713 from SciML/custom_function_expansion

Create function for expanding e.g. `mm` expressions, and use for e.g. HomotopyContinuation

Author: TorkelE

Project.toml

@@ -17,6 +17,7 @@ Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 RuntimeGeneratedFunctions = "7e49a35a-f44a-4d26-94aa-eba1b4ca6b47"
+Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
@@ -42,6 +43,7 @@ Parameters = "0.12"
 Reexport = "0.2, 1.0"
 Requires = "1.0"
 RuntimeGeneratedFunctions = "0.5.12"
+Setfield = "1"
 SymbolicUtils = "1.0.3"
 Symbolics = "5.0.3"
 Unitful = "1.12.4"

ext/CatalystHomotopyContinuationExtension.jl

@@ -4,6 +4,7 @@ module CatalystHomotopyContinuationExtension
 using Catalyst
 import ModelingToolkit as MT
 import HomotopyContinuation as HC
+import Setfield: @set
 import Symbolics: unwrap, wrap, Rewriters, symtype, issym, istree
 
 # Creates and exports hc_steady_states function.

ext/CatalystHomotopyContinuationExtension/homotopy_continuation_extension.jl

@@ -32,6 +32,7 @@ gives
 ```
 
 Notes:
+- Homotopy-based steady state finding only works when all rates are rational polynomials (e.g. constant, linear, mm, or hill functions).
 ```
   """
 function Catalyst.hc_steady_states(rs::ReactionSystem, ps; filter_negative=true, neg_thres=-1e-20, u0=[], kwargs...)
@@ -43,14 +44,14 @@ end
 
 # For a given reaction system, parameter values, and initial conditions, find the polynomial that HC solves to find steady states.
 function steady_state_polynomial(rs::ReactionSystem, ps, u0)
+    rs = Catalyst.expand_registered_functions(rs)
     ns = convert(NonlinearSystem, rs; remove_conserved = true)
     pre_varmap = [symmap_to_varmap(rs,u0)..., symmap_to_varmap(rs,ps)...]
     conservationlaw_errorcheck(rs, pre_varmap)
     p_vals = ModelingToolkit.varmap_to_vars(pre_varmap, parameters(ns); defaults = ModelingToolkit.defaults(ns))
     p_dict  = Dict(parameters(ns) .=> p_vals)
     eqs_pars_funcs = vcat(equations(ns), conservedequations(rs))
-    eqs_funcs = map(eq -> substitute(eq.rhs - eq.lhs, p_dict), eqs_pars_funcs)
-    eqs = [deregister([mm, mmr, hill, hillr, hillar], eq) for eq in eqs_funcs]
+    eqs = map(eq -> substitute(eq.rhs - eq.lhs, p_dict), eqs_pars_funcs)
     eqs_intexp = make_int_exps.(eqs)
     return Catalyst.to_multivariate_poly(remove_denominators.(eqs_intexp))
 end
@@ -63,27 +64,6 @@ function conservationlaw_errorcheck(rs, pre_varmap)
         error("The system has conservation laws but initial conditions were not provided for some species.")
 end
 
-# Unfolds a function (like mm or hill). 
-function deregister(fs::Vector{T}, expr) where T
-    for f in fs
-        expr = deregister(f, expr) 
-    end
-    return expr
-end
-# Provided by Shashi Gowda.
-deregister(f, expr) = wrap(Rewriters.Postwalk(Rewriters.PassThrough(___deregister(f)))(unwrap(expr)))
-function ___deregister(f)
-    (expr) ->
-    if istree(expr) && operation(expr) == f
-        args = arguments(expr)
-        invoke_with = map(args) do a
-            t = typeof(a)
-            issym(a) || istree(a) ? wrap(a) => symtype(a) : a => typeof(a)
-        end 
-        invoke(f, Tuple{last.(invoke_with)...}, first.(invoke_with)...)
-    end
-end
-
 # Parses and expression and return a version where any exponents that are Float64 (but an int, like 2.0) are turned into Int64s.
 make_int_exps(expr) = wrap(Rewriters.Postwalk(Rewriters.PassThrough(___make_int_exps))(unwrap(expr))).val
 function ___make_int_exps(expr)

src/Catalyst.jl

@@ -4,7 +4,7 @@ $(DocStringExtensions.README)
 module Catalyst
 
 using DocStringExtensions
-using SparseArrays, DiffEqBase, Reexport
+using SparseArrays, DiffEqBase, Reexport, Setfield
 using LaTeXStrings, Latexify, Requires
 using JumpProcesses: JumpProcesses,
                      JumpProblem, MassActionJump, ConstantRateJump,

src/reactionsystem.jl

@@ -1335,6 +1335,7 @@ function Base.convert(::Type{<:NonlinearSystem}, rs::ReactionSystem; name = name
                          as_odes = false, include_zero_odes)
     error_if_constraint_odes(NonlinearSystem, fullrs)
     eqs, sts, ps, obs, defs = addconstraints!(eqs, fullrs, ists, ispcs; remove_conserved)
+
     NonlinearSystem(eqs, sts, ps;
                     name,
                     observed = obs,

src/registered_functions.jl

@@ -107,3 +107,41 @@ function Symbolics.derivative(::typeof(hillar), args::NTuple{5, Any}, ::Val{5})
      (args[4]^args[5]) * log(args[4])) /
     (args[1]^args[5] + args[2]^args[5] + args[4]^args[5])^2
 end
+
+"""
+expand_registered_functions(expr)
+
+Takes an expression, and expands registered function expressions. E.g. `mm(X,v,K)` is replaced with v*X/(X+K). Currently supported functions: `mm`, `mmr`, `hill`, `hillr`, and `hill`.
+"""
+function expand_registered_functions(expr)
+    istree(expr) || return expr
+    args = arguments(expr)
+    if operation(expr) == Catalyst.mm
+        return args[2]*args[1]/(args[1] + args[3])
+    elseif operation(expr) == Catalyst.mmr
+        return args[2]*args[3]/(args[1] + args[3])
+    elseif operation(expr) == Catalyst.hill
+        return args[2]*(args[1]^args[4])/((args[1])^args[4] + (args[3])^args[4])
+    elseif operation(expr) == Catalyst.hillr
+        return args[2]*(args[3]^args[4])/((args[1])^args[4] + (args[3])^args[4])
+    elseif operation(expr) == Catalyst.hillar
+        return args[3]*(args[1]^args[5])/((args[1])^args[5] + (args[2])^args[5] + (args[4])^args[5])
+    end
+    for i = 1:length(args)
+        args[i] = expand_registered_functions(args[i])
+    end
+    return expr
+end
+# If applied to a Reaction, return a reaction with its rate modified.
+function expand_registered_functions(rx::Reaction)
+    Reaction(expand_registered_functions(rx.rate), rx.substrates, rx.products, rx.substoich, rx.prodstoich, rx.netstoich, rx.only_use_rate)
+end
+# If applied to a Equation, returns it with it applied to lhs and rhs
+function expand_registered_functions(eq::Equation)
+    return expand_registered_functions(eq.lhs) ~ expand_registered_functions(eq.rhs)
+end
+# If applied to a ReactionSystem, applied function to all Reactions and other Equations, and return updated system.
+function expand_registered_functions(rs::ReactionSystem)
+    rs = @set rs.eqs = [Catalyst.expand_registered_functions(eq) for eq in rs.eqs]
+    return @set rs.rxs = [Catalyst.expand_registered_functions(rx) for rx in rs.rxs]
+end

test/dsl/custom_functions.jl

@@ -6,7 +6,7 @@ using ModelingToolkit: get_states, get_ps
 using StableRNGs
 rng = StableRNG(12345)
 
-### Tests Cutom Functions ###
+### Tests Custom Functions ###
 let
     new_hill(x, v, k, n) = v * x^n / (k^n + x^n)
     new_poly(x, p1, p2) = 0.5 * p1 * x^2
@@ -204,3 +204,65 @@ let
                   v * (X^n) * ((K^n + Y^n) * log(X) - (K^n) * log(K) - (Y^n) * log(Y)) /
                   (K^n + X^n + Y^n)^2)
 end
+
+### Tests Current Function Expansion ###
+
+# Tests `ReactionSystem`s.
+let
+    @variables t 
+    @species x(t) y(t)
+    @parameters k v n 
+    rs1 = @reaction_network rs begin
+        mm(x, v, k), 0 --> x
+        mmr(x, v, k), 0 --> x
+        hill(x, v, k, n), 0 --> x
+        hillr(x, v, k, n), 0 --> x
+        hillar(x, y, v, k, n), 0 --> x    
+    end
+    rs2 = @reaction_network rs begin
+        v * x / (x + k), 0 --> x
+        v * k / (x + k), 0 --> x
+        v * (x^n) / (x^n + k^n), 0 --> x
+        v * (k^n) / (x^n + k^n), 0 --> x
+        v * (x^n) / (x^n + y^n + k^n), 0 --> x    
+    end
+
+    @test Catalyst.expand_registered_functions(rs1) == rs2
+end
+
+# Tests `Reaction`s.
+let
+    @variables t 
+    @species x(t) y(t)
+    @parameters k v n 
+    
+    r1 = @reaction mm(x, v, k), 0 --> x
+    r2 = @reaction mmr(x, v, k), 0 --> x
+    r3 = @reaction hill(x, v, k, n), 0 --> x
+    r4 = @reaction hillr(x, v, k, n), 0 --> x
+    r5 = @reaction hillar(x, y, v, k, n), 0 --> x + y
+    
+    @test isequal(Catalyst.expand_registered_functions(r1).rate, v * x / (x + k))
+    @test isequal(Catalyst.expand_registered_functions(r2).rate, v * k / (x + k))
+    @test isequal(Catalyst.expand_registered_functions(r3).rate, v * (x^n) / (x^n + k^n))
+    @test isequal(Catalyst.expand_registered_functions(r4).rate, v * (k^n) / (x^n + k^n))
+    @test isequal(Catalyst.expand_registered_functions(r5).rate, v * (x^n) / (x^n + y^n + k^n))
+end
+
+# Tests `Equation`s.
+let
+    @variables T X(T) Y(T)
+    @parameters K V N 
+    
+    eq1 = 0 ~ mm(X, V, K)
+    eq2 = 0 ~ mmr(X, V, K)
+    eq3 = 0 ~ hill(X, V, K, N)
+    eq4 = 0 ~ hillr(X, V, K, N)
+    eq5 = 0 ~ hillar(X, Y, V, K, N)
+    
+    @test isequal(Catalyst.expand_registered_functions(eq1), 0 ~ V * X / (X + K))
+    @test isequal(Catalyst.expand_registered_functions(eq2), 0 ~ V * K / (X + K))
+    @test isequal(Catalyst.expand_registered_functions(eq3), 0 ~ V * (X^N) / (X^N + K^N))
+    @test isequal(Catalyst.expand_registered_functions(eq4), 0 ~ V * (K^N) / (X^N + K^N))
+    @test isequal(Catalyst.expand_registered_functions(eq5), 0 ~ V * (X^N) / (X^N + Y^N + K^N))
+end

test/extensions/homotopy_continuation.jl

@@ -79,11 +79,12 @@ end
 
 # Tests that non-scalar reaction rates work.
 # Tests that rational polynomial steady state systems work. 
+# Tests that Hill function is correctly expanded even if nested.
 # Test filter_negative=false works.
-# tests than non-integer exponents throws an error.
+# Tests than non-integer exponents throws an error.
 let 
     rs = @reaction_network begin
-        0.1 + hill(X,v,K,n), 0 --> X
+        v*(0.1/v + hill(X,1,K,n)), 0 --> X
         d, X --> 0
     end
     ps = [:v => 5.0, :K => 2.5, :n => 3, :d => 1.0]
@@ -92,7 +93,7 @@ let
     f = ODEFunction(convert(ODESystem, rs))
     @test length(sss) == 4
     for ss in sss
-        @test f(sss[1], last.(ps), 0.0)[1] == 0.0
+        @test isapprox(f(sss[1], last.(ps), 0.0)[1], 0.0; atol=1e-12)
     end
 
     @test_throws Exception hc_steady_states(rs, [:v => 5.0, :K => 2.5, :n => 2.7, :d => 1.0]; show_progress=false)