basic version of inplicit tau leap is done

Author: sivasathyaseeelan

src/JumpProcesses.jl

@@ -20,6 +20,8 @@ using Base.Threads: Threads, @threads
 using Base.FastMath: add_fast
 using Setfield: @set, @set!
 
+using SimpleNonlinearSolve
+
 # Import functions we extend from Base
 import Base: size, getindex, setindex!, length, similar, show, merge!, merge
 

src/simple_regular_solve.jl

@@ -69,7 +69,7 @@ end
 SimpleImplicitTauLeaping(; epsilon=0.05) = SimpleImplicitTauLeaping(epsilon)
 
 function compute_hor(nu)
-    hor = zeros(Int, size(nu, 2))
+    hor = zeros(Int64, size(nu, 2))
     for j in 1:size(nu, 2)
         hor[j] = sum(abs.(nu[:, j])) > maximum(abs.(nu[:, j])) ? 2 : 1
     end
@@ -88,8 +88,8 @@ end
 
 function compute_tau_explicit(u, rate_cache, nu, hor, p, t, epsilon, rate)
     rate(rate_cache, u, p, t)
-    mu = zeros(length(u))
-    sigma2 = zeros(length(u))
+    mu = zeros(Float64, length(u))
+    sigma2 = zeros(Float64, length(u))
     tau = Inf
     for i in 1:length(u)
         for j in 1:size(nu, 2)
@@ -111,21 +111,20 @@ function compute_tau_implicit(u, rate_cache, nu, p, t, rate)
     for i in 1:length(u)
         sum_nu_a = 0.0
         for j in 1:size(nu, 2)
-            if nu[i, j] < 0  # Only sum negative stoichiometry
+            if nu[i, j] < 0
                 sum_nu_a += abs(nu[i, j]) * rate_cache[j]
             end
         end
-        if sum_nu_a > 0 && u[i] > 0  # Avoid division by zero
+        if sum_nu_a > 0 && u[i] > 0
             tau = min(tau, u[i] / sum_nu_a)
         end
     end
     return tau
 end
 
 function implicit_tau_step(u_prev, t_prev, tau, rate_cache, counts, nu, p, rate, numjumps)
-    # Nonlinear system: F(u_new) = u_new - u_prev - sum(nu_j * (k_j - tau * (a_j(u_prev) - a_j(u_new)))) = 0
-    function f(u_new)
-        rate_new = zeros(Float64, numjumps)
+    function f(u_new, p)
+        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
@@ -134,41 +133,14 @@ function implicit_tau_step(u_prev, t_prev, tau, rate_cache, counts, nu, p, rate,
         return residual
     end
 
-    # Numerical Jacobian
-    function compute_jacobian(u_new)
-        n = length(u_new)
-        J = zeros(Float64, n, n)
-        h = 1e-6
-        f_u = f(u_new)
-        for j in 1:n
-            u_pert = copy(u_new)
-            u_pert[j] += h
-            f_pert = f(u_pert)
-            J[:, j] = (f_pert - f_u) / h
-        end
-        return J
-    end
+    u_new = float.(u_prev + sum(nu[:, j] * counts[j] for j in 1:numjumps))
+    prob = NonlinearProblem{false}(f, u_new, p)
+    sol = solve(prob, SimpleNewtonRaphson(), abstol=1e-6, maxiters=100)
 
-    # Inline Newton-Raphson
-    u_new = float.(u_prev + sum(nu[:, j] * counts[j] for j in 1:numjumps))  # Initial guess: explicit step
-    tol = 1e-6
-    maxiters = 100
-    for iter in 1:maxiters
-        F = f(u_new)
-        if norm(F) < tol
-            return round.(Int, max.(u_new, 0.0))  # Converged
-        end
-        J = compute_jacobian(u_new)
-        if abs(det(J)) < 1e-10  # Check for singular Jacobian
-            return nothing
-        end
-        delta = J \ F
-        u_new -= delta
-        if any(isnan.(u_new)) || any(isinf.(u_new))
-            return nothing
-        end
+    if sol.retcode != ReturnCode.Success || any(isnan.(sol.u)) || any(isinf.(sol.u))
+        return nothing
     end
-    return nothing  # Failed to converge
+    return round.(Int64, max.(sol.u, 0.0))
 end
 
 function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping; seed=nothing, dtmin=1e-10, saveat=nothing)
@@ -211,7 +183,6 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping;
     while t[end] < t_end
         u_prev = u[end]
         t_prev = t[end]
-        # Recompute stoichiometry
         for j in 1:numjumps
             fill!(counts_temp, 0)
             counts_temp[j] = 1
@@ -221,11 +192,10 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping;
         rate(rate_cache, u_prev, p, t_prev)
         tau_prime = compute_tau_explicit(u_prev, rate_cache, nu, hor, p, t_prev, epsilon, rate)
         tau_double_prime = compute_tau_implicit(u_prev, rate_cache, nu, p, t_prev, rate)
-        # Cao et al. (2007): Use tau_prime for explicit, tau_double_prime for implicit
         use_implicit = false
-        tau = tau_prime  # Default to explicit
-        if tau_double_prime < tau_prime && any(u_prev .< 10)  # Implicit if populations are low
-            tau = tau_double_prime
+        tau = tau_prime
+        if any(u_prev .< 10)
+            tau = min(tau_double_prime, tau_prime) # Tighter cap for accuracy
             use_implicit = true
         end
         tau = max(tau, dtmin)
