CaoTauLeaping Implementation #603
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,25 +1,39 @@ | ||
| @cache mutable struct TauLeapingCache{uType, rateType, jumpRateType} <: StochasticDiffEqMutableCache | ||
| u::uType | ||
| uprev::uType | ||
| tmp::uType | ||
| rate::rateType | ||
| newrate::rateType | ||
| EEstcache::jumpRateType | ||
| end | ||
|
|
||
| function alg_cache(alg::TauLeaping, prob, u, ΔW, ΔZ, p, rate_prototype, noise_rate_prototype, jump_rate_prototype, ::Type{uEltypeNoUnits}, ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, f, t, dt, ::Type{Val{true}}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits} | ||
| tmp = zero(u) | ||
| rate = zero(jump_rate_prototype) | ||
| newrate = zero(jump_rate_prototype) | ||
| EEstcache = zero(jump_rate_prototype) | ||
| TauLeapingCache(u, uprev, tmp, rate, newrate, EEstcache) | ||
| end | ||
|
|
||
| @cache mutable struct CaoTauLeapingCache{uType, rateType, muType} <: StochasticDiffEqMutableCache | ||
| u::uType | ||
| uprev::uType | ||
| tmp::uType | ||
| rate::rateType | ||
| mu::muType | ||
| sigma2::muType | ||
| end | ||
|
|
||
| function alg_cache(alg::CaoTauLeaping, prob, u, ΔW, ΔZ, p, rate_prototype, noise_rate_prototype, jump_rate_prototype, ::Type{uEltypeNoUnits}, ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, f, t, dt, ::Type{Val{true}}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits} | ||
| tmp = zero(u) | ||
| rate = zero(jump_rate_prototype) | ||
| mu = zero(u) | ||
| sigma2 = zero(u) | ||
| CaoTauLeapingCache(u, uprev, tmp, rate, mu, sigma2) | ||
| end | ||
|
|
||
| struct TauLeapingConstantCache <: StochasticDiffEqConstantCache end | ||
| struct CaoTauLeapingConstantCache <: StochasticDiffEqConstantCache end | ||
|
|
||
| alg_cache(alg::TauLeaping, prob, u, ΔW, ΔZ, p, rate_prototype, noise_rate_prototype, jump_rate_prototype, ::Type{uEltypeNoUnits}, ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, f, t, dt, ::Type{Val{false}}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits} = TauLeapingConstantCache() | ||
| alg_cache(alg::CaoTauLeaping, prob, u, ΔW, ΔZ, p, rate_prototype, noise_rate_prototype, jump_rate_prototype, ::Type{uEltypeNoUnits}, ::Type{uBottomEltypeNoUnits}, ::Type{tTypeNoUnits}, uprev, f, t, dt, ::Type{Val{false}}) where {uEltypeNoUnits, uBottomEltypeNoUnits, tTypeNoUnits} = CaoTauLeapingConstantCache() |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -19,27 +19,85 @@ end | |
|
|
||
|
|
||
| function stepsize_controller!(integrator::SDEIntegrator, alg::TauLeaping) | ||
| nothing # Post-leap adjustment happens in perform_step! | ||
| end | ||
|
|
||
| function step_accept_controller!(integrator::SDEIntegrator, alg::TauLeaping) | ||
| if alg.adaptive | ||
| integrator.q = min(integrator.opts.gamma / integrator.EEst, integrator.opts.qmax) | ||
| return integrator.dt * integrator.q | ||
| else | ||
| return integrator.dt | ||
| end | ||
| end | ||
|
|
||
| function step_reject_controller!(integrator::SDEIntegrator, alg::TauLeaping) | ||
| if alg.adaptive | ||
| integrator.dt = integrator.opts.gamma * integrator.dt / integrator.EEst | ||
| end | ||
|
Comment on lines
+35
to
+37
|
||
| end | ||
|
|
||
| # CaoTauLeaping: Pre-leap τ computation | ||
| function stepsize_controller!(integrator::SDEIntegrator, alg::CaoTauLeaping) | ||
| if !alg.adaptive | ||
| return | ||
| end | ||
|
|
||
| @unpack u, p, t, P, opts, c = integrator | ||
| cache = integrator.cache | ||
|
|
||
| P === nothing && error("CaoTauLeaping requires a JumpProblem with a RegularJump") | ||
|
|
||
| # Handle both constant and mutable caches | ||
| if isa(cache, CaoTauLeapingConstantCache) | ||
| rate = P.cache.rate(u, p, t) # Compute propensities directly | ||
| mu = zero(u) | ||
| sigma2 = zero(u) | ||
| else # CaoTauLeapingCache | ||
| @unpack mu, sigma2, rate = cache | ||
| P.cache.rate(rate, u, p, t) # Compute propensities into cache | ||
| fill!(mu, zero(eltype(mu))) | ||
| fill!(sigma2, zero(eltype(sigma2))) | ||
| end | ||
|
|
||
| # Infer ν_ij using c by applying unit counts for each reaction | ||
| num_reactions = length(rate) | ||
| ν = zeros(eltype(u), length(u), num_reactions) | ||
| unit_counts = zeros(eltype(rate), num_reactions) | ||
| for j in 1:num_reactions | ||
| unit_counts[j] = 1 | ||
| c(ν[:, j], u, p, t, unit_counts, nothing) # ν[:, j] is the change vector for reaction j | ||
| unit_counts[j] = 0 # Reset | ||
| end | ||
|
Comment on lines
+63
to
+71
There was a problem hiding this comment. This is fine for now, but we really should be storing that differently in the jumps. @isaacsas this is worth thinking about. There was a problem hiding this comment. I think it would be worthwhile to make leaping work with individually defined jumps too. Split-step methods require that anyways. But that wouldn’t solve this issue for general non-mass action jumps. There was a problem hiding this comment. and we'd need two implementations of each leaping? I'm thinking the interface should maybe require more information in the jump building, under the assumption users will be using something like Catalyst rather than building most jumps by hand. To do some of this optimally we have to shift been a few different representations, where for a user this would be redundant information to the solver it's different views that have different optimal uses. There was a problem hiding this comment. For methods that work on the full vector at once, can’t we can manually build the needed function from the list of jumps? So those methods can stick with the current interface, but then we could support methods that really need jump-by-jump rate/affect access too. But yes, Catalyst should build the needed inputs for different representations. There was a problem hiding this comment.
You get a lot of compile time trying to inline all of the small functions into the large one. Also there's a lack of fusion. So it's quite slow to do it that way. I think the point is, some times you need jump-by-jump information, and sometimes you want the aggregate information (all tau leap steppers will use an aggregate function a lot of the time), and those are the same information in two very different computational representations, and building one from the other isn't easy to do from the implementation but it is trivial to do symbolically. There was a problem hiding this comment. Definitely, so it should be preferred that methods that work on the full vector get a single input function, but having a fallback would be useful. Catalyst/MTK, should just generate both inputs (or whichever is most appropriate for the selected integrator -- we can add a function to query the preferred input type). |
||
|
|
||
| # Compute μ_i and σ_i^2 | ||
| for i in eachindex(u) | ||
| for j in 1:num_reactions | ||
| ν_ij = ν[i, j] | ||
| mu[i] += ν_ij * rate[j] | ||
|
There was a problem hiding this comment. The formula uses the derivative of the rate function, not the rate itself. |
||
| sigma2[i] += ν_ij^2 * rate[j] | ||
| end | ||
| end | ||
|
|
||
| # Compute τ per species | ||
| ϵ = alg.epsilon | ||
| τ_vals = similar(u, Float64) | ||
| for i in eachindex(u) | ||
| max_term = max(ϵ * u[i], 1.0) | ||
|
There was a problem hiding this comment. This is supposed to use the sum of rates term, not u[i], also why the 1.0? There was a problem hiding this comment. this is missing the g_i factor |
||
| τ1 = abs(mu[i]) > 0 ? max_term / abs(mu[i]) : Inf | ||
| τ2 = sigma2[i] > 0 ? max_term^2 / sigma2[i] : Inf | ||
| τ_vals[i] = min(τ1, τ2) | ||
| end | ||
|
|
||
| τ = min(minimum(τ_vals), opts.dtmax) | ||
| integrator.dt = max(τ, opts.dtmin) | ||
| integrator.EEst = 1.0 | ||
| end | ||
|
|
||
| function step_accept_controller!(integrator::SDEIntegrator, alg::CaoTauLeaping) | ||
| return integrator.dt | ||
| end | ||
|
|
||
| function step_reject_controller!(integrator::SDEIntegrator, alg::CaoTauLeaping) | ||
| error("CaoTauLeaping should never reject steps") | ||
| end | ||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,40 +1,66 @@ | ||
| # Perform Step | ||
| @muladd function perform_step!(integrator, cache::TauLeapingConstantCache) | ||
| @unpack t, dt, uprev, u, p, P, c = integrator | ||
|
|
||
| P === nothing && error("TauLeaping requires a JumpProblem with a RegularJump") | ||
| P.dt = dt | ||
| tmp = c(uprev, p, t, P.dW, nothing) | ||
| integrator.u = uprev .+ tmp | ||
|
|
||
| if integrator.alg.adaptive | ||
| oldrate = P.cache.currate | ||
| newrate = P.cache.rate(integrator.u, p, t + dt) | ||
| EEstcache = @. abs(newrate - oldrate) / max(50 * integrator.opts.reltol * oldrate, integrator.rate_constants / integrator.dt) | ||
| integrator.EEst = integrator.opts.internalnorm(EEstcache, t) | ||
| if integrator.EEst <= 1 | ||
| P.cache.currate = newrate | ||
| end | ||
| else | ||
| integrator.EEst = 1.0 | ||
| end | ||
| end | ||
|
|
||
| @muladd function perform_step!(integrator, cache::TauLeapingCache) | ||
| @unpack t, dt, uprev, u, p, P, c = integrator | ||
| @unpack tmp, rate, newrate, EEstcache = cache | ||
|
|
||
| P === nothing && error("TauLeaping requires a JumpProblem with a RegularJump") | ||
| P.dt = dt | ||
| c(tmp, uprev, p, t, P.dW, nothing) | ||
| @.. u = uprev + tmp | ||
|
|
||
| if integrator.alg.adaptive | ||
| P.cache.rate(newrate, u, p, t + dt) | ||
| P.cache.rate(rate, uprev, p, t) | ||
| @.. EEstcache = abs(newrate - rate) / max(50 * integrator.opts.reltol * rate, integrator.rate_constants / integrator.dt) | ||
| integrator.EEst = integrator.opts.internalnorm(EEstcache, t) | ||
| if integrator.EEst <= 1 | ||
| P.cache.currate .= newrate | ||
| end | ||
| else | ||
| integrator.EEst = 1.0 | ||
| end | ||
| end | ||
|
|
||
| @muladd function perform_step!(integrator, cache::CaoTauLeapingConstantCache) | ||
| @unpack t, dt, uprev, u, p, P, c = integrator | ||
|
|
||
| P === nothing && error("CaoTauLeaping requires a JumpProblem with a RegularJump") | ||
| P.dt = dt | ||
| tmp = c(uprev, p, t, P.dW, nothing) | ||
| integrator.u = uprev .+ tmp | ||
|
|
||
| integrator.EEst = 1.0 | ||
| end | ||
|
|
||
| @muladd function perform_step!(integrator, cache::CaoTauLeapingCache) | ||
| @unpack t, dt, uprev, u, p, P, c = integrator | ||
| @unpack tmp = cache | ||
|
|
||
| P === nothing && error("CaoTauLeaping requires a JumpProblem with a RegularJump") | ||
| P.dt = dt | ||
| c(tmp, uprev, p, t, P.dW, nothing) | ||
| @.. u = uprev + tmp | ||
|
|
||
| integrator.EEst = 1.0 | ||
| end |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
what are these doing here?