@@ -56,12 +56,11 @@ struct ImplicitTauLeaping <: DiffEqBase.DEAlgorithm
5656 nc:: Int # Critical reaction threshold
5757 nstiff:: Float64 # Stiffness threshold for switching
5858 delta:: Float64 # Partial equilibrium threshold
59- equilibrium_pairs:: Vector{Tuple{Int,Int}} # Reversible reaction pairs
6059end
6160
6261# Default constructor
63- ImplicitTauLeaping (; epsilon= 0.05 , nc= 10 , nstiff= 100.0 , delta= 0.05 , equilibrium_pairs = [( 1 , 2 )] ) =
64- ImplicitTauLeaping (epsilon, nc, nstiff, delta, equilibrium_pairs )
62+ ImplicitTauLeaping (; epsilon= 0.05 , nc= 10 , nstiff= 100.0 , delta= 0.05 ) =
63+ ImplicitTauLeaping (epsilon, nc, nstiff, delta)
6564
6665function DiffEqBase. solve (jump_prob:: JumpProblem , alg:: ImplicitTauLeaping ; seed= nothing )
6766 # Boilerplate setup
@@ -91,7 +90,6 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::ImplicitTauLeaping; seed=
9190 nc = alg. nc
9291 nstiff = alg. nstiff
9392 delta = alg. delta
94- equilibrium_pairs = alg. equilibrium_pairs
9593 t_end = tspan[2 ]
9694
9795 # Compute stoichiometry matrix from c function
@@ -108,17 +106,30 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::ImplicitTauLeaping; seed=
108106 end
109107 nu = compute_stoichiometry (c, u0, numjumps)
110108
109+ # Detect reversible reaction pairs
110+ function find_reversible_pairs (nu)
111+ pairs = Vector {Tuple{Int,Int}} ()
112+ for j in 1 : numjumps
113+ for k in (j+ 1 ): numjumps
114+ if nu[:, j] == - nu[:, k]
115+ push! (pairs, (j, k))
116+ end
117+ end
118+ end
119+ return pairs
120+ end
121+ equilibrium_pairs = find_reversible_pairs (nu)
122+
111123 # Helper function to compute g_i (approximation from Cao et al., 2006)
112124 function compute_gi (u, nu, i)
113- # Simplified g_i: highest order of reaction involving species i
114125 max_order = 1.0
115126 for j in 1 : numjumps
116127 if abs (nu[i, j]) > 0
117- # Approximate reaction order based on propensity (heuristic)
118128 rate (rate_cache, u, p, t[end ])
119129 if rate_cache[j] > 0
120- order = 1.0 # Assume first-order for simplicity
121- if j == 1 # For SIR infection (S*I), assume second-order
130+ order = 1.0
131+ # Heuristic: if reaction involves multiple species, assume higher order
132+ if sum (abs .(nu[:, j])) > abs (nu[i, j])
122133 order = 2.0
123134 end
124135 max_order = max (max_order, order)
@@ -218,14 +229,15 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::ImplicitTauLeaping; seed=
218229 if norm (residual) < tol
219230 break
220231 end
221- # Approximate Jacobian
232+ # Approximate Jacobian (diagonal approximation for simplicity)
222233 J = Diagonal (ones (length (u_new)))
223234 for j in 1 : numjumps
224235 for i in 1 : length (u_new)
225- if j == 1 && i in [1 , 2 ] # Infection: β*S*I
226- J[i, i] += nu[i, j] * tau * p[1 ] * (i == 1 ? u_new[2 ] : u_new[1 ])
227- elseif j == 2 && i == 2 # Recovery: ν*I
228- J[i, i] += nu[i, j] * tau * p[2 ]
236+ # Heuristic derivative: assume linear or quadratic propensity
237+ rate (rate_new, u_new, p, t_prev + tau)
238+ if rate_new[j] > 0
239+ # Estimate derivative based on stoichiometry
240+ J[i, i] += nu[i, j] * tau * rate_new[j] / max (u_new[i], 1.0 )
229241 end
230242 end
231243 end
@@ -266,7 +278,7 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::ImplicitTauLeaping; seed=
266278 method = use_implicit ? :implicit : :explicit
267279
268280 # Check if tau1 is too small
269- if tau1 < 10 / a0
281+ if a0 > 0 && tau1 < 10 / a0
270282 # Use SSA for a few steps
271283 steps = method == :implicit ? 10 : 100
272284 for _ in 1 : steps
0 commit comments