SimpleAdaptiveTauLeaping solver

src/simple_regular_solve.jl

@@ -1,5 +1,11 @@
 struct SimpleTauLeaping <: DiffEqBase.DEAlgorithm end
 
+struct SimpleAdaptiveTauLeaping{T <: AbstractFloat} <: DiffEqBase.DEAlgorithm
+    epsilon::T  # Error control parameter
+end
+
+SimpleAdaptiveTauLeaping(; epsilon=0.05) = SimpleAdaptiveTauLeaping(epsilon)
+
 function validate_pure_leaping_inputs(jump_prob::JumpProblem, alg)
     if !(jump_prob.aggregator isa PureLeaping)
         @warn "When using $alg, please pass PureLeaping() as the aggregator to the \
@@ -14,6 +20,19 @@ function validate_pure_leaping_inputs(jump_prob::JumpProblem, alg)
     jump_prob.regular_jump !== nothing    
 end
 
+function validate_pure_leaping_inputs(jump_prob::JumpProblem, alg::SimpleAdaptiveTauLeaping)
+    if !(jump_prob.aggregator isa PureLeaping)
+        @warn "When using $alg, please pass PureLeaping() as the aggregator to the \
+        JumpProblem, i.e. call JumpProblem(::DiscreteProblem, PureLeaping(),...). \
+        Passing $(jump_prob.aggregator) is deprecated and will be removed in the next breaking release."
+    end
+    isempty(jump_prob.jump_callback.continuous_callbacks) &&
+    isempty(jump_prob.jump_callback.discrete_callbacks) &&
+    isempty(jump_prob.constant_jumps) &&
+    isempty(jump_prob.variable_jumps) &&
+    jump_prob.massaction_jump !== nothing
+end
+
 function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping;
         seed = nothing,
         dt = error("dt is required for SimpleTauLeaping."))
@@ -62,6 +81,164 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping;
         interp = DiffEqBase.ConstantInterpolation(t, u))
 end
 
