Merge pull request #707 from SciML/BifurcationKitExtension

TorkelE · web-flow · commit 847077b2c9f3 · 2023-12-03T19:55:45.000-05:00
Bifurcation kit extension
diff --git a/HISTORY.md b/HISTORY.md
@@ -24,6 +24,35 @@
   ```
   X's value will be `1.0` in the first vertex, but `0.0` in the remaining one (the system have 25 vertexes in total). SInce th parameters `p` and `d` are part of the non-spatial reaction network, their values are tied to vertexes. However, if the `D` parameter (which governs diffusion between vertexes) is given several values, these will instead correspond to the specific edges (and transportation along those edges.)
 
+- Add a CatalystBifurcationKitExtension, permitting BifurcationKit's `BifurcationProblem`s to be created from Catalyst reaction networks. Example usage:
+```julia
+using Catalyst
+wilhelm_2009_model = @reaction_network begin
+    k1, Y --> 2X
+    k2, 2X --> X + Y
+    k3, X + Y --> Y
+    k4, X --> 0
+    k5, 0 --> X
+end
+
+
+using BifurcationKit
+bif_par = :k1
+u_guess = [:X => 5.0, :Y => 2.0]
+p_start = [:k1 => 4.0, :k2 => 1.0, :k3 => 1.0, :k4 => 1.5, :k5 => 1.25]
+plot_var = :X
+bprob = BifurcationProblem(wilhelm_2009_model, u_guess, p_start, bif_par; plot_var=plot_var)
+
+p_span = (2.0, 20.0)
+opts_br = ContinuationPar(p_min = p_span[1], p_max = p_span[2], max_steps=1000)
+
+bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside=true)
+
+using Plots
+plot(bif_dia; xguide="k1", yguide="X")
+```
+- Automatically handles elimination of conservation laws for computing bifurcation diagrams.
+- Updated Bifurcation documentation with respect to this new feature.
 
 ## Catalyst 13.5
 - Added a CatalystHomotopyContinuationExtension extension, which exports the `hc_steady_state` function if HomotopyContinuation is exported. `hc_steady_state` finds the steady states of a reaction system using the homotopy continuation method. This feature is only available for julia versions 1.9+. Example: 
diff --git a/Project.toml b/Project.toml
@@ -24,21 +24,25 @@ Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [weakdeps]
+BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
 
 [extensions]
+CatalystBifurcationKitExtension = "BifurcationKit"
 CatalystHomotopyContinuationExtension = "HomotopyContinuation"
 
 [compat]
+BifurcationKit = "0.3"
 DataStructures = "0.18"
 DiffEqBase = "6.83.0"
 DocStringExtensions = "0.8, 0.9"
 Graphs = "1.4"
 JumpProcesses = "9.3.2"
+HomotopyContinuation = "2.9"
 LaTeXStrings = "1.3.0"
 Latexify = "0.14, 0.15, 0.16"
 MacroTools = "0.5.5"
-ModelingToolkit = "8.66"
+ModelingToolkit = "8.73"
 Parameters = "0.12"
 Reexport = "0.2, 1.0"
 Requires = "1.0"
@@ -50,6 +54,7 @@ Unitful = "1.12.4"
 julia = "1.9"
 
 [extras]
+BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
 DomainSets = "5b8099bc-c8ec-5219-889f-1d9e522a28bf"
 Graphviz_jll = "3c863552-8265-54e4-a6dc-903eb78fde85"
 HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
@@ -68,4 +73,4 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
 
 [targets]
-test = ["DomainSets", "Graphviz_jll", "HomotopyContinuation", "NonlinearSolve", "OrdinaryDiffEq", "Random", "SafeTestsets", "SciMLBase", "SciMLNLSolve", "StableRNGs", "Statistics", "SteadyStateDiffEq", "StochasticDiffEq", "Test", "Unitful"]
+test = ["BifurcationKit", "DomainSets", "Graphviz_jll", "HomotopyContinuation", "NonlinearSolve", "OrdinaryDiffEq", "Random", "SafeTestsets", "SciMLBase", "SciMLNLSolve", "StableRNGs", "Statistics", "SteadyStateDiffEq", "StochasticDiffEq", "Test", "Unitful"]
diff --git a/docs/src/api/catalyst_api.md b/docs/src/api/catalyst_api.md
@@ -180,7 +180,7 @@ netstoichmat
 reactionrates
 ```
 
