PureLeaping aggregator

Author: isaacsas

src/JumpProcesses.jl

@@ -121,6 +121,13 @@ include("SSA_stepper.jl")
 export SSAStepper
 
 # leaping: 
+"""
+Aggregator to indicate that individual jumps should also be handled via the leaping
+algorithm that is passed to solve.
+"""
+struct PureLeaping <: AbstractAggregatorAlgorithm end
+export PureLeaping
+
 include("simple_regular_solve.jl")
 export SimpleTauLeaping, EnsembleGPUKernel
 

src/problem.jl

@@ -302,6 +302,48 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS
         jump_cbs, crjs, cvrjs, jumps.regular_jump, maj, rng, solkwargs)
 end
 
+# Special dispatch for PureLeaping aggregator - bypasses all aggregation
+function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet;
+        save_positions = prob isa DiffEqBase.AbstractDiscreteProblem ?
+                         (false, true) : (true, true),
+        rng = DEFAULT_RNG, scale_rates = true, useiszero = true,
+        spatial_system = nothing, hopping_constants = nothing,
+        callback = nothing, kwargs...)
+
+    # Validate no spatial systems (not currently supported)
+    (spatial_system !== nothing || hopping_constants !== nothing) &&
+        error("PureLeaping does not currently support spatial problems.")
+
+    # Initialize the MassActionJump rate constants with the user parameters
+    if using_params(jumps.massaction_jump)
+        rates = jumps.massaction_jump.param_mapper(prob.p)
+        maj = MassActionJump(rates, jumps.massaction_jump.reactant_stoch,
+            jumps.massaction_jump.net_stoch,
+            jumps.massaction_jump.param_mapper; scale_rates = scale_rates,
+            useiszero = useiszero,
+            nocopy = true)
+    else
+        maj = jumps.massaction_jump
+    end
+
+    # For PureLeaping, all jumps are handled by the tau-leaping solver
+    # No discrete jump aggregation or variable rate callbacks are created
+    disc_agg = nothing
+    jump_cbs = CallbackSet()
+    
+    # Store all jump types for access by tau-leaping solver
+    crjs = jumps.constant_jumps
+    vrjs = jumps.variable_jumps
+    
+    iip = isinplace_jump(prob, jumps.regular_jump)
+    solkwargs = make_kwarg(; callback)
+
+    JumpProblem{iip, typeof(prob), typeof(aggregator), typeof(jump_cbs), 
+        typeof(disc_agg), typeof(crjs), typeof(vrjs), typeof(jumps.regular_jump),
+        typeof(maj), typeof(rng), typeof(solkwargs)}(prob, aggregator, disc_agg,
+        jump_cbs, crjs, vrjs, jumps.regular_jump, maj, rng, solkwargs)
+end
+
 aggregator(jp::JumpProblem{iip, P, A}) where {iip, P, A} = A
 
 @inline function extend_tstops!(tstops, jp::JumpProblem) 

test/regular_jumps.jl

@@ -39,3 +39,87 @@ jumps = JumpSet(rj)
 prob = DiscreteProblem([999, 1, 0], (0.0, 250.0))
 jump_prob = JumpProblem(prob, Direct(), rj; rng = rng)
 sol = solve(jump_prob, SimpleTauLeaping(); dt = 1.0)
+
+# Test PureLeaping aggregator functionality
+@testset "PureLeaping Aggregator Tests" begin
+    # Test with MassActionJump
+    u0 = [10, 5, 0]
+    tspan = (0.0, 10.0)
+    p = [0.1, 0.2]
+    prob = DiscreteProblem(u0, p, tspan)
+    
+    # Create MassActionJump
+    reactant_stoich = [[1 => 1], [1 => 2]]
+    net_stoich = [[1 => -1, 2 => 1], [1 => -2, 3 => 1]]
+    rates = [0.1, 0.05]
+    maj = MassActionJump(rates, reactant_stoich, net_stoich)
+    
+    # Test PureLeaping JumpProblem creation
+    jp_pure = JumpProblem(prob, PureLeaping(), JumpSet(maj))
+    @test jp_pure.aggregator isa PureLeaping
+    @test jp_pure.discrete_jump_aggregation === nothing
+    @test jp_pure.massaction_jump !== nothing
+    @test length(jp_pure.jump_callback.discrete_callbacks) == 0
+    
+    # Test with ConstantRateJump
+    rate(u, p, t) = p[1] * u[1]
+    affect!(integrator) = (integrator.u[1] -= 1; integrator.u[3] += 1)
+    crj = ConstantRateJump(rate, affect!)
+    
+    jp_pure_crj = JumpProblem(prob, PureLeaping(), JumpSet(crj))
+    @test jp_pure_crj.aggregator isa PureLeaping
+    @test jp_pure_crj.discrete_jump_aggregation === nothing
+    @test length(jp_pure_crj.constant_jumps) == 1
+    
+    # Test with VariableRateJump
+    vrate(u, p, t) = t * p[1] * u[1]
+    vaffect!(integrator) = (integrator.u[1] -= 1; integrator.u[3] += 1)
+    vrj = VariableRateJump(vrate, vaffect!)
+    
+    jp_pure_vrj = JumpProblem(prob, PureLeaping(), JumpSet(vrj))
+    @test jp_pure_vrj.aggregator isa PureLeaping
+    @test jp_pure_vrj.discrete_jump_aggregation === nothing
+    @test length(jp_pure_vrj.variable_jumps) == 1
+    
+    # Test with RegularJump
+    function rj_rate(out, u, p, t)
+        out[1] = p[1] * u[1]
+    end
+    
+    function rj_c(dc, u, p, t, mark)
+        dc[1, 1] = -1
+        dc[3, 1] = 1
+    end
+    
+    rj_dc = zeros(3, 1)
+    regj = RegularJump(rj_rate, rj_c, rj_dc; constant_c = true)
+    
+    jp_pure_regj = JumpProblem(prob, PureLeaping(), JumpSet(regj))
+    @test jp_pure_regj.aggregator isa PureLeaping
+    @test jp_pure_regj.discrete_jump_aggregation === nothing
+    @test jp_pure_regj.regular_jump !== nothing
+    
+    # Test mixed jump types
+    mixed_jumps = JumpSet(maj, crj, vrj, regj)
+    jp_pure_mixed = JumpProblem(prob, PureLeaping(), mixed_jumps)
+    @test jp_pure_mixed.aggregator isa PureLeaping
+    @test jp_pure_mixed.discrete_jump_aggregation === nothing
+    @test jp_pure_mixed.massaction_jump !== nothing
+    @test length(jp_pure_mixed.constant_jumps) == 1
+    @test length(jp_pure_mixed.variable_jumps) == 1
+    @test jp_pure_mixed.regular_jump !== nothing
+    
+    # Test spatial system error
+    spatial_sys = CartesianGrid((2, 2))
+    hopping_consts = [1.0]
+    @test_throws ErrorException JumpProblem(prob, PureLeaping(), JumpSet(maj); 
+                                          spatial_system = spatial_sys)
+    @test_throws ErrorException JumpProblem(prob, PureLeaping(), JumpSet(maj); 
+                                          hopping_constants = hopping_consts)
+    
+    # Test MassActionJump with parameter mapping
+    param_mapper = MassActionJumpParamMapper([1, 2])
+    maj_params = MassActionJump(reactant_stoich, net_stoich, param_mapper)
+    jp_params = JumpProblem(prob, PureLeaping(), JumpSet(maj_params))
+    @test jp_params.massaction_jump.scaled_rates == p
+end