Merge pull request #70 from augustinas1/u0map

Support initial state mapping and streamline ODEProblem

Author: augustinas1

Project.toml

@@ -34,9 +34,10 @@ TupleTools = "1.2.0"
 julia = "1.6"
 
 [extras]
+OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a"
 JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "SafeTestsets", "JumpProcesses"]
+test = ["Test", "SafeTestsets", "JumpProcesses", "OrdinaryDiffEqTsit5"]

docs/src/api/momentclosure_api.md

@@ -31,7 +31,13 @@ bernoulli_moment_eqs
 ```@docs
 moment_closure
 ClosedMomentEquations
+```
+
+## Solving Moment Equations
+
+```@docs
 deterministic_IC
+ODEProblem(::MomentEquations, ::Tuple, ::Tuple, ::Tuple)
 ```
 
 ## Basic Accessor Functions
@@ -43,6 +49,7 @@ We also define a few accessor functions that return system information from the
 * `ModelingToolkit.get_iv(sys::MomentEquations)`: The independent variable used in the system.
 * `ModelingToolkit.get_ps(sys::MomentEquations)`: The parameters of the system.
 * `ModelingToolkit.unknowns(sys::MomentEquations)`: The set of unknowns (moments) in the equations.
+* `Catalyst.speciesmap(sys::MomentEquations)`: The dictionary mapping the chemical species in a `Catalyst.ReactionSystem` to their index within the corresponding moment equations.
 * `MomentClosure.get_closure(sys::ClosedMomentEquations)`: The dictionary of moment closure functions for each higher order moment.
 
 ## [Displaying Equations and Closures](@id visualisation_api)

docs/src/tutorials/LMA_example.md

