refactor

Author: sivasathyaseeelan

src/simple_regular_solve.jl

@@ -61,20 +61,169 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping;
         interp = DiffEqBase.ConstantInterpolation(t, u))
 end
 
-# Define the SimpleImplicitTauLeaping algorithm
 struct SimpleImplicitTauLeaping <: DiffEqBase.DEAlgorithm
     epsilon::Float64  # Error control parameter
     nc::Int          # Critical reaction threshold
     nstiff::Float64  # Stiffness threshold for switching
     delta::Float64   # Partial equilibrium threshold
 end
 
-# Default constructor
 SimpleImplicitTauLeaping(; epsilon=0.05, nc=10, nstiff=100.0, delta=0.05) = 
     SimpleImplicitTauLeaping(epsilon, nc, nstiff, delta)
 
-function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping; seed = nothing)
-    # Boilerplate setup
+# Compute stoichiometry matrix from c function
+function compute_stoichiometry(c, u, numjumps, p, t)
+    nu = zeros(Int, length(u), numjumps)
+    for j in 1:numjumps
+        counts = zeros(numjumps)
+        counts[j] = 1
+        du = similar(u)
+        c(du, u, p, t, counts, nothing)
+        nu[:, j] = round.(Int, du)
+    end
+    return nu
+end
+
+# Detect reversible reaction pairs
+function find_reversible_pairs(nu)
+    pairs = Vector{Tuple{Int,Int}}()
+    for j in 1:size(nu, 2)
+        for k in (j+1):size(nu, 2)
+            if nu[:, j] == -nu[:, k]
+                push!(pairs, (j, k))
+            end
+        end
+    end
+    return pairs
+end
+
+# Compute g_i (approximation from Cao et al., 2006)
+function compute_gi(u, nu, i, rate, rate_cache, p, t)
+    max_order = 1.0
+    for j in 1:size(nu, 2)
+        if abs(nu[i, j]) > 0
+            rate(rate_cache, u, p, t)
+            if rate_cache[j] > 0
+                order = 1.0
+                if sum(abs.(nu[:, j])) > abs(nu[i, j])
+                    order = 2.0
+                end
+                max_order = max(max_order, order)
+            end
+        end
+    end
+    return max_order
+end
+
+# Tau-selection for explicit method (Equation 8)
+function compute_tau_explicit(u, rate_cache, nu, p, t, epsilon, rate)
+    rate(rate_cache, u, p, t)
+    mu = zeros(length(u))
+    sigma2 = zeros(length(u))
+    tau = Inf
+    for i in 1:length(u)
+        for j in 1:size(nu, 2)
+            mu[i] += nu[i, j] * rate_cache[j]
+            sigma2[i] += nu[i, j]^2 * rate_cache[j]
+        end
+        gi = compute_gi(u, nu, i, rate, rate_cache, p, t)
+        bound = max(epsilon * u[i] / gi, 1.0)
+        mu_term = abs(mu[i]) > 0 ? bound / abs(mu[i]) : Inf
+        sigma_term = sigma2[i] > 0 ? bound^2 / sigma2[i] : Inf
+        tau = min(tau, mu_term, sigma_term)
+    end
+    return max(tau, 1e-10)
+end
+
+# Partial equilibrium check (Equation 13)
+function is_partial_equilibrium(rate_cache, j_plus, j_minus, delta)
+    a_plus = rate_cache[j_plus]
+    a_minus = rate_cache[j_minus]
+    return abs(a_plus - a_minus) <= delta * min(a_plus, a_minus)
+end
+
+# Tau-selection for implicit method (Equation 14)
+function compute_tau_implicit(u, rate_cache, nu, p, t, epsilon, rate, equilibrium_pairs, delta)
+    rate(rate_cache, u, p, t)
+    mu = zeros(length(u))
+    sigma2 = zeros(length(u))
+    non_equilibrium = trues(size(nu, 2))
+    for (j_plus, j_minus) in equilibrium_pairs
+        if is_partial_equilibrium(rate_cache, j_plus, j_minus, delta)
+            non_equilibrium[j_plus] = false
+            non_equilibrium[j_minus] = false
+        end
+    end
+    tau = Inf
+    for i in 1:length(u)
+        for j in 1:size(nu, 2)
+            if non_equilibrium[j]
+                mu[i] += nu[i, j] * rate_cache[j]
+                sigma2[i] += nu[i, j]^2 * rate_cache[j]
+            end
+        end
+        gi = compute_gi(u, nu, i, rate, rate_cache, p, t)
+        bound = max(epsilon * u[i] / gi, 1.0)
+        mu_term = abs(mu[i]) > 0 ? bound / abs(mu[i]) : Inf
+        sigma_term = sigma2[i] > 0 ? bound^2 / sigma2[i] : Inf
+        tau = min(tau, mu_term, sigma_term)
+    end
+    return max(tau, 1e-10)
+end
+
+# Identify critical reactions
+function identify_critical_reactions(u, rate_cache, nu, nc)
+    critical = falses(size(nu, 2))
+    for j in 1:size(nu, 2)
+        if rate_cache[j] > 0
+            Lj = Inf
+            for i in 1:length(u)
+                if nu[i, j] < 0
+                    Lj = min(Lj, floor(Int, u[i] / abs(nu[i, j])))
+                end
+            end
+            if Lj < nc
+                critical[j] = true
+            end
+        end
+    end
+    return critical
+end
+
+# Implicit tau-leaping step using NonlinearSolve
+function implicit_tau_step(u_prev, t_prev, tau, rate_cache, counts, nu, p, rate, numjumps)
+    # Define the nonlinear system: F(u_new) = u_new - u_prev - sum(nu_j * (counts_j - tau * a_j(u_prev) + tau * a_j(u_new))) = 0
+    function f(u_new, params)
+        rate_new = zeros(eltype(u_new), numjumps)
+        rate(rate_new, u_new, p, t_prev + tau)
+        residual = u_new - u_prev
+        for j in 1:numjumps
+            residual -= nu[:, j] * (counts[j] - tau * rate_cache[j] + tau * rate_new[j])
+        end
+        return residual
+    end
+    
+    # Initial guess
+    u_new = copy(u_prev)
+    
+    # Solve the nonlinear system
+    prob = NonlinearProblem(f, u_new, nothing)
+    sol = solve(prob, NewtonRaphson())
+    
+    # Check for convergence and numerical stability
+    if sol.retcode != ReturnCode.Success || any(isnan.(sol.u)) || any(isinf.(sol.u))
+        return round.(Int, max.(u_prev, 0.0))  # Revert to previous state
+    end
+    
+    return round.(Int, max.(sol.u, 0.0))
+end
+
+# Down-shifting condition (Equation 19)
+function use_down_shifting(t, tau_im, tau_ex, a0, t_end)
+    return a0 > 0 && t + tau_im >= t_end - 100 * (tau_ex + 1 / a0)
+end
+
+function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping; seed=nothing)
     @assert isempty(jump_prob.jump_callback.continuous_callbacks)
     @assert isempty(jump_prob.jump_callback.discrete_callbacks)
     prob = jump_prob.prob
