Merge pull request #520 from isaacsas/add_leap_aggregator

Add leap aggregator

Author: isaacsas

Project.toml

@@ -45,7 +45,7 @@ FunctionWrappers = "1.1"
 Graphs = "1.11"
 KernelAbstractions = "0.9"
 LinearAlgebra = "1"
-LinearSolve = "2, 3"
+LinearSolve = "3"
 OrdinaryDiffEq = "6"
 Pkg = "1"
 PoissonRandom = "0.4"

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, alg) ||
+        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/JumpProcesses.jl

@@ -110,6 +110,13 @@ export ExtendedJumpArray
 include("variable_rate.jl")
 export VariableRateAggregator, VR_FRM, VR_Direct, VR_DirectFW
 
+"""
+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
+
 # core problem and timestepping
 include("problem.jl")
 export JumpProblem, SplitCoupledJumpProblem

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) 

src/simple_regular_solve.jl

@@ -1,12 +1,24 @@
 struct SimpleTauLeaping <: DiffEqBase.DEAlgorithm end
 
+function validate_pure_leaping_inputs(jump_prob::JumpProblem, alg)
+    if !(jump_prob.aggregator isa PureLeaping)
+        @warn "When using $alg, please pass PureLeaping() as the aggregator to the \
+        JumpProblem, i.e. call JumpProblem(::DiscreteProblem, PureLeaping(),...). \
+        Passing $(jump_prob.aggregator) is deprecated and will be removed in the next breaking release."
+    end
+    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, alg) ||
+        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 +74,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)

test/regular_jumps.jl

@@ -8,23 +8,7 @@ function regular_rate(out, u, p, t)
     out[2] = 0.01u[2]
 end
 
-function regular_c(dc, u, p, t, mark)
-    dc[1, 1] = -1
-    dc[2, 1] = 1
-    dc[2, 2] = -1
-    dc[3, 2] = 1
-end
-
-dc = zeros(3, 2)
-
-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)
-sol = solve(jump_prob, SimpleTauLeaping(); dt = 1.0)
-
-const _dc = zeros(3, 2)
+const dc = zeros(3, 2)
 dc[1, 1] = -1
 dc[2, 1] = 1
 dc[2, 2] = -1
@@ -37,5 +21,92 @@ 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
+@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, tspan, p)
+    
+    # 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
+    
+    rj_dc = zeros(3, 1)
+    rj_dc[1, 1] = -1
+    rj_dc[3, 1] = 1
+    
+    function rj_c(du, u, p, t, counts, mark)
+        mul!(du, rj_dc, counts)
+    end
+    
+    regj = RegularJump(rj_rate, rj_c, 1)
+    
+    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(; massaction_jumps = maj, constant_jumps = (crj,), 
+        variable_jumps = (vrj,), regular_jumps = 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
+    maj_params = MassActionJump(reactant_stoich, net_stoich; param_idxs = [1, 2])
+    jp_params = JumpProblem(prob, PureLeaping(), JumpSet(maj_params))
+    scaled_rates = [p[1], p[2]/2]
+    @test jp_params.massaction_jump.scaled_rates == scaled_rates
+end