@@ -71,14 +71,12 @@ Note that the results agree with Eqs. (1) and (2) (after a corresponding substit
 ```julia
 using OrdinaryDiffEq, Sundials, Plots
 
-# [g, p] ordered as in `speciesmap(rn_nonlinear)`
-u₀ = [1.0, 0.001]
-p = Dict(:σ_b => 0.004, :σ_u => 0.25, :ρ_b => 25.0, :ρ_u => 60.0)
+u0map = [:g => 1.0, :p => 0.001]
+pmap = Dict(:σ_b => 0.004, :σ_u => 0.25, :ρ_b => 25.0, :ρ_u => 60.0)
 tspan = (0., 15.)
 dt = 0.1
 
-u₀map = deterministic_IC(u₀, LMA_eqs)
-oprob_LMA = ODEProblem(LMA_eqs, u₀map, tspan, pmap)
+oprob_LMA = ODEProblem(LMA_eqs, u0map, tspan, pmap)
 sol_LMA = solve(oprob_LMA, CVODE_BDF(), saveat=dt)
 
 plot(sol_LMA, idxs=[2], label="LMA", ylabel="⟨p⟩", xlabel="time", fmt="svg")

docs/src/tutorials/P53_system_example.md

@@ -59,10 +59,10 @@ rn = @reaction_network begin
 end
 
 # parameters
-p = [:k₁ => 90, :k₂ => 0.002, :k₃ => 1.7, :k₄ => 1.1, :k₅ => 0.93, :k₆ => 0.96, :k₇ => 0.01]
+pmap = [:k₁ => 90, :k₂ => 0.002, :k₃ => 1.7, :k₄ => 1.1, :k₅ => 0.93, :k₆ => 0.96, :k₇ => 0.01]
 
 # initial molecule numbers [x, y₀, y]
-u₀ = [70, 30, 60]
+u0map = [:x=> 70, :y₀ => 30, :y => 60]
 ```
 
 Let's first simulate the reaction network using SSA in order to have a reference point of the real system dynamics. We choose a relatively long simulation time span in order to clearly see how the molecule numbers converge to their steady-state values and opt for $5 \times 10^4$ SSA realisations:
@@ -73,7 +73,7 @@ tspan = (0., 200.)
 # constructing the discrete jump problem using DifferentialEquations
 jsys = convert(JumpSystem, rn, combinatoric_ratelaws=false)
 jsys = complete(jsys)
-dprob = DiscreteProblem(jsys, u₀, tspan, p)
+dprob = DiscreteProblem(jsys, u0map, tspan, pmap)
 
 jprob = JumpProblem(jsys, dprob, Direct(), save_positions=(false, false))
 ensembleprob  = EnsembleProblem(jprob)
@@ -123,10 +123,7 @@ for q in 3:6
     eqs = generate_central_moment_eqs(rn, 2, q, combinatoric_ratelaws=false)
     for (closure, plt) in zip(closures, plts)
         closed_eqs = moment_closure(eqs, closure)
-
-        u₀map = deterministic_IC(u₀, closed_eqs)
-        oprob = ODEProblem(closed_eqs, u₀map, tspan, p)
-
+        oprob = ODEProblem(closed_eqs, u0map, tspan, pmap)
         sol = solve(oprob, Tsit5(), saveat=0.1)
         plt = plot!(plt, sol, idxs=[1], lw=3, label  = "q = "*string(q))
     end
@@ -171,12 +168,9 @@ for q in [4,6]
     eqs = generate_central_moment_eqs(rn, 2, q, combinatoric_ratelaws=false)
     for closure in closures
         closed_eqs = moment_closure(eqs, closure)
-
-        u₀map = deterministic_IC(u₀, closed_eqs)
-        oprob = ODEProblem(closed_eqs, u₀map, tspan, p)
+        oprob = ODEProblem(closed_eqs, u0map, tspan, pmap)
         sol = solve(oprob, Tsit5(), saveat=0.1)
-
-        # index of M₂₀₀ can be checked with `u₀map` or `closed_eqs.odes.states`
+        # index of M₂₀₀ can be checked with `unknowns(closed_eqs)`
         plt = plot!(plt, sol, idxs=[4], lw=3, label  = closure*" q = "*string(q))
     end
 end
@@ -202,11 +196,8 @@ for (q, plt_m, plt_v) in zip(q_vals, plt_means, plt_vars)
 
     eqs = generate_central_moment_eqs(rn, 3, q, combinatoric_ratelaws=false)
     for closure in closures
-
         closed_eqs = moment_closure(eqs, closure)
-
-        u₀map = deterministic_IC(u₀, closed_eqs)
-        oprob = ODEProblem(closed_eqs, u₀map, tspan, p)
+        oprob = ODEProblem(closed_eqs, u0map, tspan, pmap)
 
         sol = solve(oprob, Tsit5(), saveat=0.1)
         plt_m = plot!(plt_m, sol, idxs=[1], label = closure)    
@@ -276,9 +267,7 @@ oprobs = Dict()
 
 for closure in closures
     closed_eqs = moment_closure(eqs, closure)
-
-    u₀map = deterministic_IC(u₀, closed_eqs)
-    oprobs[closure] = ODEProblem(closed_eqs, u₀map, tspan, p)
+    oprobs[closure] = ODEProblem(closed_eqs, u0map, tspan, pmap)
     sol = solve(oprobs[closure], Tsit5(), saveat=0.1)
 
     plt = plot!(plt, sol, idxs=[1], label = closure)    

docs/src/tutorials/common_issues.md

@@ -13,8 +13,8 @@ rn = @reaction_network begin
   (c₃*Ω, c₄), 0 ↔ X
 end
 
-p = [:c₁ => 0.9, :c₂ => 2, :c₃ => 1, :c₄ => 1, :Ω => 100]
-u₀ = [1, 1]
+pmap = [:c₁ => 0.9, :c₂ => 2, :c₃ => 1, :c₄ => 1, :Ω => 100]
+u0map = [:X => 1, :Y => 1]
 tspan = (0., 100.)
 
 raw_eqs = generate_raw_moment_eqs(rn, 2, combinatoric_ratelaws=false)
@@ -23,8 +23,7 @@ As we have seen earlier, second-order moment expansion using normal closure appr
 ```julia
 closed_raw_eqs = moment_closure(raw_eqs, "zero")
 
-u₀map = deterministic_IC(u₀, closed_raw_eqs)
-oprob = ODEProblem(closed_raw_eqs, u₀map, tspan, p)
+oprob = ODEProblem(closed_raw_eqs, u0map, tspan, pmap)
 sol = solve(oprob, Tsit5(), saveat=0.1)
 
 plot(sol, idxs=[1,2], lw=2)
@@ -37,8 +36,7 @@ Let's apply log-normal closure next:
 ```julia
 closed_raw_eqs = moment_closure(raw_eqs, "log-normal")
 
-u₀map = deterministic_IC(u₀, closed_raw_eqs)
-oprob = ODEProblem(closed_raw_eqs, u₀map, tspan, p)
+oprob = ODEProblem(closed_raw_eqs, u0map, tspan, pmap)
 sol = solve(oprob, Tsit5(), saveat=0.1)
 
 plot(sol, idxs=[1,2], lw=2, legend=:bottomright)
@@ -52,8 +50,7 @@ Normal closure is also quite fragile. This can be seen by simply including the c
 raw_eqs = generate_raw_moment_eqs(rn, 2, combinatoric_ratelaws=true)
 closed_raw_eqs = moment_closure(raw_eqs, "normal")
 
-u₀map = deterministic_IC(u₀, closed_raw_eqs)
-oprob = ODEProblem(closed_raw_eqs, u₀map, tspan, p)
+oprob = ODEProblem(closed_raw_eqs, u0map, tspan, pmap)
 sol = solve(oprob, Tsit5(), saveat=0.1)
 
 plot(sol, idxs=[1,2], lw=2)
@@ -65,8 +62,7 @@ Nevertheless, this can be improved upon by increasing the order of moment expans
 raw_eqs = generate_raw_moment_eqs(rn, 3, combinatoric_ratelaws=true)
 closed_raw_eqs = moment_closure(raw_eqs, "normal")
 
-u₀map = deterministic_IC(u₀, closed_raw_eqs)
-oprob = ODEProblem(closed_raw_eqs, u₀map, tspan, p)
+oprob = ODEProblem(closed_raw_eqs, u0map, tspan, pmap)
 sol = solve(oprob, Tsit5(), saveat=0.1)
 
 plot(sol, idxs=[1,2], lw=2, legend=:bottomright)
@@ -78,8 +74,7 @@ Some dampening in the system is now visible. Increasing the expansion order to `
 raw_eqs = generate_raw_moment_eqs(rn, 4, combinatoric_ratelaws=true)
 closed_raw_eqs = moment_closure(raw_eqs, "normal")
 
-u₀map = deterministic_IC(u₀, closed_raw_eqs)
-oprob = ODEProblem(closed_raw_eqs, u₀map, tspan, p)
+oprob = ODEProblem(closed_raw_eqs, u0map, tspan, pmap)
 sol = solve(oprob, Tsit5(), saveat=0.1)
 
 plot(sol, idxs=[1,2], lw=2)
@@ -91,8 +86,7 @@ For dessert, we consider unphysical divergent trajectories—a frequent problem
 raw_eqs = generate_raw_moment_eqs(rn, 2, combinatoric_ratelaws=true)
 closed_raw_eqs = moment_closure(raw_eqs, "log-normal")
 
-u₀map = deterministic_IC(u₀, closed_raw_eqs)
-oprob = ODEProblem(closed_raw_eqs, u₀map, tspan, p)
+oprob = ODEProblem(closed_raw_eqs, u0map, tspan, pmap)
 sol = solve(oprob, Rodas4P(), saveat=0.1)
 
 plot(sol, idxs=[1,2], lw=2)
@@ -104,8 +98,7 @@ In contrast to normal closure, increasing the expansion order makes the problem
 raw_eqs = generate_raw_moment_eqs(rn, 3, combinatoric_ratelaws=true)
 closed_raw_eqs = moment_closure(raw_eqs, "log-normal")
 
-u₀map = deterministic_IC(u₀, closed_raw_eqs)
-oprob = ODEProblem(closed_raw_eqs, u₀map, tspan, p)
+oprob = ODEProblem(closed_raw_eqs, u0map, tspan, pmap)
 sol = solve(oprob, Rodas4P(), saveat=0.1)
 
 plot(sol, idxs=[1,2], lw=2)

docs/src/tutorials/derivative_matching_example.md

@@ -18,9 +18,9 @@ rn = @reaction_network begin
 end
 
 # parameter values
-p = [:c₁ => 1.0, :c₂ => 1.0]
+pmap = [:c₁ => 1.0, :c₂ => 1.0]
 # initial conditions
-u0 = [20, 10]
+u0map = [:x₁ => 20, :x₂ => 10]
 # time interval to solve on
 tspan = (0., 0.5)
 ```
@@ -46,8 +46,7 @@ Note that all closure functions are consistent with the ones shown in Table II o
 ```julia
 using OrdinaryDiffEqTsit5
 
-u0map = deterministic_IC(u0, dm2_eqs) # assuming deterministic initial conditions
-oprob = ODEProblem(dm2_eqs, u0map, tspan, p)
+oprob = ODEProblem(dm2_eqs, u0map, tspan, pmap)
 dm2_sol = solve(oprob, Tsit5(), saveat=0.01)
 ```
 Now the question is how can we extract the time evolution of the cumulant $\kappa_{03}$. Firstly, note that using the standard moment relationships it can be expressed in terms of raw moments as:
@@ -108,8 +107,7 @@ unknowns(dm3_eqs.odes)
 ```
 and solve the moment equations, computing the required cumulant:
 ```julia
-u0map = deterministic_IC(u0, dm3_eqs)
-oprob = ODEProblem(dm3_eqs, u0map, tspan, p)
+oprob = ODEProblem(dm3_eqs, u0map, tspan, pmap)
 dm3_sol = solve(oprob, Tsit5(), saveat=0.01, abstol=1e-8, reltol=1e-8)
 
 μ₀₁ = dm3_sol[2,:]
@@ -122,8 +120,7 @@ Note that we could have also obtained $\kappa_{03}$ estimate in an easier way by
 central_eqs3 = generate_central_moment_eqs(rn, 3)
 dm3_central_eqs = moment_closure(central_eqs3, "derivative matching")
 
-u0map = deterministic_IC(u0, dm3_central_eqs)
-oprob = ODEProblem(dm3_central_eqs, u0map, tspan, p)
+oprob = ODEProblem(dm3_central_eqs, u0map, tspan, pmap)
 dm3_central_sol = solve(oprob, Tsit5(), saveat=0.01, abstol=1e-8, reltol=1e-8)
 
 # check that the two estimates are equivalent
@@ -136,7 +133,7 @@ The last ingredient we need for a proper comparison between the second and third
 ```julia
 using JumpProcesses
 
-dprob = DiscreteProblem(rn, u0, tspan, p)
+dprob = DiscreteProblem(rn, u0map, tspan, pmap)
 jprob = JumpProblem(rn, dprob, Direct(), save_positions=(false, false))
 
 ensembleprob  = EnsembleProblem(jprob)

docs/src/tutorials/geometric_reactions+conditional_closures.md

@@ -142,8 +142,8 @@ pmap = [:k_on => k_on_val,
           :γ_p => γ_p_val,
           :b => mean_b]
 
-# initial gene state and protein number, order [g, p]
-u₀ = [1, 1]
+# initial gene state and protein number
+u0map = [:g => 1, :P => 1]
 
 # time interval to solve on
 tspan = (0., 6.0)
@@ -155,7 +155,7 @@ jsys = convert(JumpSystem, rn; combinatoric_ratelaws=false)
 jsys = complete(jsys)
 
 # create a discrete problem setting the simulation parameters
-dprob = DiscreteProblem(u₀, tspan, pmap)
+dprob = DiscreteProblem(u0map, tspan, pmap)
 
 # create a JumpProblem compatible with ReactionSystemMod
 jprob = JumpProblem(rn, dprob, Direct(), save_positions=(false, false))
@@ -174,9 +174,6 @@ We continue to solve the moment equations for each closure:
 plt_m = plot()   # plot mean protein number
 plt_std = plot() # plot ssd of protein number
 
-# construct the initial molecule number mapping
-u₀map = deterministic_IC(u₀, dm_eqs)
-
 # solve moment ODEs for each closure and plot the results
 for closure in ["normal", "derivative matching",
                 "conditional gaussian", "conditional derivative matching"]
@@ -185,7 +182,7 @@ for closure in ["normal", "derivative matching",
     closed_eqs = moment_closure(eqs, closure, binary_vars)
 
     # solve the system of moment ODEs
-    oprob = ODEProblem(closed_eqs, u₀map, tspan, pmap)
+    oprob = ODEProblem(closed_eqs, u0map, tspan, pmap)
     sol = solve(oprob, AutoTsit5(Rosenbrock23()), saveat=0.01)
 
     # μ₀₁ is 2nd and μ₀₂ is 4th element in sol
@@ -273,7 +270,7 @@ pmap = [:kx_on => kx_on_val,
           :b_y => mean_b_y]
 
 # initial gene state and protein number, order [g_x, g_y, x, y]
-u₀ = [1, 1, 1, 1]
+u0map = [:g_x => 1, :g_y => 1, :x => 1, :y => 1]
 
 # time interval to solve on
 tspan = (0., 12.0)
@@ -284,7 +281,7 @@ We can run SSA as follows:
 ```julia
 jsys = convert(JumpSystem, rn, combinatoric_ratelaws=false)
 jsys = complete(jsys)
-dprob = DiscreteProblem(jsys, u₀, tspan, pmap)
+dprob = DiscreteProblem(jsys, u0map, tspan, pmap)
 jprob = JumpProblem(jsys, dprob, Direct(), save_positions=(false, false))
 
 ensembleprob  = EnsembleProblem(jprob)
@@ -302,9 +299,7 @@ plt_std = plot() # plot ssd of activator protein number
 for closure in ["derivative matching", "conditional derivative matching"]
 
     closed_eqs = moment_closure(eqs, closure, binary_vars)
-
-    u₀map = deterministic_IC(u₀, closed_eqs)
-    oprob = ODEProblem(closed_eqs, u₀map, tspan, pmap)
+    oprob = ODEProblem(closed_eqs, u0map, tspan, pmap)
     sol = solve(oprob, Tsit5(), saveat=0.1)
 
     # μ₀₀₀₁ is the 4th and μ₀₀₀₂ is the 12th element in sol (can check with closed_eqs.odes.states)

docs/src/tutorials/parameter_estimation_SDE.md

@@ -109,8 +109,7 @@ Now we approach the same model identification problem via MAs in the hope of cut
 using MomentClosure
 
 LV_moments = moment_closure(generate_raw_moment_eqs(LV, 2), "log-normal")
-u0map = deterministic_IC(last.(u0), LV_moments)
-prob_MA = ODEProblem(LV_moments, u0map, (0.0, Tf), zeros(5))
+prob_MA = ODEProblem(LV_moments, u0, (0.0, Tf), zeros(5))
 psetter_MA! = setp(prob_MA, (γ1, γ2, γ3, γ4, γ5))
 
 function obj_MA(p)

docs/src/tutorials/time-dependent_propensities.md

@@ -48,16 +48,13 @@ closed_raw_eqs = moment_closure(raw_eqs, "normal")
 p = [:c₁ => 0.9, :c₂ => 2., :c₃ => 1., :c₄ => 1., :Ω => 5., :ω => 1., :τ => 40.]
 
 # initial molecule numbers of species [X, Y]
-u₀ = [1., 1.]
-
-# deterministic initial conditions
-u₀map = deterministic_IC(u₀, closed_raw_eqs)
+u0map = [1., 1.]
 
 # time interval to solve one on
 tspan = (0., 100.)
 
 # convert the closed raw moment equations into a DifferentialEquations ODEProblem
-oprob = ODEProblem(closed_raw_eqs, u₀map, tspan, p)
+oprob = ODEProblem(closed_raw_eqs, u0map, tspan, pmap)
 
 # solve using Tsit5 solver
 sol = solve(oprob, Tsit5(), saveat=0.2)
@@ -69,7 +66,7 @@ It would be great to compare our results to the true dynamics. Using Differentia
 ```julia
 using JumpProcesses
 
-jinputs = JumpInputs(rn, u₀, tspan, p, combinatoric_ratelaws=false)
+jinputs = JumpInputs(rn, u0map, tspan, pmap, combinatoric_ratelaws=false)
 jprob = JumpProblem(jinputs, Direct())
 ```
 Note that now we have to provide an ODE solver to `solve` in order to integrate over the time-dependent propensities.