@@ -103,160 +252,12 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping;
     delta = alg.delta
     t_end = tspan[2]
 
-    # Compute stoichiometry matrix from c function
-    function compute_stoichiometry(c, u, numjumps)
-        nu = zeros(Int, length(u), numjumps)
-        for j in 1:numjumps
-            counts = zeros(numjumps)
-            counts[j] = 1
-            du = similar(u)
-            c(du, u, p, t[1], counts, nothing)
-            nu[:, j] = round.(Int, du)
-        end
-        return nu
-    end
-    nu = compute_stoichiometry(c, u0, numjumps)
+    # Compute stoichiometry matrix
+    nu = compute_stoichiometry(c, u0, numjumps, p, t[1])
 
     # Detect reversible reaction pairs
-    function find_reversible_pairs(nu)
-        pairs = Vector{Tuple{Int,Int}}()
-        for j in 1:numjumps
-            for k in (j+1):numjumps
-                if nu[:, j] == -nu[:, k]
-                    push!(pairs, (j, k))
-                end
-            end
-        end
-        return pairs
-    end
     equilibrium_pairs = find_reversible_pairs(nu)
 
-    # Helper function to compute g_i (approximation from Cao et al., 2006)
-    function compute_gi(u, nu, i)
-        max_order = 1.0
-        for j in 1:numjumps
-            if abs(nu[i, j]) > 0
-                rate(rate_cache, u, p, t[end])
-                if rate_cache[j] > 0
-                    order = 1.0
-                    if sum(abs.(nu[:, j])) > abs(nu[i, j])
-                        order = 2.0
-                    end
-                    max_order = max(max_order, order)
-                end
-            end
-        end
-        return max_order
-    end
-    
-    # Tau-selection for explicit method (Equation 8)
-    function compute_tau_explicit(u, rate_cache, nu, p, t)
-        rate(rate_cache, u, p, t)
-        mu = zeros(length(u))
-        sigma2 = zeros(length(u))
-        tau = Inf
-        for i in 1:length(u)
-            for j in 1:numjumps
-                mu[i] += nu[i, j] * rate_cache[j]
-                sigma2[i] += nu[i, j]^2 * rate_cache[j]
-            end
-            gi = compute_gi(u, nu, i)
-            bound = max(epsilon * u[i] / gi, 1.0)
-            mu_term = abs(mu[i]) > 0 ? bound / abs(mu[i]) : Inf
-            sigma_term = sigma2[i] > 0 ? bound^2 / sigma2[i] : Inf
-            tau = min(tau, mu_term, sigma_term)
-        end
-        return max(tau, 1e-10)
-    end
-    
-    # Partial equilibrium check (Equation 13)
-    function is_partial_equilibrium(rate_cache, j_plus, j_minus)
-        a_plus = rate_cache[j_plus]
-        a_minus = rate_cache[j_minus]
-        return abs(a_plus - a_minus) <= delta * min(a_plus, a_minus)
-    end
-    
-    # Tau-selection for implicit method (Equation 14)
-    function compute_tau_implicit(u, rate_cache, nu, p, t)
-        rate(rate_cache, u, p, t)
-        mu = zeros(length(u))
-        sigma2 = zeros(length(u))
-        non_equilibrium = trues(numjumps)
-        for (j_plus, j_minus) in equilibrium_pairs
-            if is_partial_equilibrium(rate_cache, j_plus, j_minus)
-                non_equilibrium[j_plus] = false
-                non_equilibrium[j_minus] = false
-            end
-        end
-        tau = Inf
-        for i in 1:length(u)
-            for j in 1:numjumps
-                if non_equilibrium[j]
-                    mu[i] += nu[i, j] * rate_cache[j]
-                    sigma2[i] += nu[i, j]^2 * rate_cache[j]
-                end
-            end
-            gi = compute_gi(u, nu, i)
-            bound = max(epsilon * u[i] / gi, 1.0)
-            mu_term = abs(mu[i]) > 0 ? bound / abs(mu[i]) : Inf
-            sigma_term = sigma2[i] > 0 ? bound^2 / sigma2[i] : Inf
-            tau = min(tau, mu_term, sigma_term)
-        end
-        return max(tau, 1e-10)
-    end
-    
-    # Identify critical reactions
-    function identify_critical_reactions(u, rate_cache, nu)
-        critical = falses(numjumps)
-        for j in 1:numjumps
-            if rate_cache[j] > 0
-                Lj = Inf
-                for i in 1:length(u)
-                    if nu[i, j] < 0
-                        Lj = min(Lj, floor(Int, u[i] / abs(nu[i, j])))
-                    end
-                end
-                if Lj < nc
-                    critical[j] = true
-                end
-            end
-        end
-        return critical
-    end
-    
-    # Implicit tau-leaping step using NonlinearSolve
-    function implicit_tau_step(u_prev, t_prev, tau, rate_cache, counts, nu, p)
-        # Define the nonlinear system: F(u_new) = u_new - u_prev - sum(nu_j * (counts_j - tau * a_j(u_prev) + tau * a_j(u_new))) = 0
-        function f(u_new, params)
-            rate_new = similar(rate_cache, eltype(u_new))
-            rate(rate_new, u_new, p, t_prev + tau)
-            residual = u_new - u_prev
-            for j in 1:numjumps
-                residual -= nu[:, j] * (counts[j] - tau * rate_cache[j] + tau * rate_new[j])
-            end
-            return residual
-        end
-        
-        # Initial guess
-        u_new = copy(u_prev)
-        
-        # Solve the nonlinear system
-        prob = NonlinearProblem(f, u_new, nothing)
-        sol = solve(prob, NewtonRaphson())
-        
-        # Check for convergence and numerical stability
-        if sol.retcode != ReturnCode.Success || any(isnan.(sol.u)) || any(isinf.(sol.u))
-            return round.(Int, max.(u_prev, 0.0))  # Revert to previous state
-        end
-        
-        return round.(Int, max.(sol.u, 0.0))
-    end
-    
-    # Down-shifting condition (Equation 19)
-    function use_down_shifting(t, tau_im, tau_ex, a0, t_end)
-        return a0 > 0 && t + tau_im >= t_end - 100 * (tau_ex + 1 / a0)
-    end
-    
     # Main simulation loop
     while t[end] < t_end
         u_prev = u[end]