-## Functions to extend or modify a network
+## (Functions to extend or modify a network)(@id api_network_extension_and_modification)
 `ReactionSystem`s can be programmatically extended using
 [`@add_reactions`](@ref), [`addspecies!`](@ref), [`addparam!`](@ref) and
 [`addreaction!`](@ref), or using [`ModelingToolkit.extend`](@ref) and
diff --git a/docs/src/catalyst_applications/bifurcation_diagrams.md b/docs/src/catalyst_applications/bifurcation_diagrams.md
diff --git a/docs/src/catalyst_applications/homotopy_continuation.md b/docs/src/catalyst_applications/homotopy_continuation.md
@@ -38,7 +38,6 @@ The order of the species in the output vectors are the same as in `species(wilhe
 
 It should be noted that the steady state multivariate polynomials produced by reaction systems may have both imaginary and negative roots, which are filtered away by `hc_steady_states`. If you want the negative roots, you can use the `hc_steady_states(wilhelm_2009_model, ps; filter_negative=false)` argument.
 
-
 ## [Systems with conservation laws](@id homotopy_continuation_conservation_laws)
 Some systems are under-determined, and have an infinite number of possible steady states. These are typically systems containing a conservation
 law, e.g.
diff --git a/ext/CatalystBifurcationKitExtension.jl b/ext/CatalystBifurcationKitExtension.jl
@@ -0,0 +1,10 @@
+module CatalystBifurcationKitExtension
+
+# Fetch packages.
+using Catalyst
+import BifurcationKit as BK
+
+# Creates and exports hc_steady_states function.
+include("CatalystBifurcationKitExtension/bifurcation_kit_extension.jl")
+
+end
diff --git a/ext/CatalystBifurcationKitExtension/bifurcation_kit_extension.jl b/ext/CatalystBifurcationKitExtension/bifurcation_kit_extension.jl
@@ -0,0 +1,21 @@
+### Dispatch for BifurcationKit BifurcationProblems ###
+
+# Creates a BifurcationProblem, using a ReactionSystem as an input.
+function BK.BifurcationProblem(rs::ReactionSystem, u0_bif, ps, bif_par, args...; 
+                               plot_var=nothing, record_from_solution=BK.record_sol_default, jac=true, u0=[], kwargs...)  
+
+    # Converts symbols to symbolics.
+    (bif_par isa Symbol) && (bif_par = rs.var_to_name[bif_par])
+    (plot_var isa Symbol) && (plot_var = rs.var_to_name[plot_var])
+    ((u0_bif isa Vector{<:Pair{Symbol,<:Any}}) || (u0_bif isa Dict{Symbol, <:Any})) && (u0_bif = symmap_to_varmap(rs, u0_bif))
+    ((ps isa Vector{<:Pair{Symbol,<:Any}}) || (ps isa Dict{Symbol, <:Any})) && (ps = symmap_to_varmap(rs, ps))
+    ((u0 isa Vector{<:Pair{Symbol,<:Any}}) || (u0 isa Dict{Symbol, <:Any})) && (u0 = symmap_to_varmap(rs, u0))
+
+    # Creates NonlinearSystem.
+    Catalyst.conservationlaw_errorcheck(rs, vcat(ps, u0))
+    nsys = convert(NonlinearSystem, rs; remove_conserved=true, defaults=Dict(u0))
+
+    # Makes BifurcationProblem (this call goes through the ModelingToolkit-based BifurcationKit extension).
+    return BK.BifurcationProblem(nsys, u0_bif, ps, bif_par, args...; plot_var=plot_var, 
+                                 record_from_solution=record_from_solution, jac=jac, kwargs...)
+end
diff --git a/ext/CatalystHomotopyContinuationExtension/homotopy_continuation_extension.jl b/ext/CatalystHomotopyContinuationExtension/homotopy_continuation_extension.jl
@@ -47,7 +47,7 @@ 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)
+    Catalyst.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))
diff --git a/src/networkapi.jl b/src/networkapi.jl
@@ -1277,6 +1277,18 @@ Compute conserved quantities for a system with the given conservation laws.
 """
 conservedquantities(state, cons_laws) = cons_laws * state
 
+
+# If u0s are not given while conservation laws are present, throws an error.
+# Used in HomotopyContinuation and BifurcationKit extensions.
+# Currently only checks if any u0s are given
+# (not whether these are enough for computing conserved quantitites, this will yield a less informative error).
+function conservationlaw_errorcheck(rs, pre_varmap)
+    vars_with_vals = Set(p[1] for p in pre_varmap)
+    any(s -> s in vars_with_vals, species(rs)) && return
+    isempty(conservedequations(Catalyst.flatten(rs))) || 
+        error("The system has conservation laws but initial conditions were not provided for some species.")
+end
+
 ######################## reaction network operators #######################
 
 """
diff --git a/src/reactionsystem.jl b/src/reactionsystem.jl
@@ -1287,6 +1287,7 @@ Keyword args and default values:
 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)),
                       kwargs...)
     spatial_convert_err(rs::ReactionSystem, ODESystem)
     fullrs = Catalyst.flatten(rs)
@@ -1299,7 +1300,7 @@ function Base.convert(::Type{<:ODESystem}, rs::ReactionSystem; name = nameof(rs)
     ODESystem(eqs, get_iv(fullrs), sts, ps;
               observed = obs,
               name,
-              defaults = defs,
+              defaults = _merge(defaults,defs),
               checks,
               continuous_events = MT.get_continuous_events(fullrs),
               discrete_events = MT.get_discrete_events(fullrs),
@@ -1326,6 +1327,7 @@ Keyword args and default values:
 function Base.convert(::Type{<:NonlinearSystem}, 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)),
                       kwargs...)
     spatial_convert_err(rs::ReactionSystem, NonlinearSystem)
     fullrs = Catalyst.flatten(rs)
@@ -1339,7 +1341,7 @@ function Base.convert(::Type{<:NonlinearSystem}, rs::ReactionSystem; name = name
     NonlinearSystem(eqs, sts, ps;
                     name,
                     observed = obs,
-                    defaults = defs,
+                    defaults = _merge(defaults,defs),
                     checks,
                     kwargs...)
 end
@@ -1375,6 +1377,7 @@ function Base.convert(::Type{<:SDESystem}, rs::ReactionSystem;
                       noise_scaling = nothing, 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)),
                       kwargs...)
     spatial_convert_err(rs::ReactionSystem, SDESystem)
 
@@ -1435,7 +1438,9 @@ Notes:
 """
 function Base.convert(::Type{<:JumpSystem}, rs::ReactionSystem; name = nameof(rs),
                       combinatoric_ratelaws = get_combinatoric_ratelaws(rs),
-                      remove_conserved = nothing, checks = false, kwargs...)
+                      remove_conserved = nothing, checks = false, 
+                      default_u0 = Dict(), default_p = Dict(), defaults = _merge(Dict(default_u0), Dict(default_p)),
+                      kwargs...)
     spatial_convert_err(rs::ReactionSystem, JumpSystem)
 
     (remove_conserved !== nothing) &&
@@ -1457,7 +1462,7 @@ function Base.convert(::Type{<:JumpSystem}, rs::ReactionSystem; name = nameof(rs
     JumpSystem(eqs, get_iv(flatrs), sts, ps;
                observed = MT.observed(flatrs),
                name,
-               defaults = MT.defaults(flatrs),
+               defaults = _merge(defaults,MT.defaults(flatrs)),
                checks,
                discrete_events = MT.discrete_events(flatrs),
                kwargs...)
diff --git a/test/extensions/bifurcation_kit.jl b/test/extensions/bifurcation_kit.jl
@@ -0,0 +1,173 @@
+### Fetch Packages ###
+using BifurcationKit, Catalyst, Test
+
+# Sets rnd number.
+using StableRNGs
+rng = StableRNG(12345)
+
+### Run Tests ###
+
+# Brusselator extended with conserved species.
+# Runs full computation, checks values corresponds to known values.
+# Checks that teh correct bifurcation point is found at the correct position.
+# Checks that bifurcation diagrams can be computed for systems with conservation laws.
+# Checks that bifurcation diagrams can be computed for systems with default values.
+# Checks that bifurcation diagrams can be computed for systems with non-constant rate.
+# Checks that not providing conserved species throws and appropriate error.
+let 
+    # Create model.
+    extended_brusselator = @reaction_network begin
+        @species W(t) = 2.0
+        @parameters k2 = 0.5
+        A, ∅ → X
+        1, 2X + Y → 3X
+        B, X → Y
+        1, X → ∅
+        (k1*Y, k2), V <--> W
+    end
+    @unpack A, B, k1 = extended_brusselator
+    u0_guess = [:X => 1.0, :Y => 1.0, :V => 0.0, :W => 0.0]
+    p_start = [A => 1.0, B => 4.0, k1 => 0.1]
+    
+    # Computes bifurcation diagram.
+    bprob = BifurcationProblem(extended_brusselator, u0_guess, p_start, :B; plot_var = :V, u0 = [:V => 1.0])
+    p_span = (0.1, 6.0)
+    opts_br = ContinuationPar(dsmin = 0.0001, dsmax = 0.001, ds = 0.0001, max_steps = 10000, p_min = p_span[1], p_max = p_span[2], n_inversion = 4)
+    bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
+    
+    # Checks computed V values are correct (Formula: V = k2*(V0+W0)/(k1*Y+k2), where Y=2*B.)
+    B_vals = getfield.(bif_dia.γ.branch, :param)
+    V_vals = getfield.(bif_dia.γ.branch, :x)
+    @test all(V_vals .≈ 0.5*(1.0+2.0) ./ (0.1 .* 2*B_vals .+ 0.5))
+    
+    # Checks that the bifurcation point is correct.
+    @test length(bif_dia.γ.specialpoint) == 3 # Includes start and end point.
+    hopf_bif_point = filter(sp -> sp.type == :hopf, bif_dia.γ.specialpoint)[1]
+    @test isapprox(hopf_bif_point.param, 1.5, atol=1e-5)
+
+    # Tests that an error is thrown if information of conserved species is not fully provided.
+    @test_throws Exception BifurcationProblem(extended_brusselator, u0_guess, p_start, :B; plot_var=:V, u0 = [:X => 1.0])
+end
+
+# Bistable switch.
+# Checks that the same bifurcation problem is created as for BifurcationKit.
+# Checks with Symbolics as bifurcation and plot vars.
+# Tries setting `jac=false`.
+let 
+    # Creates BifurcationProblem via Catalyst.
+    bistable_switch = @reaction_network begin
+        0.1 + hill(X,v,K,n), 0 --> X
+        d, X --> 0
+    end
+    @unpack X, v, K, n, d = bistable_switch
+    u0_guess = [X => 1.0]
+    p_start = [v => 5.0, K => 2.5, n => 3, d => 1.0]
+    bprob = BifurcationProblem(bistable_switch, u0_guess, p_start, K; jac=false, plot_var=X)
+    
+    # Creates BifurcationProblem via BifurcationKit.
+    function bistable_switch_BK(u, p)
+        X, = u
+        v, K, n, d = p
+        return [0.1 + v*(X^n)/(X^n + K^n) - d*X]
+    end
+    bprob_BK = BifurcationProblem(bistable_switch_BK, [1.0], [5.0, 2.5, 3, 1.0], (@lens _[1]); record_from_solution = (x, p) -> x[1])
+    
+    # Check the same function have been generated.
+    bprob.u0 == bprob_BK.u0
+    bprob.params == bprob_BK.params
+    for repeat = 1:20
+        u0 = rand(rng, 1)
+        p = rand(rng, 4)
+        @test bprob_BK.VF.F(u0, p) == bprob.VF.F(u0, p)
+    end
+end
+
+# Creates a system where rn is composed of 4, somewhat nested, networks.
+# Tests with defaults within nested networks.
+let
+    rn1 = @reaction_network rn1 begin
+        @parameters p=1.0
+        (p, d), 0 <--> X
+    end
+    rn2 = @reaction_network rn2 begin
+        @parameters p=2.0
+        (p, d), 0 <--> X
+    end
+    rn3 = @reaction_network rn3 begin
+        @parameters p=3.0
+        (p, d), 0 <--> X
+    end
+    rn4 = @reaction_network rn4 begin
+        @parameters p=4.0
+        (p, d), 0 <--> X
+    end
+    @named rn3 =compose(rn3, [rn4])
+    @named rn = compose(rn1, [rn2, rn3])
+
+    # Declares parameter values and initial u guess.
+    @unpack X, d = rn
+    u0_guess = [X => 1.0, rn2.X => 1.0, rn3.X => 1.0, rn3.rn4.X => 1.0]
+    p_start = [d => 1.0, rn2.d => 1.0, rn3.d => 1.0, rn3.rn4.d => 1.0]
+
+
+    # Computes bifurcation diagrams and check them (the final point at [p_value = 6] should be =[p/d]).
+    p_span = (0.1, 6.0)
+
+    let
+        # Checks top layer.
+        bprob = BifurcationProblem(rn, u0_guess, p_start, d; plot_var = X)
+        opts_br = ContinuationPar(dsmin = 0.0001, dsmax = 0.001, ds = 0.0001, max_steps = 10000, p_min = p_span[1], p_max = p_span[2], n_inversion = 4)
+        bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
+        @test bif_dia.γ.branch[end].x ≈ 1.0/6
+
+        # Checks second layer (1).
+        bprob = BifurcationProblem(rn, u0_guess, p_start, rn2.d; plot_var = rn2.X)
+        opts_br = ContinuationPar(dsmin = 0.0001, dsmax = 0.001, ds = 0.0001, max_steps = 10000, p_min = p_span[1], p_max = p_span[2], n_inversion = 4)
+        bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
+        @test bif_dia.γ.branch[end].x ≈ 2.0/6
+
+        # Checks second layer (2).
+        bprob = BifurcationProblem(rn, u0_guess, p_start, rn3.d; plot_var = rn3.X)
+        opts_br = ContinuationPar(dsmin = 0.0001, dsmax = 0.001, ds = 0.0001, max_steps = 10000, p_min = p_span[1], p_max = p_span[2], n_inversion = 4)
+        bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
+        @test bif_dia.γ.branch[end].x ≈ 3.0/6
+
+        # Checks third layer.
+        bprob = BifurcationProblem(rn, u0_guess, p_start, rn3.rn4.d; plot_var = rn3.rn4.X)
+        opts_br = ContinuationPar(dsmin = 0.0001, dsmax = 0.001, ds = 0.0001, max_steps = 10000, p_min = p_span[1], p_max = p_span[2], n_inversion = 4)
+        bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
+        @test bif_dia.γ.branch[end].x ≈ 4.0/6
+    end
+end
+
+# Tests for nested model with conservation laws.
+let
+    # Creates model.
+    rn1 = @reaction_network rn1 begin
+        (k1, k2), X1 <--> X2
+    end
+    rn2 = @reaction_network rn2 begin
+        (l1, l2), Y1 <--> Y2
+    end
+    @named rn = compose(rn1, [rn2])
+
+    # Creates input parameter and species vectors.
+    @unpack X1, X2, k1, k2 = rn1
+    u_guess = [X1 => 1.0, X2 => 1.0, rn2.Y1 => 1.0, rn2.Y2 => 1.0]
+    p_start = [k1 => 1.0, k2 => 1.0, rn2.l1 => 1.0, rn2.l2 => 1.0]
+    u0 = [X1 => 1.0, X2 => 0.0, rn2.Y1 => 1.0, rn2.Y2 => 0.0]
+
+    # Computes bifurcation diagram.
+    p_span = (0.2, 5.0)
+    bprob = BifurcationProblem(rn, u_guess, p_start, k1; plot_var = X1, u0=u0)
+    opts_br = ContinuationPar(dsmin = 0.0001, dsmax = 0.001, ds = 0.0001, max_steps = 10000, p_min = p_span[1], p_max = p_span[2], n_inversion = 4)
+    bif_dia = bifurcationdiagram(bprob, PALC(), 2, (args...) -> opts_br; bothside = true)
+
+    # Checks that the bifurcation diagram is correct.
+    xs = getfield.(bif_dia.γ.branch, :x)
+    k1s = getfield.(bif_dia.γ.branch, :param)
+    @test all(1 ./ k1s .* (1 .- xs) .≈ xs)
+
+    # Checks that there is an error if information for conserved quantities computation is not provided.
+    @test_throws Exception bprob = BifurcationProblem(rn, u_guess, p_start, k1; plot_var = X1)
+end
diff --git a/test/programmatic_model_creation/component_based_model_creation.jl b/test/programmatic_model_creation/component_based_model_creation.jl
diff --git a/test/runtests.jl b/test/runtests.jl
diff --git a/test/test_networks.jl b/test/test_networks.jl