+function compute_hor(reactant_stoch, numjumps)
+    hor = zeros(Int, numjumps)
+    for j in 1:numjumps
+        order = sum(stoch for (spec_idx, stoch) in reactant_stoch[j]; init=0)
+        if order > 3
+            error("Reaction $j has order $order, which is not supported (maximum order is 3).")
+        end
+        hor[j] = order
+    end
+    return hor
+end
+
+function compute_gi(u, nu, hor, i)
+    max_gi = 1
+    for j in 1:size(nu, 2)
+        if nu[i, j] < 0  # Species i is a substrate
+            if hor[j] == 1
+                max_gi = max(max_gi, 1)
+            elseif hor[j] == 2 || hor[j] == 3
+                stoch = abs(nu[i, j])
+                if stoch >= 2
+                    gi = 2 / stoch + 1 / (stoch - 1)
+                    max_gi = max(max_gi, ceil(Int, gi))
+                elseif stoch == 1
+                    max_gi = max(max_gi, hor[j])
+                end
+            end
+        end
+    end
+    return max_gi
+end
+
+function compute_tau_explicit(u, rate_cache, nu, hor, p, t, epsilon, rate, dtmin)
+    rate(rate_cache, u, p, t)
+    tau = Inf
+    for i in 1:length(u)
+        mu = zero(eltype(u))
+        sigma2 = zero(eltype(u))
+        for j in 1:size(nu, 2)
+            mu += nu[i, j] * rate_cache[j]
+            sigma2 += nu[i, j]^2 * rate_cache[j]
+        end
+        gi = compute_gi(u, nu, hor, i)
+        bound = max(epsilon * u[i] / gi, 1.0)
+        mu_term = abs(mu) > 0 ? bound / abs(mu) : Inf
+        sigma_term = sigma2 > 0 ? bound^2 / sigma2 : Inf
+        tau = min(tau, mu_term, sigma_term)
+    end
+    return max(tau, dtmin)
+end
+
+function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleAdaptiveTauLeaping; 
+        seed = nothing,
+        dtmin = 1e-10,
+        saveat = nothing)
+    validate_pure_leaping_inputs(jump_prob, alg) ||
+        error("SimpleAdaptiveTauLeaping can only be used with PureLeaping JumpProblem with a MassActionJump.")
+    prob = jump_prob.prob
+    rng = DEFAULT_RNG
+    (seed !== nothing) && seed!(rng, seed)
+
+    maj = jump_prob.massaction_jump
+    numjumps = get_num_majumps(maj)
+    rj = jump_prob.regular_jump
+    # Extract rates
+    rate = rj !== nothing ? rj.rate :
+        (out, u, p, t) -> begin
+            for j in 1:numjumps
+                out[j] = evalrxrate(u, j, maj)
+            end
+        end
+    c = rj !== nothing ? rj.c : nothing
+    u0 = copy(prob.u0)
+    tspan = prob.tspan
+    p = prob.p
+
+    # Initialize current state and saved history
+    u_current = copy(u0)
+    t_current = tspan[1]
+    usave = [copy(u0)]
+    tsave = [tspan[1]]
+    rate_cache = zeros(float(eltype(u0)), numjumps)
+    counts = zero(rate_cache)
+    du = similar(u0)
+    t_end = tspan[2]
+    epsilon = alg.epsilon
+
+    # Extract stochiometry once from MassActionJump
+    nu = zeros(float(eltype(u0)), length(u0), numjumps)
+    for j in 1:numjumps
+        for (spec_idx, stoch) in maj.net_stoch[j]
+            nu[spec_idx, j] = stoch
+        end
+    end
+    # Extract reactant stochiometry for hor
+    reactant_stoch = maj.reactant_stoch
+    hor = compute_hor(reactant_stoch, numjumps)
+
+    # Set up saveat_times
+    saveat_times = nothing
+    if isnothing(saveat)
+        saveat_times = Vector{typeof(tspan[1])}()
+    elseif saveat isa Number
+        saveat_times = collect(range(tspan[1], tspan[2], step=saveat))
+    else
+        saveat_times = collect(saveat)
+    end
+
+    save_idx = 1
+
+    while t_current < t_end
+        rate(rate_cache, u_current, p, t_current)
+        tau = compute_tau_explicit(u_current, rate_cache, nu, hor, p, t_current, epsilon, rate, dtmin)
+        tau = min(tau, t_end - t_current)
+        if !isempty(saveat_times) && save_idx <= length(saveat_times) && t_current + tau > saveat_times[save_idx]
+            tau = saveat_times[save_idx] - t_current
+        end
+        counts .= pois_rand.(rng, max.(rate_cache * tau, 0.0))
+        du .= 0
+        if c !== nothing
+            c(du, u_current, p, t_current, counts, nothing)
+        else
+            for j in 1:numjumps
+                for (spec_idx, stoch) in maj.net_stoch[j]
+                    du[spec_idx] += stoch * counts[j]
+                end
+            end
+        end
+        u_new = u_current + du
+        if any(<(0), u_new)
+            # Halve tau to avoid negative populations, as per Cao et al. (2006, J. Chem. Phys., DOI: 10.1063/1.2159468)
+            tau /= 2
+            continue
+        end
+        for i in eachindex(u_new)
+            u_new[i] = max(u_new[i], 0)
+        end
+        t_new = t_current + tau
+
+        # Save state if at a saveat time or if saveat is empty
+        if isempty(saveat_times) || (save_idx <= length(saveat_times) && t_new >= saveat_times[save_idx])
+            push!(usave, u_new)
+            push!(tsave, t_new)
+            if !isempty(saveat_times) && t_new >= saveat_times[save_idx]
+                save_idx += 1
+            end
+        end
+
+        u_current = u_new
+        t_current = t_new
+    end
+
+    sol = DiffEqBase.build_solution(prob, alg, tsave, usave,
+        calculate_error=false,
+        interp=DiffEqBase.ConstantInterpolation(tsave, usave))
+    return sol
+end
+
 struct EnsembleGPUKernel{Backend} <: SciMLBase.EnsembleAlgorithm
     backend::Backend
     cpu_offload::Float64
@@ -74,3 +251,5 @@ end
 function EnsembleGPUKernel()
     EnsembleGPUKernel(nothing, 0.0)
 end
+
+export SimpleTauLeaping, EnsembleGPUKernel, SimpleAdaptiveTauLeaping