test changes

Author: sivasathyaseeelan

test/regular_jumps.jl

@@ -1,35 +1,90 @@
 using JumpProcesses, DiffEqBase
-using Test, LinearAlgebra
-using StableRNGs, Plots
+using Test, LinearAlgebra, Statistics
+using StableRNGs
 rng = StableRNG(12345)
 
-# Parameters
-c1 = 1.0      # S1 -> 0
-c2 = 10.0     # S1 + S1 <- S2
-c3 = 1000.0   # S1 + S1 -> S2
-c4 = 0.1      # S2 -> S3
-p = (c1, c2, c3, c4)
-
-regular_rate_implicit = (out, u, p, t) -> begin
-    out[1] = p[1] * u[1]          # S1 -> 0
-    out[2] = p[2] * u[2]          # S1 + S1 <- S2
-    out[3] = p[3] * u[1] * (u[1] - 1) / 2  # S1 + S1 -> S2
-    out[4] = p[4] * u[2]          # S2 -> S3
-end
+Nsims = 8000
+
+# SIR model with influx
+let
+    β = 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]  # β*S*I (infection)
+        out[2] = p[2] * u[2]         # ν*I (recovery)
+        out[3] = p[3]                # influx_rate
+    end
+
+    regular_c = (dc, u, p, t, counts, mark) -> begin
+        dc .= 0.0
+        dc[1] = -counts[1] + counts[3]  # S: -infection + influx
+        dc[2] = counts[1] - counts[2]   # I: +infection - recovery
+        dc[3] = counts[2]               # R: +recovery
+    end
+
+    u0 = [999.0, 10.0, 0.0]  # S, I, R
+    tspan = (0.0, 250.0)
+
+    prob_disc = DiscreteProblem(u0, tspan, p)
+    rj = RegularJump(regular_rate, regular_c, 3)
+    jump_prob = JumpProblem(prob_disc, Direct(), rj)
+
+    sol = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(), EnsembleSerial(); trajectories = Nsims, dt = 1.0)
+    mean_simple = mean(sol.u[i][1,end] for i in 1:Nsims)
+
+    sol = solve(EnsembleProblem(jump_prob), SimpleImplicitTauLeaping(), EnsembleSerial(); trajectories = Nsims)
+    mean_implicit = mean(sol.u[i][1,end] for i in 1:Nsims)
+
+    sol = solve(EnsembleProblem(jump_prob), SimpleAdaptiveTauLeaping(), EnsembleSerial(); trajectories = Nsims)
+    mean_adaptive = mean(sol.u[i][1,end] for i in 1:Nsims)
 
-regular_c_implicit = (dc, u, p, t, counts, mark) -> begin
-    dc .= 0.0
-    dc[1] = -counts[1] - 2 * counts[3] + 2 * counts[2]  # S1: -decay - 2*forward + 2*backward
-    dc[2] = counts[3] - counts[2] - counts[4]           # S2: +forward - backward - decay
-    dc[3] = counts[4]                                   # S3: +decay
+    @test isapprox(mean_simple, mean_implicit, rtol=0.05)
+    @test isapprox(mean_simple, mean_adaptive, rtol=0.05)
 end
 
-u0 = [10000.0, 0.0, 0.0]  # S1, S2, S3
-tspan = (0.0, 4.0)
 
-# Create JumpProblem with proper parameter passing
-prob_disc = DiscreteProblem(u0, tspan, p)
-rj = RegularJump(regular_rate_implicit, regular_c_implicit, 4)
-jump_prob = JumpProblem(prob_disc, Direct(), rj; rng=StableRNG(12345))
-sol = solve(jump_prob, SimpleAdaptiveTauLeaping())
-plot(sol)
+# SEIR model with exposed compartment
+let
+    β = 0.3 / 1000.0
+    σ = 0.2
+    ν = 0.01
+    p = (β, σ, ν)
+
+    regular_rate = (out, u, p, t) -> begin
+        out[1] = p[1] * u[1] * u[3]  # β*S*I (infection)
+        out[2] = p[2] * u[2]         # σ*E (progression)
+        out[3] = p[3] * u[3]         # ν*I (recovery)
+    end
+
+    regular_c = (dc, u, p, t, counts, mark) -> begin
+        dc .= 0.0
+        dc[1] = -counts[1]           # S: -infection
+        dc[2] = counts[1] - counts[2] # E: +infection - progression
+        dc[3] = counts[2] - counts[3] # I: +progression - recovery
+        dc[4] = counts[3]            # R: +recovery
+    end
+
+    # Initial state
+    u0 = [999.0, 0.0, 10.0, 0.0]  # S, E, I, R
+    tspan = (0.0, 250.0)
+
+    # Create JumpProblem
+    prob_disc = DiscreteProblem(u0, tspan, p)
+    rj = RegularJump(regular_rate, regular_c, 3)
+    jump_prob = JumpProblem(prob_disc, Direct(), rj; rng=StableRNG(12345))
+
+    sol = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(), EnsembleSerial(); trajectories = Nsims, dt = 1.0)
+    mean_simple = mean(sol.u[i][end,end] for i in 1:Nsims)
+
+    sol = solve(EnsembleProblem(jump_prob), SimpleImplicitTauLeaping(), EnsembleSerial(); trajectories = Nsims)
+    mean_implicit = mean(sol.u[i][end,end] for i in 1:Nsims)
+
+    sol = solve(EnsembleProblem(jump_prob), SimpleAdaptiveTauLeaping(), EnsembleSerial(); trajectories = Nsims)
+    mean_adaptive = mean(sol.u[i][end,end] for i in 1:Nsims)
+
+    @test isapprox(mean_simple, mean_implicit, rtol=0.05)
+    @test isapprox(mean_simple, mean_adaptive, rtol=0.05)
+end