@@ -266,11 +267,11 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping;
         rate(rate_cache, u_prev, p, t_prev)
 
         # Identify critical reactions
-        critical = identify_critical_reactions(u_prev, rate_cache, nu)
+        critical = identify_critical_reactions(u_prev, rate_cache, nu, nc)
 
         # Compute tau values
-        tau_ex = compute_tau_explicit(u_prev, rate_cache, nu, p, t_prev)
-        tau_im = compute_tau_implicit(u_prev, rate_cache, nu, p, t_prev)
+        tau_ex = compute_tau_explicit(u_prev, rate_cache, nu, p, t_prev, epsilon, rate)
+        tau_im = compute_tau_implicit(u_prev, rate_cache, nu, p, t_prev, epsilon, rate, equilibrium_pairs, delta)
 
         # Compute critical propensity sum
         ac0 = sum(rate_cache[critical])
@@ -325,7 +326,7 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping;
                 end
             end
             if method == :implicit
-                u_new = implicit_tau_step(u_prev, t_prev, tau, rate_cache, counts, nu, p)
+                u_new = implicit_tau_step(u_prev, t_prev, tau, rate_cache, counts, nu, p, rate, numjumps)
             else
                 c(du, u_prev, p, t_prev, counts, nothing)
                 u_new = u_prev + du
@@ -352,7 +353,7 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping;
                 end
             end
             if method == :implicit && tau > tau_ex
-                u_new = implicit_tau_step(u_prev, t_prev, tau, rate_cache, counts, nu, p)
+                u_new = implicit_tau_step(u_prev, t_prev, tau, rate_cache, counts, nu, p, rate, numjumps)
             else
                 c(du, u_prev, p, t_prev, counts, nothing)
                 u_new = u_prev + du