@@ -241,11 +211,11 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleImplicitTauLeaping;
         if use_implicit
             u_new = implicit_tau_step(u_prev, t_prev, tau, rate_cache, counts, nu, p, rate, numjumps)
             if u_new === nothing || any(u_new .< 0)
-                tau /= 2  # Halve tau if implicit fails or produces negative populations
+                tau /= 2
                 continue
             end
         elseif any(u_new .< 0)
-            tau /= 2  # Halve tau if explicit produces negative populations
+            tau /= 2
             continue
         end
         u_new = max.(u_new, 0)

test/regular_jumps.jl

@@ -5,76 +5,31 @@ rng = StableRNG(12345)
 
 Nsims = 10
 
-# @testset "SIR Model Correctness" begin
-    β = 0.1 / 1000.0
-    ν = 0.01
-    influx_rate = 1.0
-    p = (β, ν, influx_rate)
 
-    u0 = [999, 10, 0]  # Integer initial conditions
-    tspan = (0.0, 250.0)
-    prob_disc = DiscreteProblem(u0, tspan, p)
+β = 0.1 / 1000.0
+ν = 0.01
+influx_rate = 1.0
+p = (β, ν, influx_rate)
 
-    regular_rate = (out, u, p, t) -> begin
-        out[1] = p[1] * u[1] * u[2]
-        out[2] = p[2] * u[2]
-        out[3] = p[3]
-    end
-    regular_c = (dc, u, p, t, counts, mark) -> begin
-        dc .= 0
-        dc[1] = -counts[1] + counts[3]
-        dc[2] = counts[1] - counts[2]
-        dc[3] = counts[2]
-    end
-    rj = RegularJump(regular_rate, regular_c, 3)
-    jump_prob_tau = JumpProblem(prob_disc, Direct(), rj; rng=StableRNG(12345))
+function regular_rate(out, u, p, t)
+    out[1] = p[1] * u[1] * u[2]
+    out[2] = p[2] * u[2]
+    out[3] = p[3]
+end
 
-    sol_implicit = solve(EnsembleProblem(jump_prob_tau), SimpleImplicitTauLeaping(), EnsembleSerial(); trajectories=Nsims)
+regular_c = (dc, u, p, t, counts, mark) -> begin
+    dc .= 0
+    dc[1] = -counts[1] + counts[3]
+    dc[2] = counts[1] - counts[2]
+    dc[3] = counts[2]
+end
 
-# end
+u0 = [999, 5, 0]
+tspan = (0.0, 250.0)
+prob_disc = DiscreteProblem(u0, tspan, p)
 
-# @testset "SEIR Model Correctness" begin
-    β = 0.3 / 1000.0
-    σ = 0.2
-    ν = 0.01
-    p = (β, σ, ν)
 
-    rate1(u, p, t) = p[1] * u[1] * u[3]
-    rate2(u, p, t) = p[2] * u[2]
-    rate3(u, p, t) = p[3] * u[3]
-    affect1!(integrator) = (integrator.u[1] -= 1; integrator.u[2] += 1; nothing)
-    affect2!(integrator) = (integrator.u[2] -= 1; integrator.u[3] += 1; nothing)
-    affect3!(integrator) = (integrator.u[3] -= 1; integrator.u[4] += 1; nothing)
-    jumps = (ConstantRateJump(rate1, affect1!), ConstantRateJump(rate2, affect2!), ConstantRateJump(rate3, affect3!))
+rj = RegularJump(regular_rate, regular_c, 3)
+jump_prob_tau = JumpProblem(prob_disc, Direct(), rj; rng=StableRNG(12345))
 
-    u0 = [999, 0, 10, 0]  # Integer initial conditions
-    tspan = (0.0, 250.0)
-    prob_disc = DiscreteProblem(u0, tspan, p)
-    jump_prob = JumpProblem(prob_disc, Direct(), jumps...; rng=StableRNG(12345))
-
-    sol_direct = solve(EnsembleProblem(jump_prob), SSAStepper(), EnsembleSerial(); trajectories=Nsims)
-
-    regular_rate = (out, u, p, t) -> begin
-        out[1] = p[1] * u[1] * u[3]
-        out[2] = p[2] * u[2]
-        out[3] = p[3] * u[3]
-    end
-    regular_c = (dc, u, p, t, counts, mark) -> begin
-        dc .= 0
-        dc[1] = -counts[1]
-        dc[2] = counts[1] - counts[2]
-        dc[3] = counts[2] - counts[3]
-        dc[4] = counts[3]
-    end
-    rj = RegularJump(regular_rate, regular_c, 3)
-    jump_prob_tau = JumpProblem(prob_disc, Direct(), rj; rng=StableRNG(12345))
-
-    sol_implicit = solve(EnsembleProblem(jump_prob_tau), SimpleImplicitTauLeaping(), EnsembleSerial(); trajectories=Nsims, saveat=1.0)
-
-    t_points = 0:1.0:250.0
-    mean_direct_R = [mean(sol_direct[i](t)[4] for i in 1:Nsims) for t in t_points]
-    mean_implicit_R = [mean(sol_implicit[i](t)[4] for i in 1:Nsims) for t in t_points]
-
-    max_error_implicit = maximum(abs.(mean_direct_R .- mean_implicit_R))
-    # @test max_error_implicit < 0.01 * mean(mean_direct_R)
-# end
+sol_implicit = solve(EnsembleProblem(jump_prob_tau), SimpleImplicitTauLeaping(), EnsembleSerial(); trajectories=Nsims, saveat=1.0)