update GPU code too

Author: isaacsas

ext/JumpProcessesKernelAbstractionsExt.jl

@@ -22,9 +22,9 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem,
 
     jump_prob = ensembleprob.prob
 
-    # boilerplate from SimpleTauLeaping method
-    @assert isempty(jump_prob.jump_callback.continuous_callbacks) # still needs to be a regular jump
-    @assert isempty(jump_prob.jump_callback.discrete_callbacks)
+    # Validate that this is a PureLeaping JumpProblem
+    validate_pure_leaping_inputs(jump_prob) ||
+        error("SimpleTauLeaping can only be used with PureLeaping JumpProblems with only non-RegularJumps.")
     prob = jump_prob.prob
 
     probs = [remake(jump_prob) for _ in 1:trajectories]

src/simple_regular_solve.jl

@@ -1,12 +1,21 @@
 struct SimpleTauLeaping <: DiffEqBase.DEAlgorithm end
 
+function validate_pure_leaping_inputs(jump_prob::JumpProblem)
+    jump_prob.aggregator isa PureLeaping &&
+        isempty(jump_prob.jump_callback.continuous_callbacks) &&
+        isempty(jump_prob.jump_callback.discrete_callbacks) &&
+        isempty(jump_prob.constant_jumps) &&
+        isempty(jump_prob.variable_jumps) &&
+        get_num_majumps(jump_prob.massaction_jump) == 0 &&
+        jump_prob.regular_jump !== nothing    
+end
+
 function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping;
         seed = nothing,
         dt = error("dt is required for SimpleTauLeaping."))
 
-    # boilerplate from SimpleTauLeaping method
-    @assert isempty(jump_prob.jump_callback.continuous_callbacks) # still needs to be a regular jump
-    @assert isempty(jump_prob.jump_callback.discrete_callbacks)
+    validate_pure_leaping_inputs(jump_prob) ||
+        error("SimpleTauLeaping can only be used with PureLeaping JumpProblems with only non-RegularJumps.")
     prob = jump_prob.prob
     rng = DEFAULT_RNG
     (seed !== nothing) && seed!(rng, seed)
@@ -62,4 +71,3 @@ end
 function EnsembleGPUKernel()
     EnsembleGPUKernel(nothing, 0.0)
 end
-

test/gpu/regular_jumps.jl

@@ -31,7 +31,7 @@ let
 
     prob_disc = DiscreteProblem(u0, tspan, p)
     rj = RegularJump(regular_rate, regular_c, 3)
-    jump_prob = JumpProblem(prob_disc, Direct(), rj)
+    jump_prob = JumpProblem(prob_disc, PureLeaping(), rj)
 
     sol = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(),
         EnsembleGPUKernel(CUDABackend()); trajectories = Nsims, dt = 1.0)
@@ -72,7 +72,7 @@ let
     # 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))
+    jump_prob = JumpProblem(prob_disc, PureLeaping(), rj; rng = StableRNG(12345))
 
     sol = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(),
         EnsembleGPUKernel(CUDABackend()); trajectories = Nsims, dt = 1.0)
@@ -84,3 +84,33 @@ let
 
     @test isapprox(mean_kernel, mean_serial, rtol = 0.05)
 end
+
+# Test PureLeaping validation for GPU
+@testset "PureLeaping GPU Validation" begin
+    β = 0.1 / 1000.0
+    ν = 0.01
+    p = (β, ν)
+    
+    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)
+    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] # I: +infection - recovery
+        dc[3] = counts[2]            # R: +recovery
+    end
+    
+    u0 = [999.0, 10.0, 0.0]  # S, I, R
+    tspan = (0.0, 50.0)
+    prob_disc = DiscreteProblem(u0, tspan, p)
+    rj = RegularJump(regular_rate, regular_c, 2)
+    
+    # This should fail - Direct aggregator with SimpleTauLeaping
+    jump_prob_direct = JumpProblem(prob_disc, Direct(), rj)
+    @test_throws ErrorException solve(EnsembleProblem(jump_prob_direct), 
+        SimpleTauLeaping(), EnsembleGPUKernel(CUDABackend()); 
+        trajectories = 10, dt = 1.0)
+end

test/regular_jumps.jl

@@ -21,7 +21,7 @@ rj = RegularJump(regular_rate, regular_c, dc; constant_c = true)
 jumps = JumpSet(rj)
 
 prob = DiscreteProblem([999.0, 1.0, 0.0], (0.0, 250.0))
-jump_prob = JumpProblem(prob, Direct(), rj; rng = rng)
+jump_prob = JumpProblem(prob, PureLeaping(), rj; rng)
 sol = solve(jump_prob, SimpleTauLeaping(); dt = 1.0)
 
 const _dc = zeros(3, 2)
@@ -37,7 +37,7 @@ end
 rj = RegularJump(regular_rate, regular_c, 2)
 jumps = JumpSet(rj)
 prob = DiscreteProblem([999, 1, 0], (0.0, 250.0))
-jump_prob = JumpProblem(prob, Direct(), rj; rng = rng)
+jump_prob = JumpProblem(prob, PureLeaping(), rj; rng)
 sol = solve(jump_prob, SimpleTauLeaping(); dt = 1.0)
 
 # Test PureLeaping aggregator functionality