From e0498e315f25727b4d674a18106ac7cf7b065b16 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Mon, 23 Feb 2026 22:05:31 -0500 Subject: [PATCH 01/52] Move RNG from aggregators/JumpProblem to integrator Make the integrator the single source of truth for RNG (Phase 1 Step 4). Remove rng field from all aggregator structs, JumpProblem, and variable-rate callback structs. All code now reads get_rng(integrator) at runtime. Key changes: - Remove RNG type param from AbstractSSAJumpAggregator and all concrete aggregators - Remove rng from build_jump_aggregation and all aggregate() signatures - Add rng field to SSAIntegrator with has_rng/get_rng/set_rng! interface - Add resolve_rng utility for rng/seed/fallback priority resolution - Remove rng from VR_FRMEventCallback and VR_DirectEventCache - JumpProblem rng kwarg now forwards to solver via solkwargs - Simplify solve.jl: remove _derive_jump_seed, simplify __jump_init - Update SimpleTauLeaping/SimpleExplicitTauLeaping to use rng kwarg - Add rng_kwarg_tests.jl covering all three solve pathways - Update docs with RNG control section and FAQ entry - Bump version to 9.23.0, SciMLBase compat to 2.144 Co-Authored-By: Claude Opus 4.5 --- Project.toml | 7 +- docs/src/api.md | 102 ++++++++- docs/src/faq.md | 32 ++- docs/src/tutorials/simple_poisson_process.md | 9 +- src/JumpProcesses.jl | 2 +- src/SSA_stepper.jl | 31 ++- src/aggregators/ccnrm.jl | 43 ++-- src/aggregators/coevolve.jl | 40 ++-- src/aggregators/direct.jl | 34 +-- src/aggregators/directcr.jl | 24 +- src/aggregators/frm.jl | 40 ++-- src/aggregators/nrm.jl | 37 ++-- src/aggregators/rdirect.jl | 22 +- src/aggregators/rssa.jl | 22 +- src/aggregators/rssacr.jl | 20 +- src/aggregators/sortingdirect.jl | 24 +- src/aggregators/ssajump.jl | 18 +- src/problem.jl | 55 ++--- src/simple_regular_solve.jl | 17 +- src/solve.jl | 71 +++--- src/spatial/directcrdirect.jl | 24 +- src/spatial/nsm.jl | 36 +-- src/spatial/utils.jl | 10 +- src/variable_rate.jl | 113 +++++----- test/degenerate_rx_cases.jl | 20 +- test/ensemble_problems.jl | 140 ++++++------ test/extended_jump_array_remake.jl | 43 ++-- test/remake_test.jl | 51 +++-- test/rng_kwarg_tests.jl | 219 +++++++++++++++++++ test/runtests.jl | 3 +- test/ssa_callback_test.jl | 76 +++---- test/variable_rate.jl | 3 +- 32 files changed, 873 insertions(+), 515 deletions(-) create mode 100644 test/rng_kwarg_tests.jl diff --git a/Project.toml b/Project.toml index 0a2cb01ab..404b9b5ea 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "JumpProcesses" uuid = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5" +version = "9.23.0" authors = ["Chris Rackauckas "] -version = "9.22.1" [deps] ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" @@ -18,6 +18,7 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" +StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" [weakdeps] @@ -45,14 +46,14 @@ KernelAbstractions = "0.9" LinearAlgebra = "1" LinearSolve = "3" OrdinaryDiffEq = "6" -OrdinaryDiffEqCore = "1.32.0" +OrdinaryDiffEqCore = "3.9" Pkg = "1" PoissonRandom = "0.4" Random = "1" RecursiveArrayTools = "3.35" Reexport = "1.2" SafeTestsets = "0.1" -SciMLBase = "2.115" +SciMLBase = "2.144" StableRNGs = "1" StaticArrays = "1.9.8" Statistics = "1" diff --git a/docs/src/api.md b/docs/src/api.md index d8fd944b2..885be3720 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -40,9 +40,109 @@ RSSACR SortingDirect ``` -# Private API Functions +## Random Number Generator Control + +JumpProcesses supports controlling the random number generator (RNG) used for +jump sampling via the `rng` and `seed` keyword arguments to `solve` or `init`. + +### `rng` keyword argument + +Pass any `AbstractRNG` to `solve` or `init`: + +```julia +using Random, StableRNGs + +# Using a StableRNG for cross-version reproducibility +sol = solve(jprob, SSAStepper(); rng = StableRNG(1234)) + +# Using Julia's built-in Xoshiro +sol = solve(jprob, Tsit5(); rng = Xoshiro(42)) +``` + +### `seed` keyword argument + +As a shorthand, pass an integer `seed` to create a `Xoshiro` generator: + +```julia +sol = solve(jprob, SSAStepper(); seed = 1234) +# equivalent to: solve(jprob, SSAStepper(); rng = Xoshiro(1234)) +``` + +### Passing `rng` via `JumpProblem` + +The `rng` keyword can also be passed to [`JumpProblem`](@ref), where it is +stored and automatically forwarded to the solver: + +```julia +jprob = JumpProblem(dprob, Direct(), jump; rng = StableRNG(1234)) +sol = solve(jprob, SSAStepper()) # uses StableRNG(1234) +``` + +An `rng` passed directly to `solve`/`init` takes priority over one stored in +the `JumpProblem`. + +### Resolution priority + +When multiple RNG sources are provided, the following priority determines which +is used: + +| User provides | Result | +|---|---| +| `rng` via `solve`/`init` | Uses that `rng` | +| `seed` via `solve`/`init` | Creates `Xoshiro(seed)` | +| `rng` via `JumpProblem` | Uses that `rng` | +| Nothing | Uses `Random.default_rng()` (SSAStepper, ODE, tau-leaping) or a randomly-seeded `Xoshiro` (SDE) | + +When both `rng` and `seed` are passed to the same call, `rng` takes priority. +When `solve`/`init` kwargs and `JumpProblem` kwargs overlap on the same key, +`solve`/`init` kwargs take priority. + +### Behavior by solver pathway + +| Solver | Default RNG (nothing passed) | `rng` / `seed` support | +|---|---|---| +| `SSAStepper` | `Random.default_rng()` | Full support via `solve`/`init` kwargs | +| ODE solvers (e.g., `Tsit5`) | `Random.default_rng()` | Full support via `solve`/`init` kwargs | +| SDE solvers (e.g., `SRIW1`) | Randomly-seeded `Xoshiro` | Full support; `TaskLocalRNG` is auto-converted to `Xoshiro` | +| `SimpleTauLeaping` | `Random.default_rng()` | Full support via `solve` kwargs | + +!!! note + For reproducible simulations, always pass an explicit `rng` or `seed`. + The default RNG is shared global state and may produce different results + depending on prior usage. + +# Private / Developer API ```@docs ExtendedJumpArray SSAIntegrator ``` + +## Internal Dispatch Pathways + +The following table documents which code handles `solve`/`init` for each solver +type. This is relevant for developers working on JumpProcesses or its solver +backends. + +| Solver type | `__solve` handled by | `__init` handled by | Uses `__jump_init`? | +|---|---|---|---| +| `SSAStepper` | JumpProcesses (`solve.jl`) | JumpProcesses (`SSA_stepper.jl`) | No | +| ODE (e.g., `Tsit5`) | JumpProcesses (`solve.jl`) | JumpProcesses (`solve.jl`) → OrdinaryDiffEq | Yes | +| SDE (e.g., `SRIW1`) | StochasticDiffEq | StochasticDiffEq | No | +| `SimpleTauLeaping` | JumpProcesses (`simple_regular_solve.jl`, custom `DiffEqBase.solve`) | N/A | No | + +For **SSAStepper**, `rng` is resolved via `resolve_rng` in `SSA_stepper.jl`'s +`__init` and stored on the [`SSAIntegrator`](@ref). + +For **ODE solvers**, `rng` is resolved via `resolve_rng` in `__jump_init` +(`solve.jl`) and forwarded to OrdinaryDiffEq's `init`, which stores it on the +`ODEIntegrator`. + +For **SDE solvers**, StochasticDiffEq handles the full solve/init pathway +directly (JumpProcesses' ambiguity-fix `__solve` method is never dispatched to). +StochasticDiffEq has its own `_resolve_rng` that additionally handles +`TaskLocalRNG` conversion and the problem's stored seed. + +For **tau-leaping**, JumpProcesses defines a custom `DiffEqBase.solve` that +bypasses the standard `__solve`/`__init` pathway. It calls `resolve_rng` with +the `JumpProblem`'s stored `rng` kwarg as a fallback. diff --git a/docs/src/faq.md b/docs/src/faq.md index d2d0be3be..0e3dad78c 100644 --- a/docs/src/faq.md +++ b/docs/src/faq.md @@ -55,21 +55,35 @@ jset = JumpSet(; constant_jumps = cjvec, variable_jumps = vjtuple, ## How can I set the random number generator used in the jump process sampling algorithms (SSAs)? -Random number generators can be passed to `JumpProblem` via the `rng` keyword +Random number generators can be passed to `solve` or `init` via the `rng` keyword argument. Continuing the previous example: ```julia -#] add RandomNumbers -using RandomNumbers -jprob = JumpProblem(dprob, Direct(), maj, - rng = Xorshifts.Xoroshiro128Star(rand(UInt64))) +using Random +jprob = JumpProblem(dprob, Direct(), maj) +sol = solve(jprob, SSAStepper(); rng = Xoshiro(1234)) +``` + +Any `AbstractRNG` can be used. For example, to use a generator from +[StableRNGs.jl](https://github.com/JuliaRandom/StableRNGs.jl): + +```julia +using StableRNGs +sol = solve(jprob, SSAStepper(); rng = StableRNG(1234)) ``` -uses the `Xoroshiro128Star` generator from -[RandomNumbers.jl](https://github.com/JuliaRandom/RandomNumbers.jl). +A `seed` keyword argument is also supported as a shorthand for creating a `Xoshiro` +generator: `solve(jprob, SSAStepper(); seed = 1234)`. + +Alternatively, `rng` can be passed to `JumpProblem` where it will be forwarded to the +solver automatically: + +```julia +jprob = JumpProblem(dprob, Direct(), maj; rng = Xoshiro(1234)) +sol = solve(jprob, SSAStepper()) # uses Xoshiro(1234) +``` -On version 1.7 and up, JumpProcesses uses Julia's built-in random number generator by -default. On versions below 1.7 it uses `Xoroshiro128Star`. +By default, JumpProcesses uses Julia's built-in `Random.default_rng()`. ## What are these aggregators and aggregations in JumpProcesses? diff --git a/docs/src/tutorials/simple_poisson_process.md b/docs/src/tutorials/simple_poisson_process.md index ae8e0955e..11f5c5f0e 100644 --- a/docs/src/tutorials/simple_poisson_process.md +++ b/docs/src/tutorials/simple_poisson_process.md @@ -371,11 +371,10 @@ with ``N(t)`` a Poisson counting process with constant transition rate ``\lambda``, and the ``C_i`` independent and identical samples from a uniform distribution over ``\{-1,1\}``. We can simulate such a process as follows. -We first ensure that we use the same random number generator as JumpProcesses. We -can either pass one as an input to [`JumpProblem`](@ref) via the `rng` keyword -argument, and make sure it is the same one we use in our `affect!` function, or -we can just use the default generator chosen by JumpProcesses if one is not -specified, `JumpProcesses.DEFAULT_RNG`. Let's do the latter +We first ensure that we use the same random number generator as JumpProcesses. +Custom RNGs can be passed to `solve` or `init` via the `rng` keyword argument. +If no RNG is specified, JumpProcesses uses `Random.default_rng()`, which is also +available as `JumpProcesses.DEFAULT_RNG`. Let's use the default ```@example tut1 rng = JumpProcesses.DEFAULT_RNG diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl index 87c8c2c26..e32c4fb8e 100644 --- a/src/JumpProcesses.jl +++ b/src/JumpProcesses.jl @@ -40,7 +40,7 @@ using DiffEqBase: DiffEqBase, CallbackSet, ContinuousCallback, DAEFunction, ODESolution, ReturnCode, SDEFunction, SDEProblem, add_tstop!, deleteat!, isinplace, remake, savevalues!, step!, u_modified! -using SciMLBase: SciMLBase, DEIntegrator +using SciMLBase: SciMLBase, DEIntegrator, has_rng, get_rng, set_rng! abstract type AbstractJump end abstract type AbstractMassActionJump <: AbstractJump end diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index 3403b0602..c4f72906d 100644 --- a/src/SSA_stepper.jl +++ b/src/SSA_stepper.jl @@ -13,6 +13,9 @@ Highly efficient integrator for pure jump problems that involve only `ConstantRa `SSAStepper`. - Only supports a limited subset of the output controls from the common solver interface, specifically `save_start`, `save_end`, and `saveat`. + - Supports `rng` and `seed` keyword arguments in `solve`/`init` to control the random + number generator used for jump sampling. `rng` accepts any `AbstractRNG`, while `seed` + creates a `Xoshiro` generator. `rng` takes priority over `seed`. - As when using jumps with ODEs and SDEs, saving controls for whether to save each time a jump occurs are via the `save_positions` keyword argument to `JumpProblem`. Note that when choosing `SSAStepper` as the timestepper, `save_positions = (true,true)`, `(true,false)`, @@ -62,13 +65,13 @@ SciMLBase.allows_late_binding_tstops(::SSAStepper) = true """ $(TYPEDEF) -Solution objects for pure jump problems solved via `SSAStepper`. +Integrator for pure jump problems solved via `SSAStepper`. ## Fields $(FIELDS) """ -mutable struct SSAIntegrator{F, uType, tType, tdirType, P, S, CB, SA, OPT, TS} <: +mutable struct SSAIntegrator{F, uType, tType, tdirType, P, S, CB, SA, OPT, TS, R} <: AbstractSSAIntegrator{SSAStepper, Nothing, uType, tType} """The underlying `prob.f` function. Not currently used.""" f::F @@ -108,6 +111,15 @@ mutable struct SSAIntegrator{F, uType, tType, tdirType, P, S, CB, SA, OPT, TS} < alias_tstops::Bool """If true indicates we have already allocated the tstops array""" copied_tstops::Bool + """The random number generator.""" + rng::R +end + +SciMLBase.has_rng(::SSAIntegrator) = true +SciMLBase.get_rng(integrator::SSAIntegrator) = integrator.rng +function SciMLBase.set_rng!(integrator::SSAIntegrator, rng) + integrator.rng = rng + nothing end (integrator::SSAIntegrator)(t) = copy(integrator.u) @@ -198,6 +210,7 @@ function DiffEqBase.__init(jump_prob::JumpProblem, save_start = true, save_end = true, seed = nothing, + rng = nothing, alias_jump = Threads.threadid() == 1, saveat = nothing, callback = nothing, @@ -219,19 +232,13 @@ function DiffEqBase.__init(jump_prob::JumpProblem, # Check for continuous callbacks passed via kwargs (from JumpProblem constructor or solve) check_continuous_callback_error(callback) + + _rng = resolve_rng(rng, seed) + if alias_jump cb = jump_prob.jump_callback.discrete_callbacks[end] - if seed !== nothing - Random.seed!(cb.condition.rng, seed) - end else cb = deepcopy(jump_prob.jump_callback.discrete_callbacks[end]) - # Only reseed if an explicit seed is provided. This respects the user's RNG choice - # and enables reproducibility. For EnsembleProblems, use prob_func to set unique seeds - # for each trajectory if different results are needed. - if seed !== nothing - Random.seed!(cb.condition.rng, seed) - end end opts = (callback = CallbackSet(callback),) @@ -286,7 +293,7 @@ function DiffEqBase.__init(jump_prob::JumpProblem, integrator = SSAIntegrator(prob.f, copy(prob.u0), prob.tspan[1], prob.tspan[1], tdir, prob.p, sol, 1, prob.tspan[1], cb, _saveat, save_everystep, - save_end, cur_saveat, opts, _tstops, 1, false, true, alias_tstops, false) + save_end, cur_saveat, opts, _tstops, 1, false, true, alias_tstops, false, _rng) cb.initialize(cb, integrator.u, prob.tspan[1], integrator) DiffEqBase.initialize!(opts.callback, integrator.u, prob.tspan[1], integrator) if save_start diff --git a/src/aggregators/ccnrm.jl b/src/aggregators/ccnrm.jl index 379ad642a..720868c50 100644 --- a/src/aggregators/ccnrm.jl +++ b/src/aggregators/ccnrm.jl @@ -3,8 +3,8 @@ # algorithm with optimal binning, Journal of Chemical Physics 143, 074108 # (2015). doi: 10.1063/1.4928635. -mutable struct CCNRMJumpAggregation{T, S, F1, F2, RNG, DEPGR, PT} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} +mutable struct CCNRMJumpAggregation{T, S, F1, F2, DEPGR, PT} <: + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::T @@ -15,15 +15,14 @@ mutable struct CCNRMJumpAggregation{T, S, F1, F2, RNG, DEPGR, PT} <: rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG dep_gr::DEPGR ptt::PT end function CCNRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, - maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, - rng::RNG; num_specs, dep_graph = nothing, - kwargs...) where {T, S, F1, F2, RNG} + maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + num_specs, dep_graph = nothing, + kwargs...) where {T, S, F1, F2} # a dependency graph is needed and must be provided if there are constant rate jumps if dep_graph === nothing @@ -45,30 +44,31 @@ function CCNRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, binwidthconst = binwidthconst, numbinsconst = numbinsconst) # We will re-initialize this in initialize!() affecttype = F2 <: Tuple ? F2 : Any - CCNRMJumpAggregation{T, S, F1, affecttype, RNG, typeof(dg), typeof(ptt)}( + CCNRMJumpAggregation{T, S, F1, affecttype, typeof(dg), typeof(ptt)}( nj, nj, njt, et, crs, sr, maj, rs, affs!, sps, - rng, dg, ptt) + dg, ptt) end -+############################# Required Functions ############################## +############################# Required Functions ############################## # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::CCNRM, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using function wrappers rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps) build_jump_aggregation(CCNRMJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; num_specs = length(u), + rates, affects!, save_positions; num_specs = length(u), kwargs...) end # set up a new simulation and calculate the first jump / jump time function initialize!(p::CCNRMJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] - initialize_rates_and_times!(p, u, params, t) + rng = get_rng(integrator) + initialize_rates_and_times!(p, u, params, t, rng) generate_jumps!(p, integrator, u, params, t) nothing end @@ -79,7 +79,8 @@ function execute_jumps!(p::CCNRMJumpAggregation, integrator, u, params, t, affec u = update_state!(p, integrator, u, affects!) # update current jump rates and times - update_dependent_rates!(p, u, params, t) + rng = get_rng(integrator) + update_dependent_rates!(p, u, params, t, rng) nothing end @@ -88,7 +89,7 @@ end function generate_jumps!(p::CCNRMJumpAggregation, integrator, u, params, t) p.next_jump, p.next_jump_time = getfirst(p.ptt) - # Rebuild the table if no next jump is found. + # Rebuild the table if no next jump is found. if p.next_jump == 0 timestep = 1 / sum(p.cur_rates) min_time = minimum(p.ptt.times) @@ -102,7 +103,7 @@ end ######################## SSA specific helper routines ######################## # Recalculate jump rates for jumps that depend on the just executed jump (p.next_jump) -function update_dependent_rates!(p::CCNRMJumpAggregation, u, params, t) +function update_dependent_rates!(p::CCNRMJumpAggregation, u, params, t, rng) @inbounds dep_rxs = p.dep_gr[p.next_jump] (; ptt, cur_rates, rates, ma_jumps, end_time) = p num_majumps = get_num_majumps(ma_jumps) @@ -125,7 +126,7 @@ function update_dependent_rates!(p::CCNRMJumpAggregation, u, params, t) end else if cur_rates[rx] > zero(eltype(cur_rates)) - update!(ptt, rx, oldtime, t + randexp(p.rng) / cur_rates[rx]) + update!(ptt, rx, oldtime, t + randexp(rng) / cur_rates[rx]) else update!(ptt, rx, oldtime, floatmax(typeof(t))) end @@ -134,15 +135,15 @@ function update_dependent_rates!(p::CCNRMJumpAggregation, u, params, t) nothing end -# Evaluate all the rates and initialize the times in the priority table. -function initialize_rates_and_times!(p::CCNRMJumpAggregation, u, params, t) +# Evaluate all the rates and initialize the times in the priority table. +function initialize_rates_and_times!(p::CCNRMJumpAggregation, u, params, t, rng) # Initialize next-reaction times for the mass action jumps majumps = p.ma_jumps cur_rates = p.cur_rates pttdata = Vector{typeof(t)}(undef, length(cur_rates)) @inbounds for i in 1:get_num_majumps(majumps) cur_rates[i] = evalrxrate(u, i, majumps) - pttdata[i] = t + randexp(p.rng) / cur_rates[i] + pttdata[i] = t + randexp(rng) / cur_rates[i] end # Initialize next-reaction times for the constant rates @@ -150,11 +151,11 @@ function initialize_rates_and_times!(p::CCNRMJumpAggregation, u, params, t) idx = get_num_majumps(majumps) + 1 @inbounds for rate in rates cur_rates[idx] = rate(u, params, t) - pttdata[idx] = t + randexp(p.rng) / cur_rates[idx] + pttdata[idx] = t + randexp(rng) / cur_rates[idx] idx += 1 end - # Build the priority time table with the times and bin width. + # Build the priority time table with the times and bin width. timestep = 1 / sum(cur_rates) p.ptt.times = pttdata rebuild!(p.ptt, t, timestep) diff --git a/src/aggregators/coevolve.jl b/src/aggregators/coevolve.jl index dc48a5e21..ab5dc5f4f 100644 --- a/src/aggregators/coevolve.jl +++ b/src/aggregators/coevolve.jl @@ -1,8 +1,8 @@ """ Queue method. This method handles variable intensity rates. """ -mutable struct CoevolveJumpAggregation{T, S, F1, F2, RNG, GR, PQ} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} +mutable struct CoevolveJumpAggregation{T, S, F1, F2, GR, PQ} <: + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int # the next jump to execute prev_jump::Int # the previous jump that was executed next_jump_time::T # the time of the next jump @@ -13,7 +13,6 @@ mutable struct CoevolveJumpAggregation{T, S, F1, F2, RNG, GR, PQ} <: rates::F1 # vector of rate functions affects!::F2 # vector of affect functions for VariableRateJumps save_positions::Tuple{Bool, Bool} # tuple for whether to save the jumps before and/or after event - rng::RNG # random number generator dep_gr::GR # map from jumps to jumps depending on it pq::PQ # priority queue of next time lrates::F1 # vector of rate lower bound functions @@ -24,10 +23,10 @@ mutable struct CoevolveJumpAggregation{T, S, F1, F2, RNG, GR, PQ} <: end function CoevolveJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::Nothing, - maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, - rng::RNG; u::U, dep_graph = nothing, lrates, urates, + maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + u::U, dep_graph = nothing, lrates, urates, rateintervals, haslratevec, - cur_lrates::Vector{T}) where {T, S, F1, F2, RNG, U} + cur_lrates::Vector{T}) where {T, S, F1, F2, U} if dep_graph === nothing if (get_num_majumps(maj) == 0) || !isempty(urates) error("To use Coevolve a dependency graph between jumps must be supplied.") @@ -49,9 +48,9 @@ function CoevolveJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::Not pq = MutableBinaryMinHeap{T}() affecttype = F2 <: Tuple ? F2 : Any - CoevolveJumpAggregation{T, S, F1, affecttype, RNG, typeof(dg), + CoevolveJumpAggregation{T, S, F1, affecttype, typeof(dg), typeof(pq)}(nj, nj, njt, et, crs, sr, maj, - rs, affs!, sps, rng, dg, pq, + rs, affs!, sps, dg, pq, lrates, urates, rateintervals, haslratevec, cur_lrates) end @@ -98,7 +97,7 @@ end # creating the JumpAggregation structure (tuple-based variable jumps) function aggregate(aggregator::Coevolve, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; dep_graph = nothing, + ma_jumps, save_positions; dep_graph = nothing, variable_jumps = nothing, kwargs...) RateWrapper = FunctionWrappers.FunctionWrapper{typeof(t), Tuple{typeof(u), typeof(p), typeof(t)}} @@ -141,7 +140,7 @@ function aggregate(aggregator::Coevolve, u, p, t, end_time, constant_jumps, next_jump = 0 next_jump_time = typemax(t) CoevolveJumpAggregation(next_jump, next_jump_time, end_time, cur_rates, sum_rate, - ma_jumps, rates, affects!, save_positions, rng; + ma_jumps, rates, affects!, save_positions; u, dep_graph, lrates, urates, rateintervals, haslratevec, cur_lrates) end @@ -149,7 +148,8 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::CoevolveJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] - fill_rates_and_get_times!(p, u, params, t) + rng = get_rng(integrator) + fill_rates_and_get_times!(p, u, params, t, rng) generate_jumps!(p, integrator, u, params, t) nothing end @@ -160,7 +160,8 @@ function execute_jumps!(p::CoevolveJumpAggregation, integrator, u, params, t, # execute jump update_state!(p, integrator, u, affects!) # update current jump rates and times - update_dependent_rates!(p, integrator.u, integrator.p, t) + rng = get_rng(integrator) + update_dependent_rates!(p, integrator.u, integrator.p, t, rng) nothing end @@ -178,7 +179,8 @@ function accept_next_jump!(p::CoevolveJumpAggregation, integrator, u, params, t) (next_jump <= num_majumps) && return true - (; cur_rates, rates, rng, urates, cur_lrates) = p + (; cur_rates, rates, urates, cur_lrates) = p + rng = get_rng(integrator) num_cjumps = length(urates) - length(rates) uidx = next_jump - num_majumps lidx = uidx - num_cjumps @@ -225,11 +227,11 @@ function accept_next_jump!(p::CoevolveJumpAggregation, integrator, u, params, t) return false end -function update_dependent_rates!(p::CoevolveJumpAggregation, u, params, t) +function update_dependent_rates!(p::CoevolveJumpAggregation, u, params, t, rng) @inbounds deps = p.dep_gr[p.next_jump] (; cur_rates, pq) = p for (ix, i) in enumerate(deps) - ti, urate_i = next_time(p, u, params, t, i) + ti, urate_i = next_time(p, u, params, t, i, rng) update!(pq, i, ti) @inbounds cur_rates[i] = urate_i end @@ -256,8 +258,8 @@ end @inbounds return p.rates[lidx](u, params, t) end -function next_time(p::CoevolveJumpAggregation, u, params, t, i) - (; next_jump, cur_rates, ma_jumps, rates, rng, pq, urates) = p +function next_time(p::CoevolveJumpAggregation, u, params, t, i, rng) + (; next_jump, cur_rates, ma_jumps, rates, pq, urates) = p num_majumps = get_num_majumps(ma_jumps) num_cjumps = length(urates) - length(rates) uidx = i - num_majumps @@ -300,12 +302,12 @@ function next_candidate_time!(p::CoevolveJumpAggregation, u, params, t, s, lidx) end # re-evaluates all rates, recalculate all jump times, and reinit the priority queue -function fill_rates_and_get_times!(p::CoevolveJumpAggregation, u, params, t) +function fill_rates_and_get_times!(p::CoevolveJumpAggregation, u, params, t, rng) num_jumps = get_num_majumps(p.ma_jumps) + length(p.urates) p.cur_rates = zeros(typeof(t), num_jumps) jump_times = Vector{typeof(t)}(undef, num_jumps) @inbounds for i in 1:num_jumps - jump_times[i], p.cur_rates[i] = next_time(p, u, params, t, i) + jump_times[i], p.cur_rates[i] = next_time(p, u, params, t, i, rng) end p.pq = MutableBinaryMinHeap(jump_times) nothing diff --git a/src/aggregators/direct.jl b/src/aggregators/direct.jl index ab33dd842..b51e18f7c 100644 --- a/src/aggregators/direct.jl +++ b/src/aggregators/direct.jl @@ -1,5 +1,5 @@ -mutable struct DirectJumpAggregation{T, S, F1, F2, RNG} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} +mutable struct DirectJumpAggregation{T, S, F1, F2} <: + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::T @@ -10,38 +10,37 @@ mutable struct DirectJumpAggregation{T, S, F1, F2, RNG} <: rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG end function DirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, - rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, rng::RNG; - kwargs...) where {T, S, F1, F2, RNG} + rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + kwargs...) where {T, S, F1, F2} affecttype = F2 <: Tuple ? F2 : Any - DirectJumpAggregation{T, S, F1, affecttype, RNG}(nj, nj, njt, et, crs, sr, maj, rs, - affs!, sps, rng) + DirectJumpAggregation{T, S, F1, affecttype}(nj, nj, njt, et, crs, sr, maj, rs, + affs!, sps) end ############################# Required Functions ############################# # creating the JumpAggregation structure (tuple-based constant jumps) function aggregate(aggregator::Direct, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using tuples rates, affects! = get_jump_info_tuples(constant_jumps) build_jump_aggregation(DirectJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; kwargs...) + rates, affects!, save_positions; kwargs...) end # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::DirectFW, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using function wrappers rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps) build_jump_aggregation(DirectJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; kwargs...) + rates, affects!, save_positions; kwargs...) end # set up a new simulation and calculate the first jump / jump time @@ -60,9 +59,10 @@ end # calculate the next jump / jump time function generate_jumps!(p::DirectJumpAggregation, integrator, u, params, t) - p.sum_rate, ttnj = time_to_next_jump(p, u, params, t) + rng = get_rng(integrator) + p.sum_rate, ttnj = time_to_next_jump(p, u, params, t, rng) p.next_jump_time = add_fast(t, ttnj) - @inbounds p.next_jump = searchsortedfirst(p.cur_rates, rand(p.rng) * p.sum_rate) + @inbounds p.next_jump = searchsortedfirst(p.cur_rates, rand(rng) * p.sum_rate) nothing end @@ -70,7 +70,7 @@ end # tuple-based constant jumps function time_to_next_jump(p::DirectJumpAggregation{T, S, F1}, u, params, - t) where {T, S, F1 <: Tuple} + t, rng) where {T, S, F1 <: Tuple} prev_rate = zero(t) new_rate = zero(t) cur_rates = p.cur_rates @@ -96,7 +96,7 @@ function time_to_next_jump(p::DirectJumpAggregation{T, S, F1}, u, params, end @inbounds sum_rate = cur_rates[end] - sum_rate, randexp(p.rng) / sum_rate + sum_rate, randexp(rng) / sum_rate end @inline function fill_cur_rates(u, p, t, cur_rates, idx, rate, rates...) @@ -112,7 +112,7 @@ end # function wrapper-based constant jumps function time_to_next_jump(p::DirectJumpAggregation{T, S, F1}, u, params, - t) where {T, S, F1 <: AbstractArray} + t, rng) where {T, S, F1 <: AbstractArray} prev_rate = zero(t) new_rate = zero(t) cur_rates = p.cur_rates @@ -137,5 +137,5 @@ function time_to_next_jump(p::DirectJumpAggregation{T, S, F1}, u, params, end @inbounds sum_rate = cur_rates[end] - sum_rate, randexp(p.rng) / sum_rate + sum_rate, randexp(rng) / sum_rate end diff --git a/src/aggregators/directcr.jl b/src/aggregators/directcr.jl index 41d079fe8..636300bb3 100644 --- a/src/aggregators/directcr.jl +++ b/src/aggregators/directcr.jl @@ -10,9 +10,9 @@ by S. Mauch and M. Stalzer, ACM Trans. Comp. Biol. and Bioinf., 8, No. 1, 27-35 const MINJUMPRATE = 2.0^exponent(1e-12) -mutable struct DirectCRJumpAggregation{T, S, F1, F2, RNG, DEPGR, U <: PriorityTable, +mutable struct DirectCRJumpAggregation{T, S, F1, F2, DEPGR, U <: PriorityTable, W <: Function} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::T @@ -23,7 +23,6 @@ mutable struct DirectCRJumpAggregation{T, S, F1, F2, RNG, DEPGR, U <: PriorityTa rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG dep_gr::DEPGR minrate::T maxrate::T # initial maxrate only, table can increase beyond it! @@ -32,11 +31,11 @@ mutable struct DirectCRJumpAggregation{T, S, F1, F2, RNG, DEPGR, U <: PriorityTa end function DirectCRJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, - maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, - rng::RNG; num_specs, dep_graph = nothing, + maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + num_specs, dep_graph = nothing, minrate = convert(T, MINJUMPRATE), maxrate = convert(T, Inf), - kwargs...) where {T, S, F1, F2, RNG} + kwargs...) where {T, S, F1, F2} # a dependency graph is needed and must be provided if there are constant rate jumps if dep_graph === nothing @@ -63,9 +62,9 @@ function DirectCRJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, rt = PriorityTable(ratetogroup, zeros(T, 1), minrate, 2 * minrate) affecttype = F2 <: Tuple ? F2 : Any - DirectCRJumpAggregation{T, S, F1, affecttype, RNG, typeof(dg), + DirectCRJumpAggregation{T, S, F1, affecttype, typeof(dg), typeof(rt), typeof(ratetogroup)}(nj, nj, njt, et, crs, sr, maj, - rs, affs!, sps, rng, dg, + rs, affs!, sps, dg, minrate, maxrate, rt, ratetogroup) end @@ -74,13 +73,13 @@ end # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::DirectCR, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using function wrappers rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps) build_jump_aggregation(DirectCRJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; num_specs = length(u), + rates, affects!, save_positions; num_specs = length(u), kwargs...) end @@ -113,10 +112,11 @@ end # calculate the next jump / jump time function generate_jumps!(p::DirectCRJumpAggregation, integrator, u, params, t) - p.next_jump_time = t + randexp(p.rng) / p.sum_rate + rng = get_rng(integrator) + p.next_jump_time = t + randexp(rng) / p.sum_rate if p.next_jump_time < p.end_time - p.next_jump = sample(p.rt, p.cur_rates, p.rng) + p.next_jump = sample(p.rt, p.cur_rates, rng) end nothing end diff --git a/src/aggregators/frm.jl b/src/aggregators/frm.jl index 94baed375..0db30113b 100644 --- a/src/aggregators/frm.jl +++ b/src/aggregators/frm.jl @@ -1,5 +1,5 @@ -mutable struct FRMJumpAggregation{T, S, F1, F2, RNG} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} +mutable struct FRMJumpAggregation{T, S, F1, F2} <: + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::T @@ -10,40 +10,39 @@ mutable struct FRMJumpAggregation{T, S, F1, F2, RNG} <: rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG end function FRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, - affs!::F2, sps::Tuple{Bool, Bool}, rng::RNG; - kwargs...) where {T, S, F1, F2, RNG} + affs!::F2, sps::Tuple{Bool, Bool}; + kwargs...) where {T, S, F1, F2} affecttype = F2 <: Tuple ? F2 : Any - FRMJumpAggregation{T, S, F1, affecttype, RNG}(nj, nj, njt, et, crs, sr, maj, rs, - affs!, sps, rng) + FRMJumpAggregation{T, S, F1, affecttype}(nj, nj, njt, et, crs, sr, maj, rs, + affs!, sps) end ############################# Required Functions ############################# # creating the JumpAggregation structure (tuple-based constant jumps) function aggregate(aggregator::FRM, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using tuples rates, affects! = get_jump_info_tuples(constant_jumps) build_jump_aggregation( FRMJumpAggregation, u, p, t, end_time, ma_jumps, rates, affects!, - save_positions, rng; kwargs...) + save_positions; kwargs...) end # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::FRMFW, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using function wrappers rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps) build_jump_aggregation( FRMJumpAggregation, u, p, t, end_time, ma_jumps, rates, affects!, - save_positions, rng; kwargs...) + save_positions; kwargs...) end # set up a new simulation and calculate the first jump / jump time @@ -62,8 +61,9 @@ end # calculate the next jump / jump time function generate_jumps!(p::FRMJumpAggregation, integrator, u, params, t) - nextmaj, ttnmaj = next_ma_jump(p, u, params, t) - nextcrj, ttncrj = next_constant_rate_jump(p, u, params, t) + rng = get_rng(integrator) + nextmaj, ttnmaj = next_ma_jump(p, u, params, t, rng) + nextcrj, ttncrj = next_constant_rate_jump(p, u, params, t, rng) # execute reaction with minimal time if ttnmaj < ttncrj @@ -79,13 +79,13 @@ end ######################## SSA specific helper routines ######################## # mass action jumps -function next_ma_jump(p::FRMJumpAggregation, u, params, t) +function next_ma_jump(p::FRMJumpAggregation, u, params, t, rng) ttnj = typemax(typeof(t)) nextrx = zero(Int) majumps = p.ma_jumps @inbounds for i in 1:get_num_majumps(majumps) p.cur_rates[i] = evalrxrate(u, i, majumps) - dt = randexp(p.rng) / p.cur_rates[i] + dt = randexp(rng) / p.cur_rates[i] if dt < ttnj ttnj = dt nextrx = i @@ -95,15 +95,15 @@ function next_ma_jump(p::FRMJumpAggregation, u, params, t) end # tuple-based constant jumps -function next_constant_rate_jump(p::FRMJumpAggregation{T, S, F1, F2, RNG}, u, params, - t) where {T, S, F1 <: Tuple, F2 <: Tuple, RNG} +function next_constant_rate_jump(p::FRMJumpAggregation{T, S, F1, F2}, u, params, + t, rng) where {T, S, F1 <: Tuple, F2 <: Tuple} ttnj = typemax(typeof(t)) nextrx = zero(Int) if !isempty(p.rates) idx = get_num_majumps(p.ma_jumps) + 1 fill_cur_rates(u, params, t, p.cur_rates, idx, p.rates...) @inbounds for i in idx:length(p.cur_rates) - dt = randexp(p.rng) / p.cur_rates[i] + dt = randexp(rng) / p.cur_rates[i] if dt < ttnj ttnj = dt nextrx = i @@ -115,14 +115,14 @@ end # function wrapper-based constant jumps function next_constant_rate_jump(p::FRMJumpAggregation{T, S, F1}, u, params, - t) where {T, S, F1 <: AbstractArray} + t, rng) where {T, S, F1 <: AbstractArray} ttnj = typemax(typeof(t)) nextrx = zero(Int) if !isempty(p.rates) idx = get_num_majumps(p.ma_jumps) + 1 @inbounds for i in 1:length(p.rates) p.cur_rates[idx] = p.rates[i](u, params, t) - dt = randexp(p.rng) / p.cur_rates[idx] + dt = randexp(rng) / p.cur_rates[idx] if dt < ttnj ttnj = dt nextrx = idx diff --git a/src/aggregators/nrm.jl b/src/aggregators/nrm.jl index 7fbcd5964..934540448 100644 --- a/src/aggregators/nrm.jl +++ b/src/aggregators/nrm.jl @@ -1,8 +1,8 @@ # Implementation the original Next Reaction Method # Gibson and Bruck, J. Phys. Chem. A, 104 (9), (2000) -mutable struct NRMJumpAggregation{T, S, F1, F2, RNG, DEPGR, PQ} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} +mutable struct NRMJumpAggregation{T, S, F1, F2, DEPGR, PQ} <: + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::T @@ -13,15 +13,14 @@ mutable struct NRMJumpAggregation{T, S, F1, F2, RNG, DEPGR, PQ} <: rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG dep_gr::DEPGR pq::PQ end function NRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, - maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, - rng::RNG; num_specs, dep_graph = nothing, - kwargs...) where {T, S, F1, F2, RNG} + maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + num_specs, dep_graph = nothing, + kwargs...) where {T, S, F1, F2} # a dependency graph is needed and must be provided if there are constant rate jumps if dep_graph === nothing @@ -40,29 +39,30 @@ function NRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, pq = MutableBinaryMinHeap{T}() affecttype = F2 <: Tuple ? F2 : Any - NRMJumpAggregation{T, S, F1, affecttype, RNG, typeof(dg), typeof(pq)}(nj, nj, njt, et, + NRMJumpAggregation{T, S, F1, affecttype, typeof(dg), typeof(pq)}(nj, nj, njt, et, crs, sr, maj, rs, affs!, sps, - rng, dg, pq) + dg, pq) end -+############################# Required Functions ############################## +############################# Required Functions ############################## # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::NRM, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using function wrappers rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps) build_jump_aggregation(NRMJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; num_specs = length(u), + rates, affects!, save_positions; num_specs = length(u), kwargs...) end # set up a new simulation and calculate the first jump / jump time function initialize!(p::NRMJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] - fill_rates_and_get_times!(p, u, params, t) + rng = get_rng(integrator) + fill_rates_and_get_times!(p, u, params, t, rng) generate_jumps!(p, integrator, u, params, t) nothing end @@ -73,7 +73,8 @@ function execute_jumps!(p::NRMJumpAggregation, integrator, u, params, t, affects u = update_state!(p, integrator, u, affects!) # update current jump rates and times - update_dependent_rates!(p, u, params, t) + rng = get_rng(integrator) + update_dependent_rates!(p, u, params, t, rng) nothing end @@ -87,7 +88,7 @@ end ######################## SSA specific helper routines ######################## # recalculate jump rates for jumps that depend on the just executed jump (p.next_jump) -function update_dependent_rates!(p::NRMJumpAggregation, u, params, t) +function update_dependent_rates!(p::NRMJumpAggregation, u, params, t, rng) @inbounds dep_rxs = p.dep_gr[p.next_jump] (; cur_rates, rates, ma_jumps) = p num_majumps = get_num_majumps(ma_jumps) @@ -108,7 +109,7 @@ function update_dependent_rates!(p::NRMJumpAggregation, u, params, t) end else if cur_rates[rx] > zero(eltype(cur_rates)) - update!(p.pq, rx, t + randexp(p.rng) / cur_rates[rx]) + update!(p.pq, rx, t + randexp(rng) / cur_rates[rx]) else update!(p.pq, rx, typemax(t)) end @@ -118,7 +119,7 @@ function update_dependent_rates!(p::NRMJumpAggregation, u, params, t) end # reevaluate all rates, recalculate all jump times, and reinit the priority queue -function fill_rates_and_get_times!(p::NRMJumpAggregation, u, params, t) +function fill_rates_and_get_times!(p::NRMJumpAggregation, u, params, t, rng) # mass action jumps majumps = p.ma_jumps @@ -126,7 +127,7 @@ function fill_rates_and_get_times!(p::NRMJumpAggregation, u, params, t) pqdata = Vector{typeof(t)}(undef, length(cur_rates)) @inbounds for i in 1:get_num_majumps(majumps) cur_rates[i] = evalrxrate(u, i, majumps) - pqdata[i] = t + randexp(p.rng) / cur_rates[i] + pqdata[i] = t + randexp(rng) / cur_rates[i] end # constant rates @@ -134,7 +135,7 @@ function fill_rates_and_get_times!(p::NRMJumpAggregation, u, params, t) idx = get_num_majumps(majumps) + 1 @inbounds for rate in rates cur_rates[idx] = rate(u, params, t) - pqdata[idx] = t + randexp(p.rng) / cur_rates[idx] + pqdata[idx] = t + randexp(rng) / cur_rates[idx] idx += 1 end diff --git a/src/aggregators/rdirect.jl b/src/aggregators/rdirect.jl index 8376b71b9..eea283160 100644 --- a/src/aggregators/rdirect.jl +++ b/src/aggregators/rdirect.jl @@ -2,8 +2,8 @@ Direct with rejection sampling """ -mutable struct RDirectJumpAggregation{T, S, F1, F2, RNG, DEPGR} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} +mutable struct RDirectJumpAggregation{T, S, F1, F2, DEPGR} <: + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::T @@ -14,7 +14,6 @@ mutable struct RDirectJumpAggregation{T, S, F1, F2, RNG, DEPGR} <: rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG dep_gr::DEPGR max_rate::T counter::Int @@ -22,10 +21,10 @@ mutable struct RDirectJumpAggregation{T, S, F1, F2, RNG, DEPGR} <: end function RDirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, - rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, rng::RNG; + rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; num_specs, counter_threshold = length(crs), dep_graph = nothing, - kwargs...) where {T, S, F1, F2, RNG} + kwargs...) where {T, S, F1, F2} # a dependency graph is needed and must be provided if there are constant rate jumps if dep_graph === nothing if (get_num_majumps(maj) == 0) || !isempty(rs) @@ -42,9 +41,9 @@ function RDirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, m max_rate = maximum(crs) affecttype = F2 <: Tuple ? F2 : Any - return RDirectJumpAggregation{T, S, F1, affecttype, RNG, typeof(dg)}(nj, nj, njt, et, + return RDirectJumpAggregation{T, S, F1, affecttype, typeof(dg)}(nj, nj, njt, et, crs, sr, maj, rs, - affs!, sps, rng, + affs!, sps, dg, max_rate, 0, counter_threshold) end @@ -53,13 +52,13 @@ end # creating the JumpAggregation structure (tuple-based constant jumps) function aggregate(aggregator::RDirect, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using function wrappers rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps) build_jump_aggregation(RDirectJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; num_specs = length(u), + rates, affects!, save_positions; num_specs = length(u), kwargs...) end @@ -93,7 +92,8 @@ function generate_jumps!(p::RDirectJumpAggregation, integrator, u, params, t) if nomorejumps!(p, sum_rate) return nothing end - (; rng, cur_rates, max_rate) = p + rng = get_rng(integrator) + (; cur_rates, max_rate) = p num_rxs = length(cur_rates) counter = 0 @@ -105,7 +105,7 @@ function generate_jumps!(p::RDirectJumpAggregation, integrator, u, params, t) p.counter = counter p.next_jump = rx - p.next_jump_time = t + randexp(p.rng) / sum_rate + p.next_jump_time = t + randexp(rng) / sum_rate nothing end diff --git a/src/aggregators/rssa.jl b/src/aggregators/rssa.jl index 2943bf88b..f86e0b444 100644 --- a/src/aggregators/rssa.jl +++ b/src/aggregators/rssa.jl @@ -4,8 +4,8 @@ # functions of the current population sizes (i.e. u) # requires vartojumps_map and fluct_rates as JumpProblem keywords -mutable struct RSSAJumpAggregation{T, S, F1, F2, RNG, VJMAP, JVMAP, BD, U} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} +mutable struct RSSAJumpAggregation{T, S, F1, F2, VJMAP, JVMAP, BD, U} <: + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::T @@ -17,7 +17,6 @@ mutable struct RSSAJumpAggregation{T, S, F1, F2, RNG, VJMAP, JVMAP, BD, U} <: rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG vartojumps_map::VJMAP jumptovars_map::JVMAP bracket_data::BD @@ -26,10 +25,10 @@ mutable struct RSSAJumpAggregation{T, S, F1, F2, RNG, VJMAP, JVMAP, BD, U} <: end function RSSAJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, - maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, - rng::RNG; u::U, vartojumps_map = nothing, + maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + u::U, vartojumps_map = nothing, jumptovars_map = nothing, - bracket_data = nothing, kwargs...) where {T, S, F1, F2, RNG, U} + bracket_data = nothing, kwargs...) where {T, S, F1, F2, U} # a dependency graph is needed and must be provided if there are constant rate jumps if vartojumps_map === nothing if (get_num_majumps(maj) == 0) || !isempty(rs) @@ -63,10 +62,10 @@ function RSSAJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, uhigh = similar(u) affecttype = F2 <: Tuple ? F2 : Any - RSSAJumpAggregation{T, S, F1, affecttype, RNG, typeof(vtoj_map), + RSSAJumpAggregation{T, S, F1, affecttype, typeof(vtoj_map), typeof(jtov_map), typeof(bd), U}(nj, nj, njt, et, crl_bnds, crh_bnds, sr, maj, rs, affs!, sps, - rng, vtoj_map, jtov_map, bd, ulow, + vtoj_map, jtov_map, bd, ulow, uhigh) end @@ -74,13 +73,13 @@ end # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::RSSA, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using function wrappers rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps) build_jump_aggregation(RSSAJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; u = u, + rates, affects!, save_positions; u, kwargs...) end @@ -108,7 +107,8 @@ function generate_jumps!(p::RSSAJumpAggregation, integrator, u, params, t) return nothing end # next jump type - (; ma_jumps, rates, cur_rate_high, cur_rate_low, rng) = p + (; ma_jumps, rates, cur_rate_high, cur_rate_low) = p + rng = get_rng(integrator) num_majumps = get_num_majumps(ma_jumps) rerl = zero(sum_rate) diff --git a/src/aggregators/rssacr.jl b/src/aggregators/rssacr.jl index 1caf3be53..9cc47c3b8 100644 --- a/src/aggregators/rssacr.jl +++ b/src/aggregators/rssacr.jl @@ -4,9 +4,9 @@ Composition-Rejection with Rejection sampling method (RSSA-CR) const MINJUMPRATE = 2.0^exponent(1e-12) -mutable struct RSSACRJumpAggregation{F, S, F1, F2, RNG, U, VJMAP, JVMAP, BD, +mutable struct RSSACRJumpAggregation{F, S, F1, F2, U, VJMAP, JVMAP, BD, P <: PriorityTable, W <: Function} <: - AbstractSSAJumpAggregator{F, S, F1, F2, RNG} + AbstractSSAJumpAggregator{F, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::F @@ -18,7 +18,6 @@ mutable struct RSSACRJumpAggregation{F, S, F1, F2, RNG, U, VJMAP, JVMAP, BD, rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG vartojumps_map::VJMAP jumptovars_map::JVMAP bracket_data::BD @@ -31,11 +30,11 @@ mutable struct RSSACRJumpAggregation{F, S, F1, F2, RNG, U, VJMAP, JVMAP, BD, end function RSSACRJumpAggregation(nj::Int, njt::F, et::F, crs::Vector{F}, sum_rate::F, maj::S, - rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, rng::RNG; u::U, + rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; u::U, vartojumps_map = nothing, jumptovars_map = nothing, bracket_data = nothing, minrate = convert(F, MINJUMPRATE), maxrate = convert(F, Inf), - kwargs...) where {F, S, F1, F2, RNG, U} + kwargs...) where {F, S, F1, F2, U} # a dependency graph is needed and must be provided if there are constant rate jumps if vartojumps_map === nothing if (get_num_majumps(maj) == 0) || !isempty(rs) @@ -80,10 +79,10 @@ function RSSACRJumpAggregation(nj::Int, njt::F, et::F, crs::Vector{F}, sum_rate: rt = PriorityTable(ratetogroup, zeros(F, 1), minrate, 2 * minrate) affecttype = F2 <: Tuple ? F2 : Any - RSSACRJumpAggregation{typeof(njt), S, F1, affecttype, RNG, U, typeof(vtoj_map), + RSSACRJumpAggregation{typeof(njt), S, F1, affecttype, U, typeof(vtoj_map), typeof(jtov_map), typeof(bd), typeof(rt), typeof(ratetogroup)}(nj, nj, njt, et, crl_bnds, crh_bnds, - sum_rate, maj, rs, affs!, sps, rng, vtoj_map, + sum_rate, maj, rs, affs!, sps, vtoj_map, jtov_map, bd, ulow, uhigh, minrate, maxrate, rt, ratetogroup) end @@ -92,13 +91,13 @@ end # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::RSSACR, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using function wrappers rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps) build_jump_aggregation(RSSACRJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; u = u, kwargs...) + rates, affects!, save_positions; u, kwargs...) end # set up a new simulation and calculate the first jump / jump time @@ -134,7 +133,8 @@ function generate_jumps!(p::RSSACRJumpAggregation, integrator, u, params, t) return nothing end - (; rt, ma_jumps, rates, cur_rate_high, cur_rate_low, rng) = p + (; rt, ma_jumps, rates, cur_rate_high, cur_rate_low) = p + rng = get_rng(integrator) num_majumps = get_num_majumps(ma_jumps) rerl = zero(sum_rate) diff --git a/src/aggregators/sortingdirect.jl b/src/aggregators/sortingdirect.jl index f9048f039..c20ea1ae4 100644 --- a/src/aggregators/sortingdirect.jl +++ b/src/aggregators/sortingdirect.jl @@ -2,8 +2,8 @@ # "The sorting direct method for stochastic simulation of biochemical systems with varying reaction execution behavior" # Comp. Bio. and Chem., 30, pg. 39-49 (2006). -mutable struct SortingDirectJumpAggregation{T, S, F1, F2, RNG, DEPGR} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} +mutable struct SortingDirectJumpAggregation{T, S, F1, F2, DEPGR} <: + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::T @@ -14,16 +14,15 @@ mutable struct SortingDirectJumpAggregation{T, S, F1, F2, RNG, DEPGR} <: rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG dep_gr::DEPGR jump_search_order::Vector{Int} jump_search_idx::Int end function SortingDirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, - maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}, - rng::RNG; num_specs, dep_graph = nothing, - kwargs...) where {T, S, F1, F2, RNG} + maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + num_specs, dep_graph = nothing, + kwargs...) where {T, S, F1, F2} # a dependency graph is needed and must be provided if there are constant rate jumps if dep_graph === nothing @@ -42,9 +41,9 @@ function SortingDirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr # map jump idx to idx in cur_rates jtoidx = collect(1:length(crs)) affecttype = F2 <: Tuple ? F2 : Any - SortingDirectJumpAggregation{T, S, F1, affecttype, RNG, typeof(dg)}(nj, nj, njt, et, + SortingDirectJumpAggregation{T, S, F1, affecttype, typeof(dg)}(nj, nj, njt, et, crs, sr, maj, rs, - affs!, sps, rng, + affs!, sps, dg, jtoidx, zero(Int)) end @@ -53,13 +52,13 @@ end # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::SortingDirect, u, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; kwargs...) + ma_jumps, save_positions; kwargs...) # handle constant jumps using function wrappers rates, affects! = get_jump_info_fwrappers(u, p, t, constant_jumps) build_jump_aggregation(SortingDirectJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; num_specs = length(u), + rates, affects!, save_positions; num_specs = length(u), kwargs...) end @@ -92,14 +91,15 @@ end # calculate the next jump / jump time function generate_jumps!(p::SortingDirectJumpAggregation, integrator, u, params, t) - p.next_jump_time = t + randexp(p.rng) / p.sum_rate + rng = get_rng(integrator) + p.next_jump_time = t + randexp(rng) / p.sum_rate # search for next jump if p.next_jump_time < p.end_time cur_rates = p.cur_rates numjumps = length(cur_rates) jso = p.jump_search_order - rn = p.sum_rate * rand(p.rng) + rn = p.sum_rate * rand(rng) @inbounds for idx in 1:numjumps rn -= cur_rates[jso[idx]] if rn < zero(rn) diff --git a/src/aggregators/ssajump.jl b/src/aggregators/ssajump.jl index 90c260c97..0e72b2348 100644 --- a/src/aggregators/ssajump.jl +++ b/src/aggregators/ssajump.jl @@ -13,13 +13,11 @@ An aggregator interface for SSA-like algorithms. - `rates` # vector of rate functions for ConstantRateJumps - `affects!` # vector of affect functions for ConstantRateJumps - `save_positions` # tuple for whether to save the jumps before and/or after event - - `rng` # random number generator - ### Optional fields: - `dep_gr` # dependency graph, dep_gr[i] = indices of reactions that should be updated when rx i occurs. """ -abstract type AbstractSSAJumpAggregator{T, S, F1, F2, RNG} <: AbstractJumpAggregator end +abstract type AbstractSSAJumpAggregator{T, S, F1, F2} <: AbstractJumpAggregator end function DiscreteCallback(c::AbstractSSAJumpAggregator) DiscreteCallback(c, c, initialize = c, save_positions = c.save_positions) @@ -55,8 +53,8 @@ end nothing end -@inline function concretize_affects!(p::AbstractSSAJumpAggregator{T, S, F1, F2}, - ::I) where {T, S, F1, F2 <: Tuple, I <: SciMLBase.DEIntegrator} +@inline function concretize_affects!(p::AbstractSSAJumpAggregator{<:Any, <:Any, <:Any, F2}, + ::I) where {F2 <: Tuple, I <: SciMLBase.DEIntegrator} nothing end @@ -88,8 +86,8 @@ function (p::AbstractSSAJumpAggregator)(integrator::I) where {I <: SciMLBase.DEI end function (p::AbstractSSAJumpAggregator{ - T, S, F1, F2})(integrator::SciMLBase.DEIntegrator) where - {T, S, F1, F2 <: Union{Tuple, Nothing}} + <:Any, <:Any, <:Any, F2})(integrator::SciMLBase.DEIntegrator) where + {F2 <: Union{Tuple, Nothing}} execute_jumps!(p, integrator, integrator.u, integrator.p, integrator.t, p.affects!) generate_jumps!(p, integrator, integrator.u, integrator.p, integrator.t) register_next_jump_time!(integrator, p, integrator.t) @@ -112,12 +110,12 @@ end """ build_jump_aggregation(jump_agg_type, u, p, t, end_time, ma_jumps, rates, - affects!, save_positions, rng; kwargs...) + affects!, save_positions; kwargs...) Helper routine for setting up standard fields of SSA jump aggregations. """ function build_jump_aggregation(jump_agg_type, u, p, t, end_time, ma_jumps, rates, - affects!, save_positions, rng; kwargs...) + affects!, save_positions; kwargs...) # mass action jumps majumps = ma_jumps @@ -134,7 +132,7 @@ function build_jump_aggregation(jump_agg_type, u, p, t, end_time, ma_jumps, rate next_jump = 0 next_jump_time = typemax(typeof(t)) jump_agg_type(next_jump, next_jump_time, end_time, cur_rates, sum_rate, - majumps, rates, affects!, save_positions, rng; kwargs...) + majumps, rates, affects!, save_positions; kwargs...) end """ diff --git a/src/problem.jl b/src/problem.jl index 157c97a30..e09d92ddf 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -51,7 +51,6 @@ $(FIELDS) ## Keyword Arguments - - `rng`, the random number generator to use. Defaults to Julia's built-in generator. - `save_positions=(true,true)` when including variable rates and `(false,true)` for constant rates, specifies whether to save the system's state (before, after) the jump occurs. - `spatial_system`, for spatial problems the underlying spatial structure. @@ -61,14 +60,20 @@ $(FIELDS) integration interface, and treated like general `VariableRateJump`s. - `vr_aggregator`, indicates the aggregator to use for sampling variable rate jumps. Current default is `VR_FRM`. + - `rng`, a random number generator to use for jump sampling. If provided, it is stored in + the problem's solver keyword arguments and forwarded to the solver when `solve` or `init` + is called. The recommended approach is to pass `rng` directly to `solve` or `init`: + `solve(jprob, SSAStepper(); rng = my_rng)`. + - `tstops`, time stops to pass through to the solver. Can be an `AbstractVector` of times + or a callable `(p, tspan) -> times`. Please see the [tutorial page](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) in the DifferentialEquations.jl [docs](https://docs.sciml.ai/JumpProcesses/stable/) for usage examples and commonly asked questions. """ -mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggregator}, J1, - J2, J3, J4, R, K} <: DiffEqBase.AbstractJumpProblem{P, J} +mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggregator}, J1, + J2, J3, J4, K} <: DiffEqBase.AbstractJumpProblem{P, J} """The type of problem to couple the jumps to. For a pure jump process use `DiscreteProblem`, to couple to ODEs, `ODEProblem`, etc.""" prob::P """The aggregator algorithm that determines the next jump times and types for `ConstantRateJump`s and `MassActionJump`s. Examples include `Direct`.""" @@ -85,26 +90,24 @@ mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggrega regular_jump::J3 """The `MassActionJump`s.""" massaction_jump::J4 - """The random number generator to use.""" - rng::R """kwargs to pass on to solve call.""" kwargs::K end function JumpProblem(p::P, a::A, dj::J, jc::C, cj::J1, vj::J2, rj::J3, mj::J4, - rng::R, kwargs::K) where {P, A, J, C, J1, J2, J3, J4, R, K} + kwargs::K) where {P, A, J, C, J1, J2, J3, J4, K} iip = isinplace_jump(p, rj) - JumpProblem{iip, P, A, C, J, J1, J2, J3, J4, R, K}(p, a, dj, jc, cj, vj, rj, mj, - rng, kwargs) + JumpProblem{iip, P, A, C, J, J1, J2, J3, J4, K}(p, a, dj, jc, cj, vj, rj, mj, + kwargs) end ######## remaking ###### # for a problem where prob.u0 is an ExtendedJumpArray, create an ExtendedJumpArray that # aliases and resets prob.u0.jump_u while having newu0 as the new u component. -function remake_extended_u0(prob, newu0, rng) +function remake_extended_u0(prob, newu0) jump_u = prob.u0.jump_u ttype = eltype(prob.tspan) - @. jump_u = -randexp(rng, ttype) + @. jump_u = zero(ttype) ExtendedJumpArray(newu0, jump_u) end @@ -142,7 +145,7 @@ function DiffEqBase.remake(jprob::JumpProblem; u0 = missing, p = missing, error("Passed in u0 is incompatible with current u0 which has type: $(typeof(prob.u0.u)).") end - final_u0 = remake_extended_u0(prob, state_vals, jprob.rng) + final_u0 = remake_extended_u0(prob, state_vals) end newprob = DiffEqBase.remake(prob; u0 = final_u0, p, interpret_symbolicmap, use_defaults, kwargs...) else @@ -169,8 +172,8 @@ function DiffEqBase.remake(jprob::JumpProblem; u0 = missing, p = missing, end T(newprob, jprob.aggregator, jprob.discrete_jump_aggregation, jprob.jump_callback, - jprob.constant_jumps, jprob.variable_jumps, jprob.regular_jump, - jprob.massaction_jump, jprob.rng, jprob.kwargs) + jprob.constant_jumps, jprob.variable_jumps, jprob.regular_jump, + jprob.massaction_jump, jprob.kwargs) end # for updating parameters in JumpProblems to update MassActionJumps @@ -239,7 +242,7 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS vr_aggregator::VariableRateAggregator = VR_FRM(), save_positions = prob isa DiffEqBase.AbstractDiscreteProblem ? (false, true) : (true, true), - rng = DEFAULT_RNG, scale_rates = true, useiszero = true, + rng = nothing, scale_rates = true, useiszero = true, spatial_system = nothing, hopping_constants = nothing, callback = nothing, tstops = nothing, use_vrj_bounds = true, kwargs...) @@ -289,15 +292,15 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS constant_jump_callback = CallbackSet() else disc_agg = aggregate(aggregator, u, prob.p, t, end_time, crjs, maj, - save_positions, rng; kwargs...) + save_positions; kwargs...) constant_jump_callback = DiscreteCallback(disc_agg) end # handle any remaining vrjs if length(cvrjs) > 0 # Handle variable rate jumps based on vr_aggregator - new_prob, variable_jump_callback = configure_jump_problem(prob, vr_aggregator, - jumps, cvrjs; rng) + new_prob, variable_jump_callback = configure_jump_problem(prob, vr_aggregator, + jumps, cvrjs) else new_prob = prob variable_jump_callback = CallbackSet() @@ -306,19 +309,19 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS jump_cbs = CallbackSet(constant_jump_callback, variable_jump_callback) iip = isinplace_jump(prob, jumps.regular_jump) - solkwargs = tstops === nothing ? make_kwarg(; callback) : make_kwarg(; callback, tstops) + solkwargs = make_kwarg(; filter(!isnothing, (; callback, tstops, rng))...) JumpProblem{iip, typeof(new_prob), typeof(aggregator), typeof(jump_cbs), typeof(disc_agg), typeof(crjs), typeof(cvrjs), typeof(jumps.regular_jump), - typeof(maj), typeof(rng), typeof(solkwargs)}(new_prob, aggregator, disc_agg, - jump_cbs, crjs, cvrjs, jumps.regular_jump, maj, rng, solkwargs) + typeof(maj), typeof(solkwargs)}(new_prob, aggregator, disc_agg, + jump_cbs, crjs, cvrjs, jumps.regular_jump, maj, 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, + rng = nothing, scale_rates = true, useiszero = true, spatial_system = nothing, hopping_constants = nothing, callback = nothing, tstops = nothing, kwargs...) @@ -342,18 +345,18 @@ function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; # 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 = tstops === nothing ? make_kwarg(; callback) : make_kwarg(; callback, tstops) + solkwargs = make_kwarg(; filter(!isnothing, (; callback, tstops, rng))...) 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) + typeof(maj), typeof(solkwargs)}(prob, aggregator, disc_agg, + jump_cbs, crjs, vrjs, jumps.regular_jump, maj, solkwargs) end aggregator(jp::JumpProblem{iip, P, A}) where {iip, P, A} = A diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index da98ba1c4..4d817b035 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -70,13 +70,14 @@ function _process_saveat(saveat, tspan, save_start, save_end) end function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping; - seed = nothing, dt = error("dt is required for SimpleTauLeaping."), + seed = nothing, rng = nothing, + dt = error("dt is required for SimpleTauLeaping."), saveat = nothing, save_start = nothing, save_end = nothing) validate_pure_leaping_inputs(jump_prob, alg) || error("SimpleTauLeaping can only be used with PureLeaping JumpProblems with only RegularJumps.") - (; prob, rng) = jump_prob - (seed !== nothing) && seed!(rng, seed) + prob = jump_prob.prob + _rng = resolve_rng(rng, seed, get(jump_prob.kwargs, :rng, nothing)) rj = jump_prob.regular_jump rate = rj.rate # rate function rate(out,u,p,t) @@ -117,7 +118,7 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping; t_new = tprev + dt rate(rate_cache, uprev, p, tprev) rate_cache .*= dt - counts .= pois_rand.((rng,), rate_cache) + counts .= pois_rand.((_rng,), rate_cache) c(du, uprev, p, tprev, counts, mark) u_new .= du .+ uprev @@ -335,22 +336,20 @@ function simple_explicit_tau_leaping_loop!( end function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; - seed = nothing, + seed = nothing, rng = nothing, dtmin = nothing, saveat = nothing, save_start = nothing, save_end = nothing) validate_pure_leaping_inputs(jump_prob, alg) || error("SimpleExplicitTauLeaping can only be used with PureLeaping JumpProblem with a MassActionJump.") prob = jump_prob.prob - rng = jump_prob.rng + _rng = resolve_rng(rng, seed, get(jump_prob.kwargs, :rng, nothing)) tspan = prob.tspan if dtmin === nothing dtmin = 1e-10 * one(typeof(tspan[2])) end - (seed !== nothing) && seed!(rng, seed) - maj = jump_prob.massaction_jump numjumps = get_num_majumps(maj) rj = jump_prob.regular_jump @@ -394,7 +393,7 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; reactant_stoch, hor, length(u0), numjumps) simple_explicit_tau_leaping_loop!( - prob, alg, u_current, u_new, t_current, t_end, p, rng, + prob, alg, u_current, u_new, t_current, t_end, p, _rng, rate, c, nu, hor, max_hor, max_stoich, numjumps, epsilon, dtmin, saveat_times, usave, tsave, du, counts, rate_cache, rate_effective, maj, save_end) diff --git a/src/solve.jl b/src/solve.jl index 11bc16bb6..1ae5b3842 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -1,3 +1,22 @@ +""" + resolve_rng(rng, seed[, fallback_rng]) + +Resolve which RNG to use for a jump simulation. + +Priority: `rng` > `seed` (creates `Xoshiro`) > `fallback_rng` > `Random.default_rng()`. +""" +function resolve_rng(rng, seed, fallback_rng = nothing) + if rng !== nothing + rng + elseif seed !== nothing + Random.Xoshiro(seed) + elseif fallback_rng !== nothing + fallback_rng + else + Random.default_rng() + end +end + function DiffEqBase.__solve(jump_prob::DiffEqBase.AbstractJumpProblem{P}, alg::DiffEqBase.DEAlgorithm; merge_callbacks = true, kwargs...) where {P} @@ -38,53 +57,33 @@ function DiffEqBase.__init(_jump_prob::DiffEqBase.AbstractJumpProblem{P}, kwargs = DiffEqBase.merge_problem_kwargs(_jump_prob; merge_callbacks, kwargs...) __jump_init(_jump_prob, alg; kwargs...) -end +end function __jump_init(_jump_prob::DiffEqBase.AbstractJumpProblem{P}, alg; - callback = nothing, seed = nothing, + callback = nothing, seed = nothing, rng = nothing, alias_jump = Threads.threadid() == 1, kwargs...) where {P} + + _rng = resolve_rng(rng, seed) + if alias_jump jump_prob = _jump_prob - reset_jump_problem!(jump_prob, seed) else - jump_prob = resetted_jump_problem(_jump_prob, seed) + jump_prob = resetted_jump_problem(_jump_prob) end - # DDEProblems do not have a recompile_flag argument - if jump_prob.prob isa DiffEqBase.AbstractDDEProblem - # callback comes after jump consistent with SSAStepper - integrator = init(jump_prob.prob, alg; - callback = CallbackSet(jump_prob.jump_callback, callback), - kwargs...) - else - # callback comes after jump consistent with SSAStepper - integrator = init(jump_prob.prob, alg; - callback = CallbackSet(jump_prob.jump_callback, callback), - kwargs...) - end + init(jump_prob.prob, alg; + callback = CallbackSet(jump_prob.jump_callback, callback), + rng = _rng, kwargs...) end -# Derive an independent seed from the caller's seed. When a caller (e.g. StochasticDiffEq) -# passes the same seed used for its noise process, we must produce a distinct seed for the -# jump aggregator's RNG. We cannot assume the JumpProblem's stored RNG is any particular -# type, so we pass the seed through `hash` (to decorrelate from the input) and then through -# a Xoshiro draw (to ensure strong mixing regardless of the target RNG's seeding quality). -const _JUMP_SEED_SALT = 0x4a756d7050726f63 # "JumPProc" in ASCII -_derive_jump_seed(seed) = rand(Random.Xoshiro(hash(seed, _JUMP_SEED_SALT)), UInt64) - -function resetted_jump_problem(_jump_prob, seed) - jump_prob = deepcopy(_jump_prob) - if seed !== nothing && !isempty(jump_prob.jump_callback.discrete_callbacks) - rng = jump_prob.jump_callback.discrete_callbacks[1].condition.rng - Random.seed!(rng, _derive_jump_seed(seed)) - end - jump_prob +# Keep function signatures for StochasticDiffEq backward compatibility. +# The seed argument is accepted but no longer used to reseed aggregator RNGs +# (RNG state is now managed by the integrator). +function resetted_jump_problem(_jump_prob, seed = nothing) + deepcopy(_jump_prob) end -function reset_jump_problem!(jump_prob, seed) - if seed !== nothing && !isempty(jump_prob.jump_callback.discrete_callbacks) - Random.seed!(jump_prob.jump_callback.discrete_callbacks[1].condition.rng, - _derive_jump_seed(seed)) - end +function reset_jump_problem!(jump_prob, seed = nothing) + nothing end diff --git a/src/spatial/directcrdirect.jl b/src/spatial/directcrdirect.jl index f282bd1ae..0961ec25b 100644 --- a/src/spatial/directcrdirect.jl +++ b/src/spatial/directcrdirect.jl @@ -5,10 +5,10 @@ const MINJUMPRATE = 2.0^exponent(1e-12) #NOTE state vector u is a matrix. u[i,j] is species i, site j #NOTE hopping_constants is a matrix. hopping_constants[i,j] is species i, site j -mutable struct DirectCRDirectJumpAggregation{T, S, F1, F2, RNG, J, RX, HOP, DEPGR, +mutable struct DirectCRDirectJumpAggregation{T, S, F1, F2, J, RX, HOP, DEPGR, VJMAP, JVMAP, SS, U <: PriorityTable, W <: Function} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::SpatialJump{J} #some structure to identify the next event: reaction or hop prev_jump::SpatialJump{J} #some structure to identify the previous event: reaction or hop next_jump_time::T @@ -19,7 +19,6 @@ mutable struct DirectCRDirectJumpAggregation{T, S, F1, F2, RNG, J, RX, HOP, DEPG rates::F1 # legacy, not used affects!::F2 # legacy, not used save_positions::Tuple{Bool, Bool} - rng::RNG dep_gr::DEPGR #dep graph is same for each site vartojumps_map::VJMAP #vartojumps_map is same for each site jumptovars_map::JVMAP #jumptovars_map is same for each site @@ -31,11 +30,11 @@ end function DirectCRDirectJumpAggregation(nj::SpatialJump{J}, njt::T, et::T, rx_rates::RX, hop_rates::HOP, site_rates::Vector{T}, - sps::Tuple{Bool, Bool}, rng::RNG, spatial_system::SS; + sps::Tuple{Bool, Bool}, spatial_system::SS; num_specs, minrate = convert(T, MINJUMPRATE), vartojumps_map = nothing, jumptovars_map = nothing, dep_graph = nothing, - kwargs...) where {J, T, RX, HOP, RNG, SS} + kwargs...) where {J, T, RX, HOP, SS} # a dependency graph is needed if dep_graph === nothing @@ -69,12 +68,12 @@ function DirectCRDirectJumpAggregation(nj::SpatialJump{J}, njt::T, et::T, rx_rat # construct an empty initial priority table -- we'll reset this in init rt = PriorityTable(ratetogroup, zeros(T, 1), minrate, 2 * minrate) - DirectCRDirectJumpAggregation{T, Nothing, Nothing, Nothing, RNG, J, RX, HOP, + DirectCRDirectJumpAggregation{T, Nothing, Nothing, Nothing, J, RX, HOP, typeof(dg), typeof(vtoj_map), typeof(jtov_map), SS, typeof(rt), typeof(ratetogroup)}(nj, nj, njt, et, rx_rates, hop_rates, site_rates, nothing, nothing, sps, - rng, dg, vtoj_map, + dg, vtoj_map, jtov_map, spatial_system, num_specs, rt, ratetogroup) end @@ -82,7 +81,7 @@ end ############################# Required Functions ############################## # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::DirectCRDirect, starting_state, p, t, end_time, - constant_jumps, ma_jumps, save_positions, rng; hopping_constants, + constant_jumps, ma_jumps, save_positions; hopping_constants, spatial_system, kwargs...) num_species = size(starting_state, 1) majumps = ma_jumps @@ -99,7 +98,7 @@ function aggregate(aggregator::DirectCRDirect, starting_state, p, t, end_time, site_rates = zeros(typeof(end_time), num_sites(spatial_system)) DirectCRDirectJumpAggregation(next_jump, next_jump_time, end_time, rx_rates, hop_rates, - site_rates, save_positions, rng, spatial_system; + site_rates, save_positions, spatial_system; num_specs = num_species, kwargs...) end @@ -113,10 +112,11 @@ end # calculate the next jump / jump time function generate_jumps!(p::DirectCRDirectJumpAggregation, integrator, params, u, t) - p.next_jump_time = t + randexp(p.rng) / p.rt.gsum + rng = get_rng(integrator) + p.next_jump_time = t + randexp(rng) / p.rt.gsum p.next_jump_time >= p.end_time && return nothing - site = sample(p.rt, p.site_rates, p.rng) - p.next_jump = sample_jump_direct(p, site) + site = sample(p.rt, p.site_rates, rng) + p.next_jump = sample_jump_direct(p, site, rng) nothing end diff --git a/src/spatial/nsm.jl b/src/spatial/nsm.jl index 07c48da6e..acf6b56eb 100644 --- a/src/spatial/nsm.jl +++ b/src/spatial/nsm.jl @@ -3,9 +3,9 @@ ############################ NSM ################################### #NOTE state vector u is a matrix. u[i,j] is species i, site j #NOTE hopping_constants is a matrix. hopping_constants[i,j] is species i, site j -mutable struct NSMJumpAggregation{T, S, F1, F2, RNG, J, RX, HOP, DEPGR, VJMAP, JVMAP, +mutable struct NSMJumpAggregation{T, S, F1, F2, J, RX, HOP, DEPGR, VJMAP, JVMAP, PQ, SS} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::SpatialJump{J} #some structure to identify the next event: reaction or hop prev_jump::SpatialJump{J} #some structure to identify the previous event: reaction or hop next_jump_time::T @@ -15,7 +15,6 @@ mutable struct NSMJumpAggregation{T, S, F1, F2, RNG, J, RX, HOP, DEPGR, VJMAP, J rates::F1 # legacy, not used affects!::F2 # legacy, not used save_positions::Tuple{Bool, Bool} - rng::RNG dep_gr::DEPGR #dep graph is same for each site vartojumps_map::VJMAP #vartojumps_map is same for each site jumptovars_map::JVMAP #jumptovars_map is same for each site @@ -27,9 +26,9 @@ end function NSMJumpAggregation( nj::SpatialJump{J}, njt::T, et::T, rx_rates::RX, hop_rates::HOP, sps::Tuple{Bool, Bool}, - rng::RNG, spatial_system::SS; num_specs, + spatial_system::SS; num_specs, vartojumps_map = nothing, jumptovars_map = nothing, - dep_graph = nothing, kwargs...) where {J, T, RX, HOP, RNG, SS} + dep_graph = nothing, kwargs...) where {J, T, RX, HOP, SS} # a dependency graph is needed if dep_graph === nothing @@ -55,13 +54,13 @@ function NSMJumpAggregation( pq = MutableBinaryMinHeap{T}() - NSMJumpAggregation{T, Nothing, Nothing, Nothing, RNG, J, RX, HOP, typeof(dg), + NSMJumpAggregation{T, Nothing, Nothing, Nothing, J, RX, HOP, typeof(dg), typeof(vtoj_map), typeof(jtov_map), typeof(pq), SS}(nj, nj, njt, et, rx_rates, hop_rates, nothing, nothing, - sps, rng, dg, + sps, dg, vtoj_map, jtov_map, pq, @@ -72,7 +71,7 @@ end ############################# Required Functions ############################## # creating the JumpAggregation structure (function wrapper-based constant jumps) function aggregate(aggregator::NSM, starting_state, p, t, end_time, constant_jumps, - ma_jumps, save_positions, rng; hopping_constants, spatial_system, + ma_jumps, save_positions; hopping_constants, spatial_system, kwargs...) num_species = size(starting_state, 1) majumps = ma_jumps @@ -88,7 +87,7 @@ function aggregate(aggregator::NSM, starting_state, p, t, end_time, constant_jum hop_rates = HopRates(hopping_constants, spatial_system) NSMJumpAggregation(next_jump, next_jump_time, end_time, rx_rates, hop_rates, - save_positions, rng, spatial_system; num_specs = num_species, + save_positions, spatial_system; num_specs = num_species, kwargs...) end @@ -104,7 +103,8 @@ end function generate_jumps!(p::NSMJumpAggregation, integrator, params, u, t) p.next_jump_time, site = top_with_handle(p.pq) p.next_jump_time >= p.end_time && return nothing - p.next_jump = sample_jump_direct(p, site) + rng = get_rng(integrator) + p.next_jump = sample_jump_direct(p, site, rng) nothing end @@ -125,7 +125,8 @@ end reset all structs, reevaluate all rates, recalculate tentative site firing times, and reinit the priority queue """ function fill_rates_and_get_times!(aggregation::NSMJumpAggregation, integrator, t) - (; spatial_system, rx_rates, hop_rates, rng) = aggregation + (; spatial_system, rx_rates, hop_rates) = aggregation + rng = get_rng(integrator) u = integrator.u reset!(rx_rates) @@ -153,30 +154,31 @@ recalculate jump rates for jumps that depend on the just executed jump (p.prev_j """ function update_dependent_rates_and_firing_times!(p::NSMJumpAggregation, integrator, t) u = integrator.u + rng = get_rng(integrator) jump = p.prev_jump if is_hop(p, jump) source_site = jump.src target_site = jump.dst update_rates_after_hop!(p, integrator, source_site, target_site, jump.jidx) - update_site_time!(p, source_site, t) - update_site_time!(p, target_site, t) + update_site_time!(p, source_site, t, rng) + update_site_time!(p, target_site, t, rng) else site = jump.src update_rates_after_reaction!(p, integrator, site, reaction_id_from_jump(p, jump)) - update_site_time!(p, site, t) + update_site_time!(p, site, t, rng) end nothing end """ - update_site_time!(p::NSMJumpAggregation, site, t) + update_site_time!(p::NSMJumpAggregation, site, t, rng) update the time of site in the priority queue """ -function update_site_time!(p::NSMJumpAggregation, site, t) +function update_site_time!(p::NSMJumpAggregation, site, t, rng) site_rate = (total_site_rate(p.rx_rates, p.hop_rates, site)) if site_rate > zero(typeof(site_rate)) - update!(p.pq, site, t + randexp(p.rng) / site_rate) + update!(p.pq, site, t + randexp(rng) / site_rate) else update!(p.pq, site, typemax(t)) end diff --git a/src/spatial/utils.jl b/src/spatial/utils.jl index 370e86015..3be2b7378 100644 --- a/src/spatial/utils.jl +++ b/src/spatial/utils.jl @@ -23,18 +23,18 @@ end ######################## helper routines for all spatial SSAs ######################## """ - sample_jump_direct(p, site) + sample_jump_direct(p, site, rng) sample jump at site with direct method """ -function sample_jump_direct(p, site) - if rand(p.rng) * (total_site_rate(p.rx_rates, p.hop_rates, site)) < +function sample_jump_direct(p, site, rng) + if rand(rng) * (total_site_rate(p.rx_rates, p.hop_rates, site)) < total_site_rx_rate(p.rx_rates, site) - rx = sample_rx_at_site(p.rx_rates, site, p.rng) + rx = sample_rx_at_site(p.rx_rates, site, rng) return SpatialJump(site, rx + p.numspecies, site) else species_to_diffuse, - target_site = sample_hop_at_site(p.hop_rates, site, p.rng, + target_site = sample_hop_at_site(p.hop_rates, site, rng, p.spatial_system) return SpatialJump(site, species_to_diffuse, target_site) end diff --git a/src/variable_rate.jl b/src/variable_rate.jl index 48c405343..beed2a0af 100644 --- a/src/variable_rate.jl +++ b/src/variable_rate.jl @@ -28,21 +28,21 @@ Simulating a birth-death process with `VR_FRM`: ```julia using JumpProcesses, OrdinaryDiffEq -u0 = [1.0] # Initial population -p = [10.0, 0.5] # [birth rate, death rate] +u0 = [1.0] # Initial population +p = [10.0, 0.5] # [birth rate, death rate] tspan = (0.0, 10.0) -# Birth jump: ∅ → X +# Birth jump: ∅ → X birth_rate(u, p, t) = p[1] birth_affect!(integrator) = (integrator.u[1] += 1; nothing) birth_jump = VariableRateJump(birth_rate, birth_affect!) -# Death jump: X → ∅ +# Death jump: X → ∅ death_rate(u, p, t) = p[2] * u[1] death_affect!(integrator) = (integrator.u[1] -= 1; nothing) death_jump = VariableRateJump(death_rate, death_affect!) -# Problem setup +# Problem setup oprob = ODEProblem((du, u, p, t) -> du .= 0, u0, tspan, p) jprob = JumpProblem(oprob, birth_jump, death_jump; vr_aggregator = VR_FRM()) sol = solve(jprob, Tsit5()) @@ -58,26 +58,25 @@ sol = solve(jprob, Tsit5()) """ struct VR_FRM <: VariableRateAggregator end -function configure_jump_problem(prob, vr_aggregator::VR_FRM, jumps, cvrjs; - rng = DEFAULT_RNG) - new_prob = extend_problem(prob, cvrjs; rng) - variable_jump_callback = build_variable_callback(CallbackSet(), 0, cvrjs...; rng) +function configure_jump_problem(prob, vr_aggregator::VR_FRM, jumps, cvrjs) + new_prob = extend_problem(prob, cvrjs) + variable_jump_callback = build_variable_callback(CallbackSet(), 0, cvrjs...) return new_prob, variable_jump_callback end # extends prob.u0 to an ExtendedJumpArray with Njumps integrated intensity values, # of type prob.tspan -function extend_u0(prob, Njumps, rng) +function extend_u0(prob, Njumps) ttype = eltype(prob.tspan) - u0 = ExtendedJumpArray(prob.u0, [-randexp(rng, ttype) for i in 1:Njumps]) + u0 = ExtendedJumpArray(prob.u0, zeros(ttype, Njumps)) return u0 end -function extend_problem(prob::DiffEqBase.AbstractDiscreteProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::DiffEqBase.AbstractDiscreteProblem, jumps) error("General `VariableRateJump`s require a continuous problem, like an ODE/SDE/DDE/DAE problem. To use a `DiscreteProblem` bounded `VariableRateJump`s must be used. See the JumpProcesses docs.") end -function extend_problem(prob::DiffEqBase.AbstractODEProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::DiffEqBase.AbstractODEProblem, jumps) _f = SciMLBase.unwrapped_f(prob.f) if isinplace(prob) @@ -97,13 +96,13 @@ function extend_problem(prob::DiffEqBase.AbstractODEProblem, jumps; rng = DEFAUL end end - u0 = extend_u0(prob, length(jumps), rng) + u0 = extend_u0(prob, length(jumps)) f = ODEFunction{isinplace(prob)}(jump_f; sys = prob.f.sys, observed = prob.f.observed) remake(prob; f, u0) end -function extend_problem(prob::DiffEqBase.AbstractSDEProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::DiffEqBase.AbstractSDEProblem, jumps) _f = SciMLBase.unwrapped_f(prob.f) if isinplace(prob) @@ -133,13 +132,13 @@ function extend_problem(prob::DiffEqBase.AbstractSDEProblem, jumps; rng = DEFAUL end end - u0 = extend_u0(prob, length(jumps), rng) + u0 = extend_u0(prob, length(jumps)) f = SDEFunction{isinplace(prob)}(jump_f, jump_g; sys = prob.f.sys, observed = prob.f.observed) remake(prob; f, g = jump_g, u0) end -function extend_problem(prob::DiffEqBase.AbstractDDEProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::DiffEqBase.AbstractDDEProblem, jumps) _f = SciMLBase.unwrapped_f(prob.f) if isinplace(prob) @@ -159,14 +158,14 @@ function extend_problem(prob::DiffEqBase.AbstractDDEProblem, jumps; rng = DEFAUL end end - u0 = extend_u0(prob, length(jumps), rng) + u0 = extend_u0(prob, length(jumps)) f = DDEFunction{isinplace(prob)}(jump_f; sys = prob.f.sys, observed = prob.f.observed) remake(prob; f, u0) end # Not sure if the DAE one is correct: Should be a residual of sorts -function extend_problem(prob::DiffEqBase.AbstractDAEProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::DiffEqBase.AbstractDAEProblem, jumps) _f = SciMLBase.unwrapped_f(prob.f) if isinplace(prob) @@ -186,16 +185,15 @@ function extend_problem(prob::DiffEqBase.AbstractDAEProblem, jumps; rng = DEFAUL end end - u0 = extend_u0(prob, length(jumps), rng) + u0 = extend_u0(prob, length(jumps)) f = DAEFunction{isinplace(prob)}(jump_f, sys = prob.f.sys, observed = prob.f.observed) remake(prob; f, u0) end -struct VR_FRMEventCallback{F, RNG} +struct VR_FRMEventCallback{F} idx::Int affect!::F - rng::RNG end # condition: (u, t, integrator) @@ -204,20 +202,22 @@ end # affect: (integrator) function (c::VR_FRMEventCallback)(integrator) c.affect!(integrator) - integrator.u.jump_u[c.idx] = -randexp(c.rng, typeof(integrator.t)) + rng = get_rng(integrator) + integrator.u.jump_u[c.idx] = -randexp(rng, typeof(integrator.t)) nothing end # initialize: (cb, u, t, integrator) function (c::VR_FRMEventCallback)(cb, u, t, integrator) - integrator.u.jump_u[c.idx] = -randexp(c.rng, typeof(integrator.t)) + rng = get_rng(integrator) + integrator.u.jump_u[c.idx] = -randexp(rng, typeof(integrator.t)) integrator.uprev.jump_u[c.idx] = integrator.u.jump_u[c.idx] u_modified!(integrator, false) nothing end -function wrap_jump_in_callback(idx, jump; rng = DEFAULT_RNG) - cb_functor = VR_FRMEventCallback(idx, jump.affect!, rng) +function wrap_jump_in_callback(idx, jump) + cb_functor = VR_FRMEventCallback(idx, jump.affect!) ContinuousCallback(cb_functor, cb_functor; initialize = cb_functor, idxs = jump.idxs, @@ -228,15 +228,15 @@ function wrap_jump_in_callback(idx, jump; rng = DEFAULT_RNG) reltol = jump.reltol) end -function build_variable_callback(cb, idx, jump, jumps...; rng = DEFAULT_RNG) +function build_variable_callback(cb, idx, jump, jumps...) idx += 1 - new_cb = wrap_jump_in_callback(idx, jump; rng) - build_variable_callback(CallbackSet(cb, new_cb), idx, jumps...; rng = DEFAULT_RNG) + new_cb = wrap_jump_in_callback(idx, jump) + build_variable_callback(CallbackSet(cb, new_cb), idx, jumps...) end -function build_variable_callback(cb, idx, jump; rng = DEFAULT_RNG) +function build_variable_callback(cb, idx, jump) idx += 1 - CallbackSet(cb, wrap_jump_in_callback(idx, jump; rng)) + CallbackSet(cb, wrap_jump_in_callback(idx, jump)) end @inline function update_jumps!(du, u, p, t, idx, jump) @@ -268,21 +268,21 @@ Simulating a birth-death process with `VR_Direct` (default) and VR_DirectFW: ```julia using JumpProcesses, OrdinaryDiffEq -u0 = [1.0] # Initial population -p = [10.0, 0.5] # [birth rate, death rate coefficient] +u0 = [1.0] # Initial population +p = [10.0, 0.5] # [birth rate, death rate coefficient] tspan = (0.0, 10.0) -# Birth jump: ∅ → X +# Birth jump: ∅ → X birth_rate(u, p, t) = p[1] birth_affect!(integrator) = (integrator.u[1] += 1; nothing) birth_jump = VariableRateJump(birth_rate, birth_affect!) -# Death jump: X → ∅ +# Death jump: X → ∅ death_rate(u, p, t) = p[2] * u[1] death_affect!(integrator) = (integrator.u[1] -= 1; nothing) death_jump = VariableRateJump(death_rate, death_affect!) -# Problem setup +# Problem setup oprob = ODEProblem((du, u, p, t) -> du .= 0, u0, tspan, p) jprob = JumpProblem(oprob, birth_jump, death_jump; vr_aggregator = VR_Direct()) sol = solve(jprob, Tsit5()) @@ -298,21 +298,19 @@ sol = solve(jprob, Tsit5()) struct VR_Direct <: VariableRateAggregator end struct VR_DirectFW <: VariableRateAggregator end -mutable struct VR_DirectEventCache{T, RNG, F1, F2} +mutable struct VR_DirectEventCache{T, F1, F2} prev_time::T prev_threshold::T current_time::T current_threshold::T total_rate::T - rng::RNG rate_funcs::F1 affect_funcs::F2 cum_rate_sum::Vector{T} end function VR_DirectEventCache( - jumps::JumpSet, ::VR_Direct, prob, ::Type{T}; rng = DEFAULT_RNG) where {T} - initial_threshold = randexp(rng, T) + jumps::JumpSet, ::VR_Direct, prob, ::Type{T}) where {T} vjumps = jumps.variable_jumps # handle vjumps using tuples @@ -320,14 +318,13 @@ function VR_DirectEventCache( cum_rate_sum = Vector{T}(undef, length(vjumps)) - VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), typeof(affect_funcs)}(zero(T), - initial_threshold, zero(T), initial_threshold, zero(T), rng, rate_funcs, + VR_DirectEventCache{T, typeof(rate_funcs), typeof(affect_funcs)}(zero(T), + zero(T), zero(T), zero(T), zero(T), rate_funcs, affect_funcs, cum_rate_sum) end function VR_DirectEventCache( - jumps::JumpSet, ::VR_DirectFW, prob, ::Type{T}; rng = DEFAULT_RNG) where {T} - initial_threshold = randexp(rng, T) + jumps::JumpSet, ::VR_DirectFW, prob, ::Type{T}) where {T} vjumps = jumps.variable_jumps t, u = prob.tspan[1], prob.u0 @@ -337,16 +334,17 @@ function VR_DirectEventCache( cum_rate_sum = Vector{T}(undef, length(vjumps)) - VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), Any}(zero(T), - initial_threshold, zero(T), initial_threshold, zero(T), rng, rate_funcs, + VR_DirectEventCache{T, typeof(rate_funcs), Any}(zero(T), + zero(T), zero(T), zero(T), zero(T), rate_funcs, affect_funcs, cum_rate_sum) end # Initialization function for VR_DirectEventCache function initialize_vr_direct_cache!(cache::VR_DirectEventCache, u, t, integrator) + rng = get_rng(integrator) cache.prev_time = zero(integrator.t) cache.current_time = zero(integrator.t) - cache.prev_threshold = randexp(cache.rng, eltype(integrator.t)) + cache.prev_threshold = randexp(rng, eltype(integrator.t)) cache.current_threshold = cache.prev_threshold cache.total_rate = zero(integrator.t) cache.cum_rate_sum .= 0 @@ -364,8 +362,8 @@ end nothing end -@inline function concretize_vr_direct_affects!(cache::VR_DirectEventCache{T, RNG, F1, F2}, - ::I) where {T, RNG, F1, F2 <: Tuple, I <: SciMLBase.DEIntegrator} +@inline function concretize_vr_direct_affects!(cache::VR_DirectEventCache{T, F1, F2}, + ::I) where {T, F1, F2 <: Tuple, I <: SciMLBase.DEIntegrator} nothing end @@ -393,16 +391,16 @@ function build_variable_integcallback(cache::VR_DirectEventCache, jumps) save_positions, abstol, reltol) end -function configure_jump_problem(prob, ::VR_Direct, jumps, cvrjs; rng = DEFAULT_RNG) +function configure_jump_problem(prob, ::VR_Direct, jumps, cvrjs) new_prob = prob - cache = VR_DirectEventCache(jumps, VR_Direct(), prob, eltype(prob.tspan); rng) + cache = VR_DirectEventCache(jumps, VR_Direct(), prob, eltype(prob.tspan)) variable_jump_callback = build_variable_integcallback(cache, cvrjs) return new_prob, variable_jump_callback end -function configure_jump_problem(prob, ::VR_DirectFW, jumps, cvrjs; rng = DEFAULT_RNG) +function configure_jump_problem(prob, ::VR_DirectFW, jumps, cvrjs) new_prob = prob - cache = VR_DirectEventCache(jumps, VR_DirectFW(), prob, eltype(prob.tspan); rng) + cache = VR_DirectEventCache(jumps, VR_DirectFW(), prob, eltype(prob.tspan)) variable_jump_callback = build_variable_integcallback(cache, cvrjs) return new_prob, variable_jump_callback end @@ -435,8 +433,7 @@ end end function total_variable_rate( - cache::VR_DirectEventCache{ - T, RNG, F1, F2}, u, p, t) where {T, RNG, F1, F2} + cache::VR_DirectEventCache{T, F1, F2}, u, p, t) where {T, F1, F2} (; cum_rate_sum, rate_funcs) = cache sum_rate = cumsum_rates!(cum_rate_sum, u, p, t, rate_funcs) return sum_rate @@ -481,8 +478,8 @@ function (cache::VR_DirectEventCache)(u, t, integrator) return cache.current_threshold end -@generated function execute_affect!(cache::VR_DirectEventCache{T, RNG, F1, F2}, - integrator::I, idx) where {T, RNG, F1, F2 <: Tuple, I <: SciMLBase.DEIntegrator} +@generated function execute_affect!(cache::VR_DirectEventCache{T, F1, F2}, + integrator::I, idx) where {T, F1, F2 <: Tuple, I <: SciMLBase.DEIntegrator} quote (; affect_funcs) = cache Base.Cartesian.@nif $(fieldcount(F2)) i -> (i == idx) i -> (@inbounds affect_funcs[i](integrator)) i -> (@inbounds affect_funcs[fieldcount(F2)](integrator)) @@ -509,7 +506,7 @@ function (cache::VR_DirectEventCache)(integrator) end cache.total_rate = total_variable_rate_sum - rng = cache.rng + rng = get_rng(integrator) r = rand(rng) * total_variable_rate_sum @inbounds jump_idx = searchsortedfirst(cache.cum_rate_sum, r) execute_affect!(cache, integrator, jump_idx) diff --git a/test/degenerate_rx_cases.jl b/test/degenerate_rx_cases.jl index 3b211decd..55b2985e7 100644 --- a/test/degenerate_rx_cases.jl +++ b/test/degenerate_rx_cases.jl @@ -20,8 +20,8 @@ rs = [[0 => 1]] ns = [[1 => 1]] jump = MassActionJump(rate, rs, ns) prob = DiscreteProblem([100], (0.0, 100.0)) -jump_prob = JumpProblem(prob, Direct(), jump; rng = rng) -sol = solve(jump_prob, SSAStepper()) +jump_prob = JumpProblem(prob, Direct(), jump) +sol = solve(jump_prob, SSAStepper(); rng) if doprint println("mass act jump using vectors of data: last val = ", sol[end, end]) end @@ -35,8 +35,8 @@ rate = 2.0 rs = [0 => 3] # stoich power should be ignored ns = [1 => 1] jump = MassActionJump(rate, rs, ns) -jump_prob = JumpProblem(prob, Direct(), jump; rng = rng) -sol = solve(jump_prob, SSAStepper()) +jump_prob = JumpProblem(prob, Direct(), jump) +sol = solve(jump_prob, SSAStepper(); rng) if doprint println("mass act jump using scalar data: last val = ", sol[end, end]) end @@ -51,8 +51,8 @@ rs = [Vector{Pair{Int, Int}}()] ns = [[1 => 1]] jump = MassActionJump(rate, rs, ns) prob = DiscreteProblem([100], (0.0, 100.0)) -jump_prob = JumpProblem(prob, Direct(), jump; rng = rng) -sol = solve(jump_prob, SSAStepper()) +jump_prob = JumpProblem(prob, Direct(), jump) +sol = solve(jump_prob, SSAStepper(); rng) if doprint println("mass act jump using vector of Pair{Int,Int}: last val = ", sol[end, end]) end @@ -66,8 +66,8 @@ rate = 2.0 rs = Vector{Pair{Int, Int}}() ns = [1 => 1] jump = MassActionJump(rate, rs, ns) -jump_prob = JumpProblem(prob, Direct(), jump; rng = rng) -sol = solve(jump_prob, SSAStepper()) +jump_prob = JumpProblem(prob, Direct(), jump) +sol = solve(jump_prob, SSAStepper(); rng) if doprint println("mass act jump using scalar Pair{Int,Int}: last val = ", sol[end, end]) end @@ -100,8 +100,8 @@ jump_to_dep_specs = [[1], [1]] namedpars = (dep_graph = dep_graph, vartojumps_map = spec_to_dep_jumps, jumptovars_map = jump_to_dep_specs) for method in methods - local jump_prob = JumpProblem(prob, method, jump, jump2; rng = rng, namedpars...) - local sol = solve(jump_prob, SSAStepper()) + local jump_prob = JumpProblem(prob, method, jump, jump2; namedpars...) + local sol = solve(jump_prob, SSAStepper(); rng) if doplot plot!(plothand2, sol, label = ("A <-> 0, " * string(method))) diff --git a/test/ensemble_problems.jl b/test/ensemble_problems.jl index f042d515e..541e7cf12 100644 --- a/test/ensemble_problems.jl +++ b/test/ensemble_problems.jl @@ -6,30 +6,30 @@ using StableRNGs, Random # ========================================================================== # Constant-rate birth-death for SSAStepper / ODE-coupled tests -function make_ssa_jump_prob(; rng = StableRNG(12345)) +function make_ssa_jump_prob() j1 = ConstantRateJump((u, p, t) -> 10.0, integrator -> (integrator.u[1] += 1)) j2 = ConstantRateJump((u, p, t) -> 0.5 * u[1], integrator -> (integrator.u[1] -= 1)) dprob = DiscreteProblem([10], (0.0, 20.0)) - JumpProblem(dprob, Direct(), j1, j2; rng) + JumpProblem(dprob, Direct(), j1, j2) end # ODE + variable-rate jump -function make_vr_jump_prob(agg; rng = StableRNG(12345)) +function make_vr_jump_prob(agg) f!(du, u, p, t) = (du[1] = -0.1 * u[1]; nothing) oprob = ODEProblem(f!, [100.0], (0.0, 10.0)) vrj = VariableRateJump((u, p, t) -> 0.5 * u[1], integrator -> (integrator.u[1] -= 1.0)) - JumpProblem(oprob, Direct(), vrj; vr_aggregator = agg, rng) + JumpProblem(oprob, Direct(), vrj; vr_aggregator = agg) end # SDE + variable-rate jump -function make_sde_vr_jump_prob(agg; rng = StableRNG(12345)) +function make_sde_vr_jump_prob(agg) f!(du, u, p, t) = (du[1] = -0.1 * u[1]; nothing) g!(du, u, p, t) = (du[1] = 0.1 * u[1]; nothing) sprob = SDEProblem(f!, g!, [100.0], (0.0, 10.0)) vrj = VariableRateJump((u, p, t) -> 0.5 * u[1], integrator -> (integrator.u[1] -= 1.0)) - JumpProblem(sprob, Direct(), vrj; vr_aggregator = agg, rng) + JumpProblem(sprob, Direct(), vrj; vr_aggregator = agg) end # Helpers @@ -43,7 +43,7 @@ first_jump_time(traj) = traj.t[2] @testset "SSAStepper" begin jprob = make_ssa_jump_prob() sol = solve(EnsembleProblem(jprob), SSAStepper(), EnsembleSerial(); - trajectories = 3) + trajectories = 3, rng = StableRNG(12345)) times = [first_jump_time(sol.u[i]) for i in 1:3] @test allunique(times) end @@ -51,7 +51,7 @@ first_jump_time(traj) = traj.t[2] @testset "ODE + VR ($agg)" for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) jprob = make_vr_jump_prob(agg) sol = solve(EnsembleProblem(jprob), Tsit5(), EnsembleSerial(); - trajectories = 3) + trajectories = 3, rng = StableRNG(12345)) times = [first_jump_time(sol.u[i]) for i in 1:3] @test allunique(times) finals = [sol.u[i].u[end][1] for i in 1:3] @@ -76,13 +76,15 @@ end @testset "Sequential solves: different RNG streams" begin @testset "SSAStepper" begin jprob = make_ssa_jump_prob() - times = [first_jump_time(solve(jprob, SSAStepper())) for _ in 1:3] + rng = StableRNG(12345) + times = [first_jump_time(solve(jprob, SSAStepper(); rng)) for _ in 1:3] @test allunique(times) end @testset "ODE + VR ($agg)" for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) jprob = make_vr_jump_prob(agg) - sols = [solve(jprob, Tsit5()) for _ in 1:3] + rng = StableRNG(12345) + sols = [solve(jprob, Tsit5(); rng) for _ in 1:3] times = [first_jump_time(s) for s in sols] @test allunique(times) finals = [s.u[end][1] for s in sols] @@ -93,13 +95,9 @@ end # ========================================================================== # 3. Threaded ensemble: no data race on the shared JumpProblem # -# The ODE/SSA path through __jump_init receives seed=nothing from -# SciMLBase, so deepcopy'd problems on non-main threads start with -# identical RNG states. We only assert completion here — uniqueness -# requires explicit seeding (tested in section 4 below). -# -# The SDE path goes through StochasticDiffEq's __init which generates -# per-trajectory seeds, so we can additionally verify uniqueness there. +# With integrator-owned RNGs, each thread's integrator gets its own +# default_rng(). We only assert completion here — uniqueness is tested +# via explicit rng kwarg in section 4. # ========================================================================== @testset "EnsembleThreads: no data race" begin @@ -112,8 +110,6 @@ end @testset "ODE + VR ($agg)" for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) jprob = make_vr_jump_prob(agg) - # This path previously had a data race: resetted_jump_problem called - # randexp!(_jump_prob.rng, ...) on the shared original problem. sol = solve(EnsembleProblem(jprob), Tsit5(), EnsembleThreads(); trajectories = 4, save_everystep = false) @test length(sol) == 4 @@ -121,8 +117,8 @@ end @testset "SDE + VR (VR_FRM): unique trajectories" begin jprob = make_sde_vr_jump_prob(VR_FRM()) - # StochasticDiffEq generates per-trajectory seeds and passes them to - # resetted_jump_problem, so trajectories should be distinct. + # StochasticDiffEq generates per-trajectory seeds, so trajectories + # should be distinct. sol = solve(EnsembleProblem(jprob), EM(), EnsembleThreads(); trajectories = 4, dt = 0.01, save_everystep = false) @test length(sol) == 4 @@ -132,57 +128,63 @@ end end # ========================================================================== -# 4. Seed-based stream independence: resetted_jump_problem and -# reset_jump_problem! produce distinct RNG streams for different seeds -# -# This tests the mechanism that EnsembleThreads relies on (when seeds are -# provided by the caller, e.g. StochasticDiffEq) to get independent streams -# on different threads. +# 4. rng kwarg reproducibility: same rng seed → identical trajectory, +# different rng seeds → different trajectories # ========================================================================== -@testset "resetted_jump_problem: different seeds → different streams" begin - jprob = make_ssa_jump_prob() - seeds = UInt64[100, 200, 300] +@testset "rng kwarg reproducibility" begin + @testset "SSAStepper: same seed → same trajectory" begin + jprob = make_ssa_jump_prob() + sol1 = solve(jprob, SSAStepper(); rng = StableRNG(42)) + sol2 = solve(jprob, SSAStepper(); rng = StableRNG(42)) + @test sol1.t == sol2.t + @test sol1.u == sol2.u + end - # Each seed should produce a distinct aggregator RNG state - rngs = map(seeds) do s - jp = JumpProcesses.resetted_jump_problem(jprob, s) - jp.jump_callback.discrete_callbacks[1].condition.rng + @testset "SSAStepper: different seeds → different trajectories" begin + jprob = make_ssa_jump_prob() + sol1 = solve(jprob, SSAStepper(); rng = StableRNG(100)) + sol2 = solve(jprob, SSAStepper(); rng = StableRNG(200)) + sol3 = solve(jprob, SSAStepper(); rng = StableRNG(300)) + times = [first_jump_time(sol1), first_jump_time(sol2), first_jump_time(sol3)] + @test allunique(times) end - draws = [rand(rng) for rng in rngs] - @test allunique(draws) - - # Same seed should be deterministic - jp1 = JumpProcesses.resetted_jump_problem(jprob, UInt64(42)) - jp2 = JumpProcesses.resetted_jump_problem(jprob, UInt64(42)) - rng1 = jp1.jump_callback.discrete_callbacks[1].condition.rng - rng2 = jp2.jump_callback.discrete_callbacks[1].condition.rng - @test rand(rng1) == rand(rng2) -end -@testset "reset_jump_problem!: different seeds → different streams" begin - seeds = UInt64[100, 200, 300] - draws = map(seeds) do s - jp = make_ssa_jump_prob() - JumpProcesses.reset_jump_problem!(jp, s) - rand(jp.jump_callback.discrete_callbacks[1].condition.rng) + @testset "ODE + VR ($agg): same seed → same trajectory" for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) + jprob = make_vr_jump_prob(agg) + sol1 = solve(jprob, Tsit5(); rng = StableRNG(42)) + sol2 = solve(jprob, Tsit5(); rng = StableRNG(42)) + @test sol1.t ≈ sol2.t + @test sol1.u[end] ≈ sol2.u[end] + end + + @testset "ODE + VR ($agg): different seeds → different trajectories" for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) + jprob = make_vr_jump_prob(agg) + sols = [solve(jprob, Tsit5(); rng = StableRNG(s)) for s in (100, 200, 300)] + finals = [s.u[end][1] for s in sols] + @test allunique(finals) end - @test allunique(draws) end -@testset "_derive_jump_seed: decorrelates from input seed" begin - seed = UInt64(12345) - derived = JumpProcesses._derive_jump_seed(seed) - # Derived seed should differ from input - @test derived != seed - # Should be deterministic - @test derived == JumpProcesses._derive_jump_seed(seed) - # Different inputs → different outputs - @test JumpProcesses._derive_jump_seed(UInt64(1)) != JumpProcesses._derive_jump_seed(UInt64(2)) +# ========================================================================== +# 5. has_rng / get_rng / set_rng! interface on SSAIntegrator +# ========================================================================== + +@testset "SSAIntegrator RNG interface" begin + jprob = make_ssa_jump_prob() + integrator = init(jprob, SSAStepper(); rng = StableRNG(42)) + + @test SciMLBase.has_rng(integrator) + rng = SciMLBase.get_rng(integrator) + @test rng isa StableRNG + + new_rng = StableRNG(99) + SciMLBase.set_rng!(integrator, new_rng) + @test SciMLBase.get_rng(integrator) === new_rng end # ========================================================================== -# 5. Variable-rate: jump_u thresholds are unique per trajectory +# 6. Variable-rate: jump_u thresholds are unique per trajectory # # For VR_FRM, each trajectory's first jump time is determined by the initial # jump_u threshold (set to -randexp() by the VR_FRMEventCallback initialize). @@ -193,9 +195,23 @@ end @testset "VR_FRM: jump_u thresholds unique per trajectory (EnsembleSerial)" begin jprob = make_vr_jump_prob(VR_FRM()) sol = solve(EnsembleProblem(jprob), Tsit5(), EnsembleSerial(); - trajectories = 3) + trajectories = 3, rng = StableRNG(12345)) # The second time point is when the first variable-rate jump fires, # directly reflecting the initial -randexp() threshold. event_times = [sol.u[i].t[2] for i in 1:3] @test allunique(event_times) end + +# ========================================================================== +# 7. JumpProblem rng kwarg forwarded to solver +# ========================================================================== + +@testset "JumpProblem rng kwarg forwarded to solver" begin + j1 = ConstantRateJump((u, p, t) -> 1.0, integrator -> (integrator.u[1] += 1)) + dprob = DiscreteProblem([10], (0.0, 10.0)) + jprob = JumpProblem(dprob, Direct(), j1; rng = StableRNG(1)) + @test haskey(jprob.kwargs, :rng) + @test jprob.kwargs[:rng] isa StableRNG + sol = solve(jprob, SSAStepper()) + @test sol.retcode == ReturnCode.Success +end diff --git a/test/extended_jump_array_remake.jl b/test/extended_jump_array_remake.jl index 7f2168d5e..224064ae1 100644 --- a/test/extended_jump_array_remake.jl +++ b/test/extended_jump_array_remake.jl @@ -17,20 +17,21 @@ using StableRNGs vr_affect!(integrator) = (integrator.u[1] -= 1; integrator.u[2] += 1) vrj = VariableRateJump(vr_rate, vr_affect!) - jprob = JumpProblem(oprob, vrj; rng) + jprob = JumpProblem(oprob, vrj) # Verify we have ExtendedJumpArray @test jprob.prob.u0 isa ExtendedJumpArray @test jprob.prob.u0.u == [10.0, 5.0] @testset "remake with numeric Vector{Float64}" begin - original_jump_u = copy(jprob.prob.u0.jump_u) - prob2 = remake(jprob; u0 = [20.0, 10.0]) @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u == [20.0, 10.0] - # jump_u should be resampled (different from original) - @test prob2.prob.u0.jump_u != original_jump_u + # jump_u is zeroed at construction; callback initializes it at solve time + @test all(iszero, prob2.prob.u0.jump_u) + # After init the callback should set jump_u to non-zero thresholds + integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + @test any(!iszero, integrator.u.jump_u) end @testset "remake with ExtendedJumpArray (no resample)" begin @@ -47,34 +48,34 @@ using StableRNGs end @testset "remake with Symbol pairs" begin - original_jump_u = copy(jprob.prob.u0.jump_u) - # This was the FAILING case - should work after fix prob2 = remake(jprob; u0 = [:X => 25.0]) @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u[1] == 25.0 - # jump_u should be resampled - @test prob2.prob.u0.jump_u != original_jump_u + # jump_u zeroed at construction, initialized by callback at solve time + @test all(iszero, prob2.prob.u0.jump_u) + integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + @test any(!iszero, integrator.u.jump_u) end @testset "remake with multiple Symbol pairs" begin - original_jump_u = copy(jprob.prob.u0.jump_u) - prob2 = remake(jprob; u0 = [:X => 35.0, :Y => 15.0]) @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u == [35.0, 15.0] - # jump_u should be resampled - @test prob2.prob.u0.jump_u != original_jump_u + # jump_u zeroed at construction, initialized by callback at solve time + @test all(iszero, prob2.prob.u0.jump_u) + integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + @test any(!iszero, integrator.u.jump_u) end @testset "remake with Dict" begin - original_jump_u = copy(jprob.prob.u0.jump_u) - prob2 = remake(jprob; u0 = Dict(:X => 40.0)) @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u[1] == 40.0 - # jump_u should be resampled - @test prob2.prob.u0.jump_u != original_jump_u + # jump_u zeroed at construction, initialized by callback at solve time + @test all(iszero, prob2.prob.u0.jump_u) + integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + @test any(!iszero, integrator.u.jump_u) end @testset "remake with parameters only (u0 unchanged)" begin @@ -88,14 +89,14 @@ using StableRNGs end @testset "remake with both u0 and p" begin - original_jump_u = copy(jprob.prob.u0.jump_u) - prob2 = remake(jprob; u0 = [:X => 50.0], p = [:k1 => 3.0]) @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u[1] == 50.0 @test prob2.prob.p[1] == 3.0 - # jump_u should be resampled - @test prob2.prob.u0.jump_u != original_jump_u + # jump_u zeroed at construction, initialized by callback at solve time + @test all(iszero, prob2.prob.u0.jump_u) + integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + @test any(!iszero, integrator.u.jump_u) end @testset "remake preserves problem solvability" begin diff --git a/test/remake_test.jl b/test/remake_test.jl index 91ab4daba..90c863d3c 100644 --- a/test/remake_test.jl +++ b/test/remake_test.jl @@ -1,4 +1,4 @@ -using JumpProcesses, DiffEqBase, OrdinaryDiffEq +using JumpProcesses, DiffEqBase, OrdinaryDiffEq, Test using StableRNGs rng = StableRNG(12345) @@ -21,21 +21,20 @@ p = (0.1 / 1000, 0.01) tspan = (0.0, 2500.0) dprob = DiscreteProblem(u0, tspan, p) -jprob = JumpProblem(dprob, Direct(), jump, jump2, save_positions = (false, false), - rng = rng) -sol = solve(jprob, SSAStepper()) +jprob = JumpProblem(dprob, Direct(), jump, jump2, save_positions = (false, false)) +sol = solve(jprob, SSAStepper(); rng) @test sol[3, end] == 1000 u02 = [1000, 1, 0] p2 = (0.1 / 1000, 0.0) dprob2 = remake(dprob, u0 = u02, p = p2) jprob2 = remake(jprob, prob = dprob2) -sol2 = solve(jprob2, SSAStepper()) +sol2 = solve(jprob2, SSAStepper(); rng) @test sol2[2, end] == 1001 tspan2 = (0.0, 25000.0) jprob3 = remake(jprob, p = p2, tspan = tspan2) -sol3 = solve(jprob3, SSAStepper()) +sol3 = solve(jprob3, SSAStepper(); rng) @test sol3[2, end] == 1000 @test sol3.t[end] == 25000.0 @@ -46,19 +45,19 @@ ns = [[2 => -1, 3 => 1], [1 => -1, 2 => 1]] pidxs = [2, 1] maj = MassActionJump(rs, ns; param_idxs = pidxs) dprob = DiscreteProblem(u0, tspan, p) -jprob = JumpProblem(dprob, Direct(), maj, save_positions = (false, false), rng = rng) -sol = solve(jprob, SSAStepper()) +jprob = JumpProblem(dprob, Direct(), maj, save_positions = (false, false)) +sol = solve(jprob, SSAStepper(); rng) @test sol[3, end] == 1000 # update the MassActionJump dprob2 = remake(dprob, u0 = u02, p = p2) jprob2 = remake(jprob, prob = dprob2) -sol2 = solve(jprob2, SSAStepper()) +sol2 = solve(jprob2, SSAStepper(); rng) @test sol2[2, end] == 1001 tspan2 = (0.0, 25000.0) jprob3 = remake(jprob, p = p2, tspan = tspan2) -sol3 = solve(jprob3, SSAStepper()) +sol3 = solve(jprob3, SSAStepper(); rng) @test sol3[2, end] == 1000 @test sol3.t[end] == 25000.0 @@ -75,27 +74,27 @@ let rrate(u, p, t) = u[1] aaffect!(integrator) = (integrator.u[1] += 1; nothing) vrj = VariableRateJump(rrate, aaffect!) - jprob = JumpProblem(prob, vrj; vr_aggregator = VR_Direct(), rng) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, vrj; vr_aggregator = VR_Direct()) + sol = solve(jprob, Tsit5(); rng) @test all(==(0.0), sol[1, :]) - jprob = JumpProblem(prob, vrj; vr_aggregator = VR_DirectFW(), rng) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, vrj; vr_aggregator = VR_DirectFW()) + sol = solve(jprob, Tsit5(); rng) @test all(==(0.0), sol[1, :]) - jprob = JumpProblem(prob, vrj; vr_aggregator = VR_FRM(), rng) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, vrj; vr_aggregator = VR_FRM()) + sol = solve(jprob, Tsit5(); rng) @test all(==(0.0), sol[1, :]) u0 = [4.0] jprob2 = remake(jprob; u0) @test jprob2.prob.u0 isa ExtendedJumpArray @test jprob2.prob.u0.u === u0 - sol = solve(jprob2, Tsit5()) + sol = solve(jprob2, Tsit5(); rng) u = sol[1, :] @test length(u) > 2 @test all(>(u0[1]), u[3:end]) u0 = deepcopy(jprob2.prob.u0) u0.u .= 0 jprob3 = remake(jprob2; u0) - sol = solve(jprob3, Tsit5()) + sol = solve(jprob3, Tsit5(); rng) @test all(==(0.0), sol[1, :]) @test_throws ErrorException jprob4=remake(jprob, u0 = 1) end @@ -107,24 +106,24 @@ let rrate(u, p, t) = u[1] aaffect!(integrator) = (integrator.u[1] += 1; nothing) vrj = VariableRateJump(rrate, aaffect!) - jprob = JumpProblem(prob, vrj; vr_aggregator = VR_Direct(), rng) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, vrj; vr_aggregator = VR_Direct()) + sol = solve(jprob, Tsit5(); rng) @test all(==(0.0), sol[1, :]) - jprob = JumpProblem(prob, vrj; vr_aggregator = VR_DirectFW(), rng) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, vrj; vr_aggregator = VR_DirectFW()) + sol = solve(jprob, Tsit5(); rng) @test all(==(0.0), sol[1, :]) - jprob = JumpProblem(prob, vrj; vr_aggregator = VR_FRM(), rng) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, vrj; vr_aggregator = VR_FRM()) + sol = solve(jprob, Tsit5(); rng) @test all(==(0.0), sol[1, :]) u0 = [4.0] prob2 = remake(jprob.prob; u0) @test_throws ErrorException jprob2=remake(jprob; prob = prob2) - u0eja = JumpProcesses.remake_extended_u0(jprob.prob, u0, rng) + u0eja = JumpProcesses.remake_extended_u0(jprob.prob, u0) prob3 = remake(jprob.prob; u0 = u0eja) jprob3 = remake(jprob; prob = prob3) @test jprob3.prob.u0 isa ExtendedJumpArray @test jprob3.prob.u0 === u0eja - sol = solve(jprob3, Tsit5()) + sol = solve(jprob3, Tsit5(); rng) u = sol[1, :] @test length(u) > 2 @test all(>(u0[1]), u[3:end]) diff --git a/test/rng_kwarg_tests.jl b/test/rng_kwarg_tests.jl new file mode 100644 index 000000000..a619edb76 --- /dev/null +++ b/test/rng_kwarg_tests.jl @@ -0,0 +1,219 @@ +using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq, Test +using StableRNGs, Random + +# ========================================================================== +# Test that rng can be passed via JumpProblem kwargs OR solve kwargs, +# and that solve-level rng takes precedence over JumpProblem-level rng. +# +# Strategy: use different RNG *types* to verify which one the integrator +# receives. StableRNG is passed at one level and Xoshiro at another, +# then we check the type on the integrator. +# ========================================================================== + +# -------------------------------------------------------------------------- +# Problem constructors +# -------------------------------------------------------------------------- +function make_ssa_jump_prob(; kwargs...) + j1 = ConstantRateJump((u, p, t) -> 10.0, integrator -> (integrator.u[1] += 1)) + j2 = ConstantRateJump((u, p, t) -> 0.5 * u[1], integrator -> (integrator.u[1] -= 1)) + dprob = DiscreteProblem([10], (0.0, 20.0)) + JumpProblem(dprob, Direct(), j1, j2; kwargs...) +end + +function make_ode_vr_jump_prob(; kwargs...) + f!(du, u, p, t) = (du[1] = -0.1 * u[1]; nothing) + oprob = ODEProblem(f!, [100.0], (0.0, 10.0)) + vrj = VariableRateJump((u, p, t) -> 0.5 * u[1], + integrator -> (integrator.u[1] -= 1.0)) + JumpProblem(oprob, Direct(), vrj; kwargs...) +end + +function make_sde_vr_jump_prob(; kwargs...) + f!(du, u, p, t) = (du[1] = -0.1 * u[1]; nothing) + g!(du, u, p, t) = (du[1] = 0.1 * u[1]; nothing) + sprob = SDEProblem(f!, g!, [100.0], (0.0, 10.0)) + vrj = VariableRateJump((u, p, t) -> 0.5 * u[1], + integrator -> (integrator.u[1] -= 1.0)) + JumpProblem(sprob, Direct(), vrj; kwargs...) +end + +# ========================================================================== +# 1. SSAStepper: rng via JumpProblem +# ========================================================================== +@testset "SSAStepper: rng via JumpProblem kwargs" begin + xrng = Xoshiro(42) + jprob = make_ssa_jump_prob(; rng = xrng) + integrator = init(jprob, SSAStepper()) + @test SciMLBase.get_rng(integrator) isa Xoshiro + sol = solve(jprob, SSAStepper()) + @test sol.retcode == ReturnCode.Success +end + +# ========================================================================== +# 2. SSAStepper: rng via solve overrides JumpProblem +# ========================================================================== +@testset "SSAStepper: solve rng overrides JumpProblem rng" begin + jprob = make_ssa_jump_prob(; rng = Xoshiro(42)) + integrator = init(jprob, SSAStepper(); rng = StableRNG(99)) + @test SciMLBase.get_rng(integrator) isa StableRNG +end + +# ========================================================================== +# 3. SSAStepper: reproducibility via JumpProblem rng +# ========================================================================== +@testset "SSAStepper: JumpProblem rng reproducibility" begin + jprob1 = make_ssa_jump_prob(; rng = StableRNG(123)) + jprob2 = make_ssa_jump_prob(; rng = StableRNG(123)) + sol1 = solve(jprob1, SSAStepper()) + sol2 = solve(jprob2, SSAStepper()) + @test sol1.t == sol2.t + @test sol1.u == sol2.u +end + +# ========================================================================== +# 4. SSAStepper: solve rng overrides for reproducibility +# ========================================================================== +@testset "SSAStepper: solve rng override reproducibility" begin + jprob = make_ssa_jump_prob(; rng = Xoshiro(1)) + sol1 = solve(jprob, SSAStepper(); rng = StableRNG(42)) + sol2 = solve(jprob, SSAStepper(); rng = StableRNG(42)) + @test sol1.t == sol2.t + @test sol1.u == sol2.u +end + +# ========================================================================== +# 5. ODE + VR: rng via JumpProblem +# ========================================================================== +@testset "ODE + VR: rng via JumpProblem kwargs" begin + jprob = make_ode_vr_jump_prob(; rng = Xoshiro(42)) + integrator = init(jprob, Tsit5()) + @test SciMLBase.get_rng(integrator) isa Xoshiro +end + +# ========================================================================== +# 6. ODE + VR: solve rng overrides JumpProblem rng +# ========================================================================== +@testset "ODE + VR: solve rng overrides JumpProblem rng" begin + jprob = make_ode_vr_jump_prob(; rng = Xoshiro(42)) + integrator = init(jprob, Tsit5(); rng = StableRNG(99)) + @test SciMLBase.get_rng(integrator) isa StableRNG +end + +# ========================================================================== +# 7. ODE + VR: reproducibility via JumpProblem rng +# ========================================================================== +@testset "ODE + VR: JumpProblem rng reproducibility" begin + jprob1 = make_ode_vr_jump_prob(; rng = StableRNG(123)) + jprob2 = make_ode_vr_jump_prob(; rng = StableRNG(123)) + sol1 = solve(jprob1, Tsit5()) + sol2 = solve(jprob2, Tsit5()) + @test sol1.t ≈ sol2.t + @test sol1.u[end] ≈ sol2.u[end] +end + +# ========================================================================== +# 8. ODE + VR: solve rng overrides for reproducibility +# ========================================================================== +@testset "ODE + VR: solve rng override reproducibility" begin + jprob = make_ode_vr_jump_prob(; rng = Xoshiro(1)) + sol1 = solve(jprob, Tsit5(); rng = StableRNG(42)) + sol2 = solve(jprob, Tsit5(); rng = StableRNG(42)) + @test sol1.t ≈ sol2.t + @test sol1.u[end] ≈ sol2.u[end] +end + +# ========================================================================== +# 9. SDE + VR: rng via JumpProblem +# ========================================================================== +@testset "SDE + VR: rng via JumpProblem kwargs" begin + jprob = make_sde_vr_jump_prob(; rng = Xoshiro(42)) + integrator = init(jprob, EM(); dt = 0.01) + @test SciMLBase.get_rng(integrator) isa Xoshiro +end + +# ========================================================================== +# 10. SDE + VR: solve rng overrides JumpProblem rng +# ========================================================================== +@testset "SDE + VR: solve rng overrides JumpProblem rng" begin + jprob = make_sde_vr_jump_prob(; rng = Xoshiro(42)) + integrator = init(jprob, EM(); dt = 0.01, rng = StableRNG(99)) + @test SciMLBase.get_rng(integrator) isa StableRNG +end + +# ========================================================================== +# 11. SDE + VR: reproducibility via JumpProblem rng +# ========================================================================== +@testset "SDE + VR: JumpProblem rng reproducibility" begin + jprob1 = make_sde_vr_jump_prob(; rng = StableRNG(123)) + jprob2 = make_sde_vr_jump_prob(; rng = StableRNG(123)) + sol1 = solve(jprob1, EM(); dt = 0.01, save_everystep = false) + sol2 = solve(jprob2, EM(); dt = 0.01, save_everystep = false) + @test sol1.u[end] ≈ sol2.u[end] +end + +# ========================================================================== +# 12. Tau-leaping: rng via JumpProblem +# ========================================================================== +@testset "SimpleTauLeaping: rng via JumpProblem kwargs" begin + rate(out, u, p, t) = (out .= max.(u, 0); nothing) + c(du, u, p, t, counts, mark) = (du .= counts; nothing) + rj = RegularJump(rate, c, 2) + dprob = DiscreteProblem([100, 100], (0.0, 1.0)) + jprob = JumpProblem(dprob, PureLeaping(), rj; rng = StableRNG(42)) + sol1 = solve(jprob, SimpleTauLeaping(); dt = 0.01) + jprob2 = JumpProblem(dprob, PureLeaping(), rj; rng = StableRNG(42)) + sol2 = solve(jprob2, SimpleTauLeaping(); dt = 0.01) + @test sol1.u == sol2.u +end + +# ========================================================================== +# 13. Tau-leaping: solve rng overrides JumpProblem rng +# ========================================================================== +@testset "SimpleTauLeaping: solve rng overrides JumpProblem rng" begin + rate(out, u, p, t) = (out .= max.(u, 0); nothing) + c(du, u, p, t, counts, mark) = (du .= counts; nothing) + rj = RegularJump(rate, c, 2) + dprob = DiscreteProblem([100, 100], (0.0, 1.0)) + # JumpProblem has Xoshiro, solve has StableRNG + jprob = JumpProblem(dprob, PureLeaping(), rj; rng = Xoshiro(1)) + sol1 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(42)) + sol2 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(42)) + @test sol1.u == sol2.u + # Different from using the JumpProblem rng + jprob2 = JumpProblem(dprob, PureLeaping(), rj; rng = Xoshiro(1)) + sol3 = solve(jprob2, SimpleTauLeaping(); dt = 0.01) + @test sol1.u != sol3.u +end + +# ========================================================================== +# 14. has_rng / get_rng / set_rng! interface on SSAIntegrator +# ========================================================================== +@testset "SSAIntegrator RNG interface" begin + jprob = make_ssa_jump_prob() + integrator = init(jprob, SSAStepper(); rng = StableRNG(42)) + + @test SciMLBase.has_rng(integrator) + rng = SciMLBase.get_rng(integrator) + @test rng isa StableRNG + + new_rng = StableRNG(99) + SciMLBase.set_rng!(integrator, new_rng) + @test SciMLBase.get_rng(integrator) === new_rng +end + +# ========================================================================== +# 15. No rng kwarg: uses default_rng (non-reproducible but functional) +# ========================================================================== +@testset "No rng kwarg: functional solve" begin + @testset "SSAStepper" begin + jprob = make_ssa_jump_prob() + sol = solve(jprob, SSAStepper()) + @test sol.retcode == ReturnCode.Success + end + + @testset "ODE + VR" begin + jprob = make_ode_vr_jump_prob() + sol = solve(jprob, Tsit5()) + @test sol.retcode == ReturnCode.Success + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 389ed1dab..be0236c3b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -39,6 +39,7 @@ end if GROUP == "All" || GROUP == "InterfaceII" @time @safetestset "Saveat Regression test" begin include("saveat_regression.jl") end @time @safetestset "Save_positions test" begin include("save_positions.jl") end + @time @safetestset "RNG kwarg tests" begin include("rng_kwarg_tests.jl") end @time @safetestset "Ensemble Uniqueness test" begin include("ensemble_uniqueness.jl") end @time @safetestset "Thread Safety test" begin include("thread_safety.jl") end @time @safetestset "Ensemble Problem Tests" begin include("ensemble_problems.jl") end @@ -47,7 +48,7 @@ end @time @safetestset "ExtendedJumpArray remake tests" begin include("extended_jump_array_remake.jl") end @time @safetestset "Symbol based problem indexing" begin include("jprob_symbol_indexing.jl") end @time @safetestset "Long time accuracy test" begin include("longtimes_test.jl") end - @time @safetestset "Hawkes process" begin include("hawkes_test.jl") end + @time @safetestset "Hawkes process" begin include("hawkes_test.jl") end @time @safetestset "Reaction rates" begin include("spatial/reaction_rates.jl") end @time @safetestset "Hop rates" begin include("spatial/hop_rates.jl") end @time @safetestset "Topology" begin include("spatial/topology.jl") end diff --git a/test/ssa_callback_test.jl b/test/ssa_callback_test.jl index dbe9c969d..1cbf678b3 100644 --- a/test/ssa_callback_test.jl +++ b/test/ssa_callback_test.jl @@ -11,9 +11,9 @@ end jump = ConstantRateJump(rate, affect!) prob = DiscreteProblem([0.0, 0.0], (0.0, 10.0)) -jump_prob = JumpProblem(prob, Direct(), jump; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump) -sol = solve(jump_prob, SSAStepper()) +sol = solve(jump_prob, SSAStepper(); rng) @test sol.t == [0.0, 10.0] @test sol.u == [[0.0, 0.0], [0.0, 0.0]] @@ -25,13 +25,13 @@ function fuel_affect!(integrator) end cb = DiscreteCallback(condition, fuel_affect!, save_positions = (false, true)) -sol = solve(jump_prob, SSAStepper(); callback = cb, tstops = [5]) +sol = solve(jump_prob, SSAStepper(); rng, callback = cb, tstops = [5]) @test sol.t[1:2] == [0.0, 5.0] # no jumps between t=0 and t=5 @test sol(5 + 1e-10) == [100, 0] # state just after fueling before any decays can happen # test can pass callbacks via JumpProblem -jump_prob2 = JumpProblem(prob, Direct(), jump; rng = rng, callback = cb) -sol2 = solve(jump_prob2, SSAStepper(); tstops = [5]) +jump_prob2 = JumpProblem(prob, Direct(), jump; callback = cb) +sol2 = solve(jump_prob2, SSAStepper(); rng, tstops = [5]) @test sol2.t[1:2] == [0.0, 5.0] # no jumps between t=0 and t=5 @test sol2(5 + 1e-10) == [100, 0] # state just after fueling before any decays can happen @@ -48,7 +48,7 @@ finalizer_called = 0 fuel_finalize(cb, u, t, integrator) = global finalizer_called += 1 cb2 = DiscreteCallback(condition, fuel_affect!, initialize = fuel_init!, finalize = fuel_finalize) -sol = solve(jump_prob, SSAStepper(), callback = cb2) +sol = solve(jump_prob, SSAStepper(); rng, callback = cb2) for tstop in random_tstops @test tstop ∈ sol.t end @@ -62,37 +62,37 @@ maj = MassActionJump(rs, ns; param_idxs = [1, 2]) u₀ = [100, 0] tspan = (0.0, 2000.0) dprob = DiscreteProblem(u₀, tspan, p) -jprob = JumpProblem(dprob, Direct(), maj, save_positions = (false, false), rng = rng) +jprob = JumpProblem(dprob, Direct(), maj, save_positions = (false, false)) pcondit(u, t, integrator) = t == 1000.0 function paffect!(integrator) integrator.p[1] = 0.0 integrator.p[2] = 1.0 reset_aggregated_jumps!(integrator) end -sol = solve(jprob, SSAStepper(), tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) @test all(p .== [0.0, 1.0]) @test sol[1, end] == 100 p .= [1.0, 0.0] maj1 = MassActionJump([1 => 1], [1 => -1, 2 => 1]; param_idxs = 1) maj2 = MassActionJump([2 => 1], [1 => 1, 2 => -1]; param_idxs = 2) -jprob = JumpProblem(dprob, Direct(), maj1, maj2, save_positions = (false, false), rng = rng) -sol = solve(jprob, SSAStepper(), tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) +jprob = JumpProblem(dprob, Direct(), maj1, maj2, save_positions = (false, false)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) @test all(p .== [0.0, 1.0]) @test sol[1, end] == 100 p2 = [1.0, 0.0, 0.0] maj3 = MassActionJump([1 => 1], [1 => -1, 2 => 1]; param_idxs = 3) dprob = DiscreteProblem(u₀, tspan, p2) -jprob = JumpProblem(dprob, Direct(), maj1, maj2, maj3, save_positions = (false, false), rng = rng) -sol = solve(jprob, SSAStepper(), tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) +jprob = JumpProblem(dprob, Direct(), maj1, maj2, maj3, save_positions = (false, false)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) @test all(p2 .== [0.0, 1.0, 0.0]) @test sol[1, end] == 100 p2 .= [1.0, 0.0, 0.0] jprob = JumpProblem(dprob, Direct(), JumpSet(; massaction_jumps = [maj1, maj2, maj3]), - save_positions = (false, false), rng = rng) -sol = solve(jprob, SSAStepper(), tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) + save_positions = (false, false)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) @test all(p2 .== [0.0, 1.0, 0.0]) @test sol[1, end] == 100 @@ -100,8 +100,8 @@ p .= [1.0, 0.0] dprob = DiscreteProblem(u₀, tspan, p) maj4 = MassActionJump([[1 => 1], [2 => 1]], [[1 => -1, 2 => 1], [1 => 1, 2 => -1]]; param_idxs = [1, 2]) -jprob = JumpProblem(dprob, Direct(), maj4, save_positions = (false, false), rng = rng) -sol = solve(jprob, SSAStepper(), tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) +jprob = JumpProblem(dprob, Direct(), maj4, save_positions = (false, false)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) @test all(p .== [0.0, 1.0]) @test sol[1, end] == 100 @@ -109,9 +109,9 @@ sol = solve(jprob, SSAStepper(), tstops = [1000.0], callback = DiscreteCallback( p .= [1.0] dprob = DiscreteProblem(u₀, tspan, p) maj5 = MassActionJump([[1 => 2]], [[1 => -1, 2 => 1]]; param_idxs = [1]) -jprob = JumpProblem(dprob, Direct(), maj5, save_positions = (false, false), rng = rng) +jprob = JumpProblem(dprob, Direct(), maj5, save_positions = (false, false)) @test all(jprob.massaction_jump.scaled_rates .== [0.5]) -jprob = JumpProblem(dprob, Direct(), maj5, save_positions = (false, false), rng = rng, scale_rates = false) +jprob = JumpProblem(dprob, Direct(), maj5, save_positions = (false, false), scale_rates = false) @test all(jprob.massaction_jump.scaled_rates .== [1.0]) # test for https://github.com/SciML/JumpProcesses.jl/issues/239 @@ -119,11 +119,11 @@ maj6 = MassActionJump([[1 => 1], [2 => 1]], [[1 => -1, 2 => 1], [1 => 1, 2 => -1 param_idxs = [1, 2]) p = (0.1, 0.1) dprob = DiscreteProblem([10, 0], (0.0, 100.0), p) -jprob = JumpProblem(dprob, Direct(), maj6; save_positions = (false, false), rng = rng) +jprob = JumpProblem(dprob, Direct(), maj6; save_positions = (false, false)) cbtimes = [20.0, 30.0] affectpresets!(integrator) = integrator.u[1] += 10 cb = PresetTimeCallback(cbtimes, affectpresets!) -jsol = solve(jprob, SSAStepper(), saveat = 0.1, callback = cb) +jsol = solve(jprob, SSAStepper(); rng, saveat = 0.1, callback = cb) @test (jsol(20.00000000001) - jsol(19.9999999999))[1] == 10 # test periodic callbacks working, i.e. #417 @@ -134,18 +134,18 @@ let dprob = DiscreteProblem([0], (0.0, 10.0)) cbfun(integ) = (integ.u[1] += 1; nothing) cb = PeriodicCallback(cbfun, 1.0) - jprob = JumpProblem(dprob, crj; rng) - sol = solve(jprob; callback = cb) + jprob = JumpProblem(dprob, crj) + sol = solve(jprob; rng, callback = cb) @test sol[1, end] == 9 cb = PeriodicCallback(cbfun, 1.0; initial_affect = true) - jprob = JumpProblem(dprob, crj; rng) - sol = solve(jprob; callback = cb) + jprob = JumpProblem(dprob, crj) + sol = solve(jprob; rng, callback = cb) @test sol[1, end] == 10 cb = PeriodicCallback(cbfun, 1.0; initial_affect = true, final_affect = true) - jprob = JumpProblem(dprob, crj; rng) - sol = solve(jprob; callback = cb) + jprob = JumpProblem(dprob, crj) + sol = solve(jprob; rng, callback = cb) @test sol[1, end] == 11 end @@ -157,22 +157,22 @@ let dprob = DiscreteProblem([0], (0.0, 10.0)) cbfun(integ) = (integ.u[1] += 1; nothing) cb = PeriodicCallback(cbfun, 1.0) - jprob = JumpProblem(dprob, crj; rng) + jprob = JumpProblem(dprob, crj) tstops = Float64[] # tests for when aliasing system is in place - #sol = solve(jprob; callback = cb, tstops, alias_tstops = true) + #sol = solve(jprob; callback = cb, tstops, alias_tstops = true) # @test sol[1, end] == 9 - #@test tstops == 1.0:9.0 + #@test tstops == 1.0:9.0 # empty!(tstops) # sol = solve(jprob; callback = cb, tstops, alias_tstops = false) # @test sol[1, end] == 9 # @test isempty(tstops) - sol = solve(jprob; callback = cb, tstops) + sol = solve(jprob; rng, callback = cb, tstops) @test sol[1, end] == 9 @test isempty(tstops) empty!(tstops) - integ = init(jprob, SSAStepper(); callback = cb, tstops) + integ = init(jprob, SSAStepper(); rng, callback = cb, tstops) solve!(integ) @test integ.tstops !== tstops @test isempty(tstops) @@ -184,18 +184,18 @@ let affect!(integrator) = (integrator.u[1] += 1) crj = ConstantRateJump(rate, affect!) prob = DiscreteProblem([0], (0.0, 10.0), [10.0]) - jprob = JumpProblem(prob, Direct(), crj; rng) + jprob = JumpProblem(prob, Direct(), crj) # basic callable tstops my_tstops = (p, tspan) -> [3.0, 6.0] - sol = solve(jprob, SSAStepper(); tstops = my_tstops) + sol = solve(jprob, SSAStepper(); rng, tstops = my_tstops) @test sol.t[end] == 10.0 @test 3.0 ∈ sol.t @test 6.0 ∈ sol.t # parameter-dependent callable tstops param_tstops = (p, tspan) -> [p[1] / 5.0, p[1] / 2.0] - sol2 = solve(jprob, SSAStepper(); tstops = param_tstops) + sol2 = solve(jprob, SSAStepper(); rng, tstops = param_tstops) @test sol2.t[end] == 10.0 @test 2.0 ∈ sol2.t # 10.0 / 5.0 @test 5.0 ∈ sol2.t # 10.0 / 2.0 @@ -204,7 +204,7 @@ let condition(u, t, integrator) = t == 3.0 cb_affect!(integrator) = (integrator.u[1] += 1000) cb = DiscreteCallback(condition, cb_affect!) - sol3 = solve(jprob, SSAStepper(); tstops = my_tstops, callback = cb) + sol3 = solve(jprob, SSAStepper(); rng, tstops = my_tstops, callback = cb) @test sol3.t[end] == 10.0 @test 3.0 ∈ sol3.t # verify the callback fired: use findlast to get post-callback state at t=3.0 @@ -213,15 +213,15 @@ let # callable returning a tuple tuple_tstops = (p, tspan) -> (2.0, 7.0) - sol4 = solve(jprob, SSAStepper(); tstops = tuple_tstops) + sol4 = solve(jprob, SSAStepper(); rng, tstops = tuple_tstops) @test sol4.t[end] == 10.0 @test 2.0 ∈ sol4.t @test 7.0 ∈ sol4.t # callable tstops stored in JumpProblem via constructor kwarg - jprob2 = JumpProblem(prob, Direct(), crj; rng, tstops = my_tstops) + jprob2 = JumpProblem(prob, Direct(), crj; tstops = my_tstops) @test haskey(jprob2.kwargs, :tstops) - sol5 = solve(jprob2, SSAStepper()) + sol5 = solve(jprob2, SSAStepper(); rng) @test sol5.t[end] == 10.0 @test 3.0 ∈ sol5.t @test 6.0 ∈ sol5.t diff --git a/test/variable_rate.jl b/test/variable_rate.jl index ce08781db..2669b70e2 100644 --- a/test/variable_rate.jl +++ b/test/variable_rate.jl @@ -293,7 +293,6 @@ end let seed = 12345 - rng = StableRNG(seed) b = 2.0 d = 1.0 n0 = 1.0 @@ -325,7 +324,7 @@ let ode_prob = ODEProblem(ode_fxn, u0, tspan, p) tsave = range(tspan[1], tspan[2]; step = 0.1) for vr_aggregator in (VR_Direct(), VR_DirectFW(), VR_FRM()) - sjm_prob = JumpProblem(ode_prob, b_jump, d_jump; vr_aggregator, rng) + sjm_prob = JumpProblem(ode_prob, b_jump, d_jump; vr_aggregator) for alg in (Tsit5(), Rodas5P(linsolve = QRFactorization())) umean = getmean(Nsims, sjm_prob, alg, tsave, seed) From 8861d6e0df564ca72c6bcabfd495df75f801d995 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Mon, 23 Feb 2026 22:06:28 -0500 Subject: [PATCH 02/52] Revert version bump to 9.22.1 Co-Authored-By: Claude Opus 4.5 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 404b9b5ea..ef960cdf3 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "JumpProcesses" uuid = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5" -version = "9.23.0" +version = "9.22.1" authors = ["Chris Rackauckas "] [deps] From e37938b3ac741bbda0eace5f13cc0ead2d8d49d5 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 05:56:03 -0500 Subject: [PATCH 03/52] Remove rng kwarg from JumpProblem; pass rng to solve/init only Breaking: JumpProblem no longer accepts an `rng` keyword argument. Passing `rng` now throws an ArgumentError directing users to pass it to `solve` or `init` instead. This fixes inconsistent RNG priority ordering across solver pathways (SSAStepper, ODE/SDE, tau-leaping) and makes the integrator the single source of truth for RNG state. - Simplify resolve_rng to 2-arg form (remove fallback_rng) - Remove jump_prob.kwargs rng fallback from tau-leaping solvers - Update all tests to pass rng to solve/init instead of JumpProblem - Update docs and HISTORY.md with breaking change notes Co-Authored-By: Claude Opus 4.5 --- HISTORY.md | 22 +++- docs/src/api.md | 26 +---- docs/src/faq.md | 8 -- src/problem.jl | 20 ++-- src/simple_regular_solve.jl | 4 +- src/solve.jl | 8 +- test/allocations.jl | 32 +++--- test/bimolerx_test.jl | 22 ++-- test/bracketing.jl | 2 +- test/callbacks.jl | 94 ++++++++--------- test/constant_rate.jl | 20 ++-- test/ensemble_problems.jl | 11 +- test/ensemble_uniqueness.jl | 32 ++---- test/extinction_test.jl | 18 ++-- test/fp_unknowns.jl | 4 +- test/functionwrappers.jl | 10 +- test/geneexpr_test.jl | 53 ++++++---- test/gpu/regular_jumps.jl | 2 +- test/hawkes_test.jl | 14 +-- test/linearreaction_test.jl | 20 ++-- test/longtimes_test.jl | 4 +- test/monte_carlo_test.jl | 24 ++--- test/regular_jumps.jl | 54 +++++----- test/reversible_binding.jl | 9 +- test/rng_kwarg_tests.jl | 188 +++++++++++++++------------------ test/save_positions.jl | 18 ++-- test/saveat_regression.jl | 10 +- test/sir_model.jl | 8 +- test/spatial/ABC.jl | 14 +-- test/spatial/diffusion.jl | 32 +++--- test/spatial/spatial_majump.jl | 26 ++--- test/splitcoupled.jl | 52 ++++----- test/ssa_tests.jl | 17 ++- test/thread_safety.jl | 6 +- test/variable_rate.jl | 100 +++++++++--------- 35 files changed, 477 insertions(+), 507 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 1851dbd63..d010778af 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,6 +1,26 @@ # Breaking updates and feature summaries across releases -## JumpProcesses unreleased (master branch) +## 10.0 (Breaking) + + - **Breaking**: The `rng` keyword argument has been removed from + `JumpProblem`. Pass `rng` to `solve` or `init` instead: + ```julia + # Before (no longer works): + jprob = JumpProblem(dprob, Direct(), jump; rng = Xoshiro(1234)) + sol = solve(jprob, SSAStepper()) + + # After: + jprob = JumpProblem(dprob, Direct(), jump) + sol = solve(jprob, SSAStepper(); rng = Xoshiro(1234)) + ``` + - RNG state is now owned by the integrator, not the aggregator. This + eliminates data races when sharing a `JumpProblem` across threads and + ensures a single, consistent RNG priority across all solver pathways: + `rng` > `seed` > `Random.default_rng()`. + - `rng` and `seed` kwargs are fully supported on `solve`/`init` for all + solver pathways (SSAStepper, ODE, SDE, tau-leaping). + - `SSAIntegrator` now supports the `SciMLBase` RNG interface (`has_rng`, + `get_rng`, `set_rng!`). ## 9.14 diff --git a/docs/src/api.md b/docs/src/api.md index 885be3720..4a0db690b 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -68,35 +68,17 @@ sol = solve(jprob, SSAStepper(); seed = 1234) # equivalent to: solve(jprob, SSAStepper(); rng = Xoshiro(1234)) ``` -### Passing `rng` via `JumpProblem` - -The `rng` keyword can also be passed to [`JumpProblem`](@ref), where it is -stored and automatically forwarded to the solver: - -```julia -jprob = JumpProblem(dprob, Direct(), jump; rng = StableRNG(1234)) -sol = solve(jprob, SSAStepper()) # uses StableRNG(1234) -``` - -An `rng` passed directly to `solve`/`init` takes priority over one stored in -the `JumpProblem`. - ### Resolution priority -When multiple RNG sources are provided, the following priority determines which -is used: +When both `rng` and `seed` are passed to the same `solve`/`init` call, `rng` +takes priority: | User provides | Result | |---|---| | `rng` via `solve`/`init` | Uses that `rng` | | `seed` via `solve`/`init` | Creates `Xoshiro(seed)` | -| `rng` via `JumpProblem` | Uses that `rng` | | Nothing | Uses `Random.default_rng()` (SSAStepper, ODE, tau-leaping) or a randomly-seeded `Xoshiro` (SDE) | -When both `rng` and `seed` are passed to the same call, `rng` takes priority. -When `solve`/`init` kwargs and `JumpProblem` kwargs overlap on the same key, -`solve`/`init` kwargs take priority. - ### Behavior by solver pathway | Solver | Default RNG (nothing passed) | `rng` / `seed` support | @@ -144,5 +126,5 @@ StochasticDiffEq has its own `_resolve_rng` that additionally handles `TaskLocalRNG` conversion and the problem's stored seed. For **tau-leaping**, JumpProcesses defines a custom `DiffEqBase.solve` that -bypasses the standard `__solve`/`__init` pathway. It calls `resolve_rng` with -the `JumpProblem`'s stored `rng` kwarg as a fallback. +bypasses the standard `__solve`/`__init` pathway. It calls `resolve_rng` +directly with the `rng` and `seed` kwargs from the `solve` call. diff --git a/docs/src/faq.md b/docs/src/faq.md index 0e3dad78c..47d77c3cb 100644 --- a/docs/src/faq.md +++ b/docs/src/faq.md @@ -75,14 +75,6 @@ sol = solve(jprob, SSAStepper(); rng = StableRNG(1234)) A `seed` keyword argument is also supported as a shorthand for creating a `Xoshiro` generator: `solve(jprob, SSAStepper(); seed = 1234)`. -Alternatively, `rng` can be passed to `JumpProblem` where it will be forwarded to the -solver automatically: - -```julia -jprob = JumpProblem(dprob, Direct(), maj; rng = Xoshiro(1234)) -sol = solve(jprob, SSAStepper()) # uses Xoshiro(1234) -``` - By default, JumpProcesses uses Julia's built-in `Random.default_rng()`. ## What are these aggregators and aggregations in JumpProcesses? diff --git a/src/problem.jl b/src/problem.jl index e09d92ddf..718179f7b 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -60,10 +60,6 @@ $(FIELDS) integration interface, and treated like general `VariableRateJump`s. - `vr_aggregator`, indicates the aggregator to use for sampling variable rate jumps. Current default is `VR_FRM`. - - `rng`, a random number generator to use for jump sampling. If provided, it is stored in - the problem's solver keyword arguments and forwarded to the solver when `solve` or `init` - is called. The recommended approach is to pass `rng` directly to `solve` or `init`: - `solve(jprob, SSAStepper(); rng = my_rng)`. - `tstops`, time stops to pass through to the solver. Can be an `AbstractVector` of times or a callable `(p, tspan) -> times`. @@ -242,10 +238,14 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS vr_aggregator::VariableRateAggregator = VR_FRM(), save_positions = prob isa DiffEqBase.AbstractDiscreteProblem ? (false, true) : (true, true), - rng = nothing, scale_rates = true, useiszero = true, + scale_rates = true, useiszero = true, spatial_system = nothing, hopping_constants = nothing, callback = nothing, tstops = nothing, use_vrj_bounds = true, kwargs...) + if haskey(kwargs, :rng) + throw(ArgumentError("`rng` is no longer a keyword argument for `JumpProblem`. Pass `rng` to `solve` or `init` instead, e.g. `solve(jprob, SSAStepper(); rng = my_rng)`.")) + end + # initialize the MassActionJump rate constants with the user parameters if using_params(jumps.massaction_jump) rates = jumps.massaction_jump.param_mapper(prob.p) @@ -309,7 +309,7 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS jump_cbs = CallbackSet(constant_jump_callback, variable_jump_callback) iip = isinplace_jump(prob, jumps.regular_jump) - solkwargs = make_kwarg(; filter(!isnothing, (; callback, tstops, rng))...) + solkwargs = make_kwarg(; filter(!isnothing, (; callback, tstops))...) JumpProblem{iip, typeof(new_prob), typeof(aggregator), typeof(jump_cbs), typeof(disc_agg), typeof(crjs), typeof(cvrjs), typeof(jumps.regular_jump), @@ -321,10 +321,14 @@ end function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; save_positions = prob isa DiffEqBase.AbstractDiscreteProblem ? (false, true) : (true, true), - rng = nothing, scale_rates = true, useiszero = true, + scale_rates = true, useiszero = true, spatial_system = nothing, hopping_constants = nothing, callback = nothing, tstops = nothing, kwargs...) + if haskey(kwargs, :rng) + throw(ArgumentError("`rng` is no longer a keyword argument for `JumpProblem`. Pass `rng` to `solve` or `init` instead, e.g. `solve(jprob, SSAStepper(); rng = my_rng)`.")) + end + # Validate no spatial systems (not currently supported) (spatial_system !== nothing || hopping_constants !== nothing) && error("PureLeaping does not currently support spatial problems.") @@ -351,7 +355,7 @@ function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; vrjs = jumps.variable_jumps iip = isinplace_jump(prob, jumps.regular_jump) - solkwargs = make_kwarg(; filter(!isnothing, (; callback, tstops, rng))...) + solkwargs = make_kwarg(; filter(!isnothing, (; callback, tstops))...) JumpProblem{iip, typeof(prob), typeof(aggregator), typeof(jump_cbs), typeof(disc_agg), typeof(crjs), typeof(vrjs), typeof(jumps.regular_jump), diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index 4d817b035..7c8df9fb3 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -77,7 +77,7 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping; error("SimpleTauLeaping can only be used with PureLeaping JumpProblems with only RegularJumps.") prob = jump_prob.prob - _rng = resolve_rng(rng, seed, get(jump_prob.kwargs, :rng, nothing)) + _rng = resolve_rng(rng, seed) rj = jump_prob.regular_jump rate = rj.rate # rate function rate(out,u,p,t) @@ -343,7 +343,7 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; error("SimpleExplicitTauLeaping can only be used with PureLeaping JumpProblem with a MassActionJump.") prob = jump_prob.prob - _rng = resolve_rng(rng, seed, get(jump_prob.kwargs, :rng, nothing)) + _rng = resolve_rng(rng, seed) tspan = prob.tspan if dtmin === nothing diff --git a/src/solve.jl b/src/solve.jl index 1ae5b3842..db58e3e46 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -1,17 +1,15 @@ """ - resolve_rng(rng, seed[, fallback_rng]) + resolve_rng(rng, seed) Resolve which RNG to use for a jump simulation. -Priority: `rng` > `seed` (creates `Xoshiro`) > `fallback_rng` > `Random.default_rng()`. +Priority: `rng` > `seed` (creates `Xoshiro`) > `Random.default_rng()`. """ -function resolve_rng(rng, seed, fallback_rng = nothing) +function resolve_rng(rng, seed) if rng !== nothing rng elseif seed !== nothing Random.Xoshiro(seed) - elseif fallback_rng !== nothing - fallback_rng else Random.default_rng() end diff --git a/test/allocations.jl b/test/allocations.jl index 64efb7582..25bd02f1e 100644 --- a/test/allocations.jl +++ b/test/allocations.jl @@ -35,17 +35,17 @@ let u₀ = [999, 10, 0] tspan = (0.0, 250.0) dprob = DiscreteProblem(u₀, tspan, p) - jprob = JumpProblem(dprob, Direct(), maj, jump, jump2; save_positions, rng) - sol = solve(jprob, SSAStepper()) + jprob = JumpProblem(dprob, Direct(), maj, jump, jump2; save_positions) + sol = solve(jprob, SSAStepper(); rng) - al1 = @allocations solve(jprob, SSAStepper()) + al1 = @allocations solve(jprob, SSAStepper(); rng) tspan2 = (0.0, 2500.0) dprob2 = DiscreteProblem(u₀, tspan2, p) - jprob2 = JumpProblem(dprob2, Direct(), maj, jump, jump2; save_positions, rng) - sol2 = solve(jprob2, SSAStepper()) + jprob2 = JumpProblem(dprob2, Direct(), maj, jump, jump2; save_positions) + sol2 = solve(jprob2, SSAStepper(); rng) - al2 = @allocations solve(jprob2, SSAStepper()) + al2 = @allocations solve(jprob2, SSAStepper(); rng) @test al1 == al2 end @@ -56,7 +56,7 @@ let end function makeprob(; T = 100.0, alg = Direct(), save_positions = (false, false), - graphkwargs = (;), rng) + graphkwargs = (;)) r1(u, p, t) = rate(p[1], u[1], u[2], p[2]) * u[1] r2(u, p, t) = rate(p[1], u[2], u[1], p[2]) * u[2] r3(u, p, t) = p[3] * u[1] @@ -86,7 +86,7 @@ let ConstantRateJump(r3, aff3!), ConstantRateJump(r4, aff4!), ConstantRateJump(r5, aff5!), ConstantRateJump(r6, aff6!); - save_positions, rng, graphkwargs...) + save_positions, graphkwargs...) return jprob end @@ -99,15 +99,15 @@ let graphkwargs = (; dep_graph, vartojumps_map, jumptovars_map) @testset "Allocations for $agg" for agg in JumpProcesses.JUMP_AGGREGATORS - jprob1 = makeprob(; alg = agg, T = 10.0, graphkwargs, rng = StableRNG(1234)) + jprob1 = makeprob(; alg = agg, T = 10.0, graphkwargs) stepper = SSAStepper() - sol1 = solve(jprob1, stepper) - sol1 = solve(jprob1, stepper) - al1 = @allocated solve(jprob1, stepper) - jprob2 = makeprob(; alg = agg, T = 100.0, graphkwargs, rng = StableRNG(1234)) - sol2 = solve(jprob2, stepper) - sol2 = solve(jprob2, stepper) - al2 = @allocated solve(jprob2, stepper) + sol1 = solve(jprob1, stepper; rng = StableRNG(1234)) + sol1 = solve(jprob1, stepper; rng = StableRNG(1234)) + al1 = @allocated solve(jprob1, stepper; rng = StableRNG(1234)) + jprob2 = makeprob(; alg = agg, T = 100.0, graphkwargs) + sol2 = solve(jprob2, stepper; rng = StableRNG(1234)) + sol2 = solve(jprob2, stepper; rng = StableRNG(1234)) + al2 = @allocated solve(jprob2, stepper; rng = StableRNG(1234)) @test al1 == al2 end end diff --git a/test/bimolerx_test.jl b/test/bimolerx_test.jl index befeefab1..832968b77 100644 --- a/test/bimolerx_test.jl +++ b/test/bimolerx_test.jl @@ -55,10 +55,14 @@ jump_to_dep_specs = [[1, 2], [1, 2], [1, 2, 3], [1, 2, 3], [1, 3]] majumps = MassActionJump(rates, reactstoch, netstoch) # average number of proteins in a simulation -function runSSAs(jump_prob; use_stepper = true) +function runSSAs(jump_prob; use_stepper = true, rng = nothing) Psamp = zeros(Int, Nsims) for i in 1:Nsims - sol = use_stepper ? solve(jump_prob, SSAStepper()) : solve(jump_prob) + sol = if use_stepper + isnothing(rng) ? solve(jump_prob, SSAStepper()) : solve(jump_prob, SSAStepper(); rng) + else + isnothing(rng) ? solve(jump_prob) : solve(jump_prob; rng) + end Psamp[i] = sol[1, end] end mean(Psamp) @@ -72,8 +76,8 @@ if doplot for alg in SSAalgs local jump_prob = JumpProblem(prob, alg, majumps, vartojumps_map = spec_to_dep_jumps, - jumptovars_map = jump_to_dep_specs, rng = rng) - local sol = solve(jump_prob, SSAStepper()) + jumptovars_map = jump_to_dep_specs) + local sol = solve(jump_prob, SSAStepper(); rng) local plothand = plot(sol, seriestype = :steppost, reuse = false) display(plothand) end @@ -84,15 +88,15 @@ if dotestmean for (i, alg) in enumerate(SSAalgs) local jump_prob = JumpProblem(prob, alg, majumps, save_positions = (false, false), vartojumps_map = spec_to_dep_jumps, - jumptovars_map = jump_to_dep_specs, rng = rng) - means = runSSAs(jump_prob) + jumptovars_map = jump_to_dep_specs) + means = runSSAs(jump_prob; rng) relerr = abs(means - expected_avg) / expected_avg doprintmeans && println("Mean from method: ", typeof(alg), " is = ", means, ", rel err = ", relerr) @test abs(means - expected_avg) < reltol * expected_avg # test not specifying SSAStepper - means = runSSAs(jump_prob; use_stepper = false) + means = runSSAs(jump_prob; use_stepper = false, rng) relerr = abs(means - expected_avg) / expected_avg @test abs(means - expected_avg) < reltol * expected_avg end @@ -107,8 +111,8 @@ if dotestmean jset = JumpSet((), (), nothing, majump_vec) jump_prob = JumpProblem(prob, Direct(), jset, save_positions = (false, false), vartojumps_map = spec_to_dep_jumps, - jumptovars_map = jump_to_dep_specs, rng = rng) - meanval = runSSAs(jump_prob) + jumptovars_map = jump_to_dep_specs) + meanval = runSSAs(jump_prob; rng) relerr = abs(meanval - expected_avg) / expected_avg if doprintmeans println("Using individual MassActionJumps; Mean from method: ", typeof(Direct()), diff --git a/test/bracketing.jl b/test/bracketing.jl index 7a4776da4..182689432 100644 --- a/test/bracketing.jl +++ b/test/bracketing.jl @@ -49,7 +49,7 @@ t = 0.0 ### Aggregator ### mutable struct DummyAggregator{T, M, R, BD} <: - JP.AbstractSSAJumpAggregator{T, M, R, Nothing, Nothing} + JP.AbstractSSAJumpAggregator{T, M, R, Nothing} ulow::Vector{Int} uhigh::Vector{Int} cur_rate_low::Vector{T} diff --git a/test/callbacks.jl b/test/callbacks.jl index 6ac5f9547..ce45d0069 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -25,8 +25,8 @@ rng = StableRNG(12345) affect_cb!(integrator) = (cb_called[] = true) cb = ContinuousCallback(condition, affect_cb!) - jprob = JumpProblem(prob, Direct(), jump; rng, callback = cb) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, Direct(), jump; callback = cb) + sol = solve(jprob, Tsit5(); rng) @test cb_called[] @test sol.t[end] ≈ 10.0 @@ -37,8 +37,8 @@ rng = StableRNG(12345) affect_dcb!(integrator) = (dcb_called[] += 1) dcb = DiscreteCallback(condition_d, affect_dcb!) - jprob = JumpProblem(prob, Direct(), jump; rng, callback = dcb) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, Direct(), jump; callback = dcb) + sol = solve(jprob, Tsit5(); rng) @test dcb_called[] > 0 # Should have fired multiple times @@ -47,8 +47,8 @@ rng = StableRNG(12345) affect_term!(integrator) = terminate!(integrator) cb_term = ContinuousCallback(condition_term, affect_term!) - jprob = JumpProblem(prob, Direct(), jump; rng, callback = cb_term) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, Direct(), jump; callback = cb_term) + sol = solve(jprob, Tsit5(); rng) @test sol.t[end] ≈ 3.0 # Should terminate at t=3 @@ -57,8 +57,8 @@ rng = StableRNG(12345) affect_mod!(integrator) = (integrator.u[1] *= 2.0) cb_mod = ContinuousCallback(condition_mod, affect_mod!) - jprob = JumpProblem(prob, Direct(), jump; rng, callback = cb_mod) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, Direct(), jump; callback = cb_mod) + sol = solve(jprob, Tsit5(); rng) # Check that state was modified at t=5 idx = findfirst(t -> t >= 5.0, sol.t) @@ -95,8 +95,8 @@ end # Test 1: Both callbacks should fire (default merge_callbacks = true) cb1_count[] = 0 cb2_count[] = 0 - jprob = JumpProblem(prob, Direct(), jump; rng, callback = cb1) - sol = solve(jprob, Tsit5(); callback = cb2) + jprob = JumpProblem(prob, Direct(), jump; callback = cb1) + sol = solve(jprob, Tsit5(); callback = cb2, rng) @test cb1_count[] > 0 @test cb2_count[] > 0 @@ -110,8 +110,8 @@ end # Test 2: Only solve callback should fire (merge_callbacks = false) cb1_count[] = 0 cb2_count[] = 0 - jprob = JumpProblem(prob, Direct(), jump; rng, callback = cb1) - sol = solve(jprob, Tsit5(); callback = cb2, merge_callbacks = false) + jprob = JumpProblem(prob, Direct(), jump; callback = cb1) + sol = solve(jprob, Tsit5(); callback = cb2, merge_callbacks = false, rng) @test cb1_count[] == 0 # Should not fire @test cb2_count[] == 1 # Should fire exactly once @@ -139,8 +139,8 @@ end cb = ContinuousCallback(condition, affect_cb!) # Callback in JumpProblem constructor - jprob = JumpProblem(prob, Direct(), jump; rng, callback = cb) - integrator = init(jprob, Tsit5()) + jprob = JumpProblem(prob, Direct(), jump; callback = cb) + integrator = init(jprob, Tsit5(); rng) solve!(integrator) @test cb_called[] @@ -175,8 +175,8 @@ end # Create CallbackSet with both types cbset = CallbackSet(ccb, dcb) - jprob = JumpProblem(prob, Direct(), jump; rng, callback = cbset) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, Direct(), jump; callback = cbset) + sol = solve(jprob, Tsit5(); rng) @test ccb_called[] @test dcb_called[] > 0 @@ -202,8 +202,8 @@ end affect_term!(integrator) = (cb_called[] = true; terminate!(integrator)) dcb_term = DiscreteCallback(condition_term, affect_term!) - jprob = JumpProblem(dprob, Direct(), jump1; rng, callback = dcb_term) - sol = solve(jprob, SSAStepper()) + jprob = JumpProblem(dprob, Direct(), jump1; callback = dcb_term) + sol = solve(jprob, SSAStepper(); rng) @test cb_called[] # Should have fired @test sol.u[end][1] >= 5 # Should have reached threshold @@ -215,8 +215,8 @@ end affect_count!(integrator) = (dcb_counter[] += 1) dcb_count = DiscreteCallback(condition_count, affect_count!) - jprob2 = JumpProblem(dprob, Direct(), jump1; rng) - sol2 = solve(jprob2, SSAStepper(); callback = dcb_count) + jprob2 = JumpProblem(dprob, Direct(), jump1) + sol2 = solve(jprob2, SSAStepper(); callback = dcb_count, rng) @test dcb_counter[] > 0 # Should have fired at least once @@ -232,8 +232,8 @@ end affect_cb2!(integrator) = (cb2_count[] += 1) dcb2 = DiscreteCallback(condition2, affect_cb2!) - jprob3 = JumpProblem(dprob, Direct(), jump1; rng, callback = dcb1) - sol3 = solve(jprob3, SSAStepper(); callback = dcb2) + jprob3 = JumpProblem(dprob, Direct(), jump1; callback = dcb1) + sol3 = solve(jprob3, SSAStepper(); callback = dcb2, rng) @test cb1_count[] > 0 # First callback should fire @test cb2_count[] > 0 # Second callback should fire @@ -257,8 +257,8 @@ end affect_cb4!(integrator) = (cb4_called[] = true) dcb4 = DiscreteCallback(condition4, affect_cb4!) - jprob4 = JumpProblem(dprob, Direct(), jump1; rng, callback = dcb3) - sol4 = solve(jprob4, SSAStepper(); callback = dcb4, merge_callbacks = false) + jprob4 = JumpProblem(dprob, Direct(), jump1; callback = dcb3) + sol4 = solve(jprob4, SSAStepper(); callback = dcb4, merge_callbacks = false, rng) @test !cb3_called[] # First callback should NOT fire @test cb4_called[] # Second callback should fire @@ -288,8 +288,8 @@ end cb = ContinuousCallback(condition, affect_cb!) # This was broken in v9.17.0 - callback wouldn't fire - jprob = JumpProblem(prob, Direct(), jump; rng, callback = cb) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, Direct(), jump; callback = cb) + sol = solve(jprob, Tsit5(); rng) @test cb_called[] @test sol.t[end] ≈ 0.5 # Should terminate at 0.5, not run to 1.0 @@ -323,8 +323,8 @@ end affect_d!(integrator) = (dcb_called[] += 1) dcb = DiscreteCallback(condition_d, affect_d!) - jprob = JumpProblem(prob, Direct(), jump; rng, callback = ccb) - sol = solve(jprob, Tsit5(); callback = dcb) + jprob = JumpProblem(prob, Direct(), jump; callback = ccb) + sol = solve(jprob, Tsit5(); callback = dcb, rng) @test ccb_called[] # Continuous callback should fire @test dcb_called[] > 0 # Discrete callback should fire multiple times @@ -357,8 +357,8 @@ end affect_c!(integrator) = (ccb_called[] = true) ccb = ContinuousCallback(condition_c, affect_c!) - jprob = JumpProblem(prob, Direct(), jump; rng, callback = dcb) - sol = solve(jprob, Tsit5(); callback = ccb) + jprob = JumpProblem(prob, Direct(), jump; callback = dcb) + sol = solve(jprob, Tsit5(); callback = ccb, rng) @test dcb_called[] > 0 # Discrete callback should fire @test ccb_called[] # Continuous callback should fire @@ -398,8 +398,8 @@ end affect3!(integrator) = (cb3_called[] = true) ccb2 = ContinuousCallback(condition3, affect3!) - jprob = JumpProblem(prob, Direct(), jump; rng, callback = cbset) - sol = solve(jprob, Tsit5(); callback = ccb2) + jprob = JumpProblem(prob, Direct(), jump; callback = cbset) + sol = solve(jprob, Tsit5(); callback = ccb2, rng) @test cb1_called[] # First continuous callback should fire @test cb2_called[] > 0 # Discrete callback should fire @@ -421,12 +421,12 @@ end affect_cb!(integrator) = nothing ccb = ContinuousCallback(condition, affect_cb!) - jprob_ccb = JumpProblem(dprob, Direct(), jump; rng, callback = ccb) - @test_throws ErrorException solve(jprob_ccb, SSAStepper()) + jprob_ccb = JumpProblem(dprob, Direct(), jump; callback = ccb) + @test_throws ErrorException solve(jprob_ccb, SSAStepper(); rng) # Test 2: ContinuousCallback passed to solve should error - jprob = JumpProblem(dprob, Direct(), jump; rng) - @test_throws ErrorException solve(jprob, SSAStepper(); callback = ccb) + jprob = JumpProblem(dprob, Direct(), jump) + @test_throws ErrorException solve(jprob, SSAStepper(); callback = ccb, rng) # Test 3: CallbackSet with continuous callbacks passed to JumpProblem should error on solve condition_d(u, t, integrator) = true @@ -434,19 +434,19 @@ end dcb = DiscreteCallback(condition_d, affect_dcb!) cbset_with_continuous = CallbackSet(ccb, dcb) - jprob_cbset = JumpProblem(dprob, Direct(), jump; rng, callback = cbset_with_continuous) - @test_throws ErrorException solve(jprob_cbset, SSAStepper()) + jprob_cbset = JumpProblem(dprob, Direct(), jump; callback = cbset_with_continuous) + @test_throws ErrorException solve(jprob_cbset, SSAStepper(); rng) # Test 4: CallbackSet with continuous callbacks passed to solve should error - @test_throws ErrorException solve(jprob, SSAStepper(); callback = cbset_with_continuous) + @test_throws ErrorException solve(jprob, SSAStepper(); callback = cbset_with_continuous, rng) # Test 5: CallbackSet with multiple continuous callbacks should error with correct count ccb2 = ContinuousCallback(condition, affect_cb!) cbset_multi = CallbackSet(ccb, ccb2, dcb) - jprob_multi = JumpProblem(dprob, Direct(), jump; rng, callback = cbset_multi) + jprob_multi = JumpProblem(dprob, Direct(), jump; callback = cbset_multi) err = try - solve(jprob_multi, SSAStepper()) + solve(jprob_multi, SSAStepper(); rng) nothing catch e e @@ -457,18 +457,18 @@ end # Test 6: DiscreteCallbacks should work fine (no error) dcb_only = DiscreteCallback(condition_d, affect_dcb!) - jprob_dcb = JumpProblem(dprob, Direct(), jump; rng, callback = dcb_only) - sol = solve(jprob_dcb, SSAStepper()) + jprob_dcb = JumpProblem(dprob, Direct(), jump; callback = dcb_only) + sol = solve(jprob_dcb, SSAStepper(); rng) @test sol.retcode == ReturnCode.Success # Test 7: CallbackSet with only discrete callbacks should work dcb2 = DiscreteCallback(condition_d, affect_dcb!) cbset_discrete = CallbackSet(dcb_only, dcb2) - jprob_dcb2 = JumpProblem(dprob, Direct(), jump; rng, callback = cbset_discrete) - sol2 = solve(jprob_dcb2, SSAStepper()) + jprob_dcb2 = JumpProblem(dprob, Direct(), jump; callback = cbset_discrete) + sol2 = solve(jprob_dcb2, SSAStepper(); rng) @test sol2.retcode == ReturnCode.Success # Test 8: Error should also be thrown with init - @test_throws ErrorException init(jprob_ccb, SSAStepper()) - @test_throws ErrorException init(jprob, SSAStepper(); callback = ccb) + @test_throws ErrorException init(jprob_ccb, SSAStepper(); rng) + @test_throws ErrorException init(jprob, SSAStepper(); callback = ccb, rng) end diff --git a/test/constant_rate.jl b/test/constant_rate.jl index 86c237c06..635be26f8 100644 --- a/test/constant_rate.jl +++ b/test/constant_rate.jl @@ -16,37 +16,37 @@ end jump2 = ConstantRateJump(rate, affect!) prob = DiscreteProblem(1.0, (0.0, 3.0)) -jump_prob = JumpProblem(prob, Direct(), jump; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump) -sol = solve(jump_prob, FunctionMap()) +sol = solve(jump_prob, FunctionMap(); rng) # using Plots; plot(sol) prob = DiscreteProblem(10.0, (0.0, 3.0)) -jump_prob = JumpProblem(prob, Direct(), jump, jump2; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump, jump2) -sol = solve(jump_prob, FunctionMap()) +sol = solve(jump_prob, FunctionMap(); rng) # plot(sol) nums = Int[] @time for i in 1:10000 - local jump_prob = JumpProblem(prob, Direct(), jump, jump2; rng = rng) - local sol = solve(jump_prob, FunctionMap()) + local jump_prob = JumpProblem(prob, Direct(), jump, jump2) + local sol = solve(jump_prob, FunctionMap(); rng) push!(nums, sol.u[end]) end @test mean(nums) - 45 < 1 prob = DiscreteProblem(1.0, (0.0, 3.0)) -jump_prob = JumpProblem(prob, Direct(), jump, jump2; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump, jump2) -sol = solve(jump_prob, FunctionMap()) +sol = solve(jump_prob, FunctionMap(); rng) nums = Int[] @time for i in 1:10000 - local jump_prob = JumpProblem(prob, Direct(), jump, jump2; rng = rng) - local sol = solve(jump_prob, FunctionMap()) + local jump_prob = JumpProblem(prob, Direct(), jump, jump2) + local sol = solve(jump_prob, FunctionMap(); rng) push!(nums, sol.u[2]) end diff --git a/test/ensemble_problems.jl b/test/ensemble_problems.jl index 541e7cf12..7cfc7066f 100644 --- a/test/ensemble_problems.jl +++ b/test/ensemble_problems.jl @@ -206,12 +206,13 @@ end # 7. JumpProblem rng kwarg forwarded to solver # ========================================================================== -@testset "JumpProblem rng kwarg forwarded to solver" begin +@testset "JumpProblem rng kwarg throws ArgumentError" begin j1 = ConstantRateJump((u, p, t) -> 1.0, integrator -> (integrator.u[1] += 1)) dprob = DiscreteProblem([10], (0.0, 10.0)) - jprob = JumpProblem(dprob, Direct(), j1; rng = StableRNG(1)) - @test haskey(jprob.kwargs, :rng) - @test jprob.kwargs[:rng] isa StableRNG - sol = solve(jprob, SSAStepper()) + j1_local = ConstantRateJump((u, p, t) -> 1.0, integrator -> (integrator.u[1] += 1)) + dprob_local = DiscreteProblem([10], (0.0, 10.0)) + jprob = JumpProblem(dprob_local, Direct(), j1_local) + @test_throws ArgumentError JumpProblem(dprob_local, Direct(), j1_local; rng = StableRNG(1)) + sol = solve(jprob, SSAStepper(); rng = StableRNG(1)) @test sol.retcode == ReturnCode.Success end diff --git a/test/ensemble_uniqueness.jl b/test/ensemble_uniqueness.jl index adb97df30..9861231cc 100644 --- a/test/ensemble_uniqueness.jl +++ b/test/ensemble_uniqueness.jl @@ -7,35 +7,19 @@ u0 = [0] dprob = DiscreteProblem(u0, (0.0, 100.0)) -# For EnsembleProblems, use prob_func to create a new JumpProblem with unique RNG per trajectory. -# This ensures different trajectories while maintaining reproducibility. -# Generate seeds from a seeded RNG for reproducibility of ensemble results. -function make_seeded_prob_func(dprob, aggregator, jumps, base_rng) - return function prob_func(prob, i, repeat) - seed = rand(base_rng, UInt64) - JumpProblem(dprob, aggregator, jumps...; rng = StableRNG(seed)) - end -end - -# Test with FunctionMap - use prob_func to create JumpProblems with unique RNGs -rng1 = StableRNG(12345) -jump_prob = JumpProblem(dprob, Direct(), j1, j2; rng = rng1) -ensemble_rng = StableRNG(99999) # separate RNG for generating trajectory seeds -ensemble_prob = EnsembleProblem(jump_prob; - prob_func = make_seeded_prob_func(dprob, Direct(), (j1, j2), ensemble_rng)) -sol = solve(ensemble_prob, FunctionMap(), trajectories = 3) +# Test with FunctionMap - pass rng to solve so trajectories get unique sequences +jump_prob = JumpProblem(dprob, Direct(), j1, j2) +ensemble_prob = EnsembleProblem(jump_prob) +sol = solve(ensemble_prob, FunctionMap(), trajectories = 3; rng = StableRNG(12345)) @test Array(sol.u[1]) !== Array(sol.u[2]) @test Array(sol.u[1]) !== Array(sol.u[3]) @test Array(sol.u[2]) !== Array(sol.u[3]) @test eltype(sol.u[1].u[1]) == Int -# Test with SSAStepper - use prob_func to create JumpProblems with unique RNGs -rng2 = StableRNG(12345) -jump_prob = JumpProblem(dprob, Direct(), j1, j2; rng = rng2) -ensemble_rng2 = StableRNG(99999) # separate RNG for generating trajectory seeds -ensemble_prob2 = EnsembleProblem(jump_prob; - prob_func = make_seeded_prob_func(dprob, Direct(), (j1, j2), ensemble_rng2)) -sol = solve(ensemble_prob2, SSAStepper(), trajectories = 3) +# Test with SSAStepper - pass rng to solve so trajectories get unique sequences +jump_prob = JumpProblem(dprob, Direct(), j1, j2) +ensemble_prob2 = EnsembleProblem(jump_prob) +sol = solve(ensemble_prob2, SSAStepper(), trajectories = 3; rng = StableRNG(12345)) @test Array(sol.u[1]) !== Array(sol.u[2]) @test Array(sol.u[1]) !== Array(sol.u[3]) @test Array(sol.u[2]) !== Array(sol.u[3]) diff --git a/test/extinction_test.jl b/test/extinction_test.jl index 880254ba0..468c96743 100644 --- a/test/extinction_test.jl +++ b/test/extinction_test.jl @@ -21,9 +21,8 @@ algs = (JumpProcesses.JUMP_AGGREGATORS..., JumpProcesses.NullAggregator()) for n in 1:Nsims for ssa in algs - local jprob = JumpProblem(dprob, ssa, majump, save_positions = (false, false), - rng = rng) - local sol = solve(jprob, SSAStepper()) + local jprob = JumpProblem(dprob, ssa, majump, save_positions = (false, false)) + local sol = solve(jprob, SSAStepper(); rng) @test sol[1, end] == 0 @test sol.t[end] < Inf end @@ -33,9 +32,8 @@ u0 = SA[10] dprob = DiscreteProblem(u0, (0.0, 100.0), rates) for ssa in algs - local jprob = JumpProblem(dprob, ssa, majump, save_positions = (false, false), - rng = rng) - local sol = solve(jprob, SSAStepper(), saveat = 100.0) + local jprob = JumpProblem(dprob, ssa, majump, save_positions = (false, false)) + local sol = solve(jprob, SSAStepper(); saveat = 100.0, rng) @test sol[1, end] == 0 @test sol.t[end] < Inf end @@ -57,8 +55,8 @@ end et = ExtinctionTest() cb = DiscreteCallback(et, et, save_positions = (false, false)) dprob = DiscreteProblem(u0, (0.0, 1000.0), rates) -jprob = JumpProblem(dprob, Direct(), majump; save_positions = (false, false), rng = rng) -sol = solve(jprob, SSAStepper(), callback = cb, save_end = false) +jprob = JumpProblem(dprob, Direct(), majump; save_positions = (false, false)) +sol = solve(jprob, SSAStepper(); callback = cb, save_end = false, rng) @test sol.t[end] < 1000.0 # test terminate @@ -73,8 +71,8 @@ end cb = DiscreteCallback(extinction_condition2, extinction_affect!2, save_positions = (false, false)) dprob = DiscreteProblem(u0, (0.0, 1000.0), rates) -jprob = JumpProblem(dprob, majump; save_positions = (false, false), rng) -sol = solve(jprob; callback = cb, save_end = false) +jprob = JumpProblem(dprob, majump; save_positions = (false, false)) +sol = solve(jprob; callback = cb, save_end = false, rng) @test sol[1, end] == 1 @test sol.retcode == ReturnCode.Terminated @test sol.t[end] < 1000.0 diff --git a/test/fp_unknowns.jl b/test/fp_unknowns.jl index 6a8b5f6d6..8e1ecd9d7 100644 --- a/test/fp_unknowns.jl +++ b/test/fp_unknowns.jl @@ -34,11 +34,11 @@ function test(rng) Xmeans = zeros(length(SSAalgs)) Ymeans = zeros(length(SSAalgs)) for (j, agg) in enumerate(SSAalgs) - jprob = JumpProblem(dprob, agg, maj; save_positions = (false, false), rng, + jprob = JumpProblem(dprob, agg, maj; save_positions = (false, false), vartojumps_map = vtoj, jumptovars_map = jtov, dep_graph = dg, scale_rates = false) for i in 1:Nsims - sol = solve(jprob, SSAStepper()) + sol = solve(jprob, SSAStepper(); rng) Xmeans[j] += sol[1, end] Ymeans[j] += sol[2, end] end diff --git a/test/functionwrappers.jl b/test/functionwrappers.jl index 2f009ead4..0c1bc0a0b 100644 --- a/test/functionwrappers.jl +++ b/test/functionwrappers.jl @@ -12,18 +12,18 @@ let rateinterval = (u, p, t) -> 0.1) prob = DiscreteProblem([0.0], (0.0, 2.0), [1.0]) - jprob = JumpProblem(prob, Coevolve(), jump; dep_graph = [[1]], rng) + jprob = JumpProblem(prob, Coevolve(), jump; dep_graph = [[1]]) agg = jprob.discrete_jump_aggregation @test agg.affects! isa Vector{Any} - integ = init(jprob, SSAStepper()) + integ = init(jprob, SSAStepper(); rng) T = Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{typeof(integ)}}} @test agg.affects! isa T affs = agg.affects! sol_c = solve!(integ) # check the affects vector is unchanged from a second call - integ = init(jprob, SSAStepper()) + integ = init(jprob, SSAStepper(); rng) sol_c = solve!(integ) @test affs === agg.affects! @@ -31,7 +31,7 @@ let terminate_condition(u, t, integrator) = (return u[1] >= 1) terminate_affect!(integrator) = terminate!(integrator) terminate_cb = DiscreteCallback(terminate_condition, terminate_affect!) - integ2 = init(jprob, SSAStepper(); callback = terminate_cb) + integ2 = init(jprob, SSAStepper(); rng, callback = terminate_cb) T2 = Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{typeof(integ2)}}} @test T2 !== T @test agg.affects! isa T2 @@ -42,7 +42,7 @@ let solve!(integ2) # check affs2 is unchanged when solving again now - integ2 = init(jprob, SSAStepper(); callback = terminate_cb) + integ2 = init(jprob, SSAStepper(); rng, callback = terminate_cb) solve!(integ2) @test affs2 === agg.affects! end diff --git a/test/geneexpr_test.jl b/test/geneexpr_test.jl index ea55c82f9..d26d9b8a2 100644 --- a/test/geneexpr_test.jl +++ b/test/geneexpr_test.jl @@ -22,19 +22,28 @@ expected_avg = 5.926553750000000e+02 reltol = 0.01 # average number of proteins in a simulation -function runSSAs(jump_prob; use_stepper = true) +function runSSAs(jump_prob; use_stepper = true, rng = nothing) Psamp = zeros(Int, Nsims) for i in 1:Nsims - sol = use_stepper ? solve(jump_prob, SSAStepper()) : solve(jump_prob) + sol = if use_stepper + isnothing(rng) ? solve(jump_prob, SSAStepper()) : + solve(jump_prob, SSAStepper(); rng) + else + isnothing(rng) ? solve(jump_prob) : solve(jump_prob; rng) + end Psamp[i] = sol[3, end] end mean(Psamp) end -function runSSAs_ode(vrjprob) +function runSSAs_ode(vrjprob; rng = nothing) Psamp = zeros(Float64, Nsims) tsave = vrjprob.prob.tspan[2] - integrator = init(vrjprob, Tsit5(); saveat = tsave) + integrator = if isnothing(rng) + init(vrjprob, Tsit5(); saveat = tsave) + else + init(vrjprob, Tsit5(); saveat = tsave, rng) + end solve!(integrator) Psamp[1] = integrator.sol[3, end] for i in 2:Nsims @@ -94,8 +103,8 @@ if doplot for alg in SSAalgs local jump_prob = JumpProblem(prob, alg, majumps, vartojumps_map = spec_to_dep_jumps, - jumptovars_map = jump_to_dep_specs, rng = rng) - local sol = solve(jump_prob, SSAStepper()) + jumptovars_map = jump_to_dep_specs) + local sol = solve(jump_prob, SSAStepper(); rng) plot!(plothand, sol.t, sol[3, :], seriestype = :steppost) end display(plothand) @@ -106,8 +115,8 @@ if dotestmean for (i, alg) in enumerate(SSAalgs) local jump_prob = JumpProblem(prob, alg, majumps, save_positions = (false, false), vartojumps_map = spec_to_dep_jumps, - jumptovars_map = jump_to_dep_specs, rng = rng) - means = runSSAs(jump_prob) + jumptovars_map = jump_to_dep_specs) + means = runSSAs(jump_prob; rng) relerr = abs(means - expected_avg) / expected_avg doprintmeans && println("Mean from method: ", typeof(alg), " is = ", means, ", rel err = ", relerr) @@ -118,8 +127,8 @@ if dotestmean let alg = Direct() jump_prob = JumpProblem(prob, alg, majumps, save_positions = (false, false), vartojumps_map = spec_to_dep_jumps, - jumptovars_map = jump_to_dep_specs, rng = rng) - means = runSSAs(jump_prob; use_stepper = false) + jumptovars_map = jump_to_dep_specs) + means = runSSAs(jump_prob; use_stepper = false, rng) @test abs(means - expected_avg) < reltol * expected_avg end @@ -128,8 +137,8 @@ if dotestmean for alg in (Direct(), RSSA()) jump_probf = JumpProblem(probf, alg, majumps, save_positions = (false, false), vartojumps_map = spec_to_dep_jumps, - jumptovars_map = jump_to_dep_specs, rng = rng) - means = runSSAs(jump_probf) + jumptovars_map = jump_to_dep_specs) + means = runSSAs(jump_probf; rng) relerr = abs(means - expected_avg) / expected_avg doprintmeans && println("Mean from method (Float64 u0): ", typeof(alg), " is = ", means, ", rel err = ", relerr) @@ -139,11 +148,11 @@ end # no-aggregator tests jump_prob = JumpProblem(prob, majumps; save_positions = (false, false), - vartojumps_map = spec_to_dep_jumps, jumptovars_map = jump_to_dep_specs, rng) -@test abs(runSSAs(jump_prob) - expected_avg) < reltol * expected_avg + vartojumps_map = spec_to_dep_jumps, jumptovars_map = jump_to_dep_specs) +@test abs(runSSAs(jump_prob; rng) - expected_avg) < reltol * expected_avg -jump_prob = JumpProblem(prob, majumps, save_positions = (false, false), rng = rng) -@test abs(runSSAs(jump_prob) - expected_avg) < reltol * expected_avg +jump_prob = JumpProblem(prob, majumps, save_positions = (false, false)) +@test abs(runSSAs(jump_prob; rng) - expected_avg) < reltol * expected_avg # crj/vrj accuracy test # k1, DNA --> mRNA + DNA @@ -187,20 +196,20 @@ let VariableRateJump(r6, a6!, save_positions = (false, false))) prob = DiscreteProblem(u0, (0.0, tf), rates) - crjprob = JumpProblem(prob, crjs; save_positions = (false, false), rng) - @test abs(runSSAs(crjprob) - expected_avg) < reltol * expected_avg + crjprob = JumpProblem(prob, crjs; save_positions = (false, false)) + @test abs(runSSAs(crjprob; rng) - expected_avg) < reltol * expected_avg # vrjs are very slow so test on a shorter time span and compare to the crjs prob = DiscreteProblem(u0, (0.0, tf / 5), rates) - crjprob = JumpProblem(prob, crjs; save_positions = (false, false), rng) - crjmean = runSSAs(crjprob) + crjprob = JumpProblem(prob, crjs; save_positions = (false, false)) + crjmean = runSSAs(crjprob; rng) f(du, u, p, t) = (du .= 0; nothing) oprob = ODEProblem(f, u0f, (0.0, tf / 5), rates) for vr_agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) vrjprob = JumpProblem( - oprob, vrjs; vr_aggregator = vr_agg, save_positions = (false, false), rng) - vrjmean = runSSAs_ode(vrjprob) + oprob, vrjs; vr_aggregator = vr_agg, save_positions = (false, false)) + vrjmean = runSSAs_ode(vrjprob; rng) @test abs(vrjmean - crjmean) < reltol * crjmean end end diff --git a/test/gpu/regular_jumps.jl b/test/gpu/regular_jumps.jl index 5e60ee2d4..b35276cf4 100644 --- a/test/gpu/regular_jumps.jl +++ b/test/gpu/regular_jumps.jl @@ -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, PureLeaping(), rj; rng = StableRNG(12345)) + jump_prob = JumpProblem(prob_disc, PureLeaping(), rj) sol = solve(EnsembleProblem(jump_prob), SimpleTauLeaping(), EnsembleGPUKernel(CUDABackend()); trajectories = Nsims, dt = 1.0) diff --git a/test/hawkes_test.jl b/test/hawkes_test.jl index adf5a83dc..bae9ff2c5 100644 --- a/test/hawkes_test.jl +++ b/test/hawkes_test.jl @@ -65,7 +65,7 @@ function hawkes_problem(p, agg::Coevolve; u = [0.0], tspan = (0.0, 50.0), save_positions = (false, true), g = [[1]], h = [[]], uselrate = true, kwargs...) dprob = DiscreteProblem(u, tspan, p) jumps = hawkes_jump(u, g, h; uselrate) - jprob = JumpProblem(dprob, agg, jumps...; dep_graph = g, save_positions, rng) + jprob = JumpProblem(dprob, agg, jumps...; dep_graph = g, save_positions) return jprob end @@ -78,7 +78,7 @@ function hawkes_problem(p, agg; u = [0.0], tspan = (0.0, 50.0), save_positions = (false, true), g = [[1]], h = [[]], vr_aggregator = VR_FRM(), kwargs...) oprob = ODEProblem(f!, u, tspan, p) jumps = hawkes_jump(u, g, h) - jprob = JumpProblem(oprob, agg, jumps...; vr_aggregator, save_positions, rng) + jprob = JumpProblem(oprob, agg, jumps...; vr_aggregator, save_positions) return jprob end @@ -119,7 +119,7 @@ for (i, alg) in enumerate(algs) sols = Vector{ODESolution}(undef, Nsims) for n in 1:Nsims reset_history!(h) - sols[n] = solve(jump_prob, stepper) + sols[n] = solve(jump_prob, stepper; rng) end if alg isa Coevolve @@ -137,12 +137,12 @@ let alg = Coevolve() for vr_aggregator in (VR_FRM(), VR_Direct(), VR_DirectFW()) oprob = ODEProblem(f!, u0, tspan, p) jumps = hawkes_jump(u0, g, h) - jprob = JumpProblem(oprob, alg, jumps...; vr_aggregator, dep_graph = g, rng) + jprob = JumpProblem(oprob, alg, jumps...; vr_aggregator, dep_graph = g) @test ((jprob.variable_jumps === nothing) || isempty(jprob.variable_jumps)) sols = Vector{ODESolution}(undef, Nsims) for n in 1:Nsims reset_history!(h) - sols[n] = solve(jprob, Tsit5()) + sols[n] = solve(jprob, Tsit5(); rng) end λs = permutedims(mapreduce((sol) -> empirical_rate(sol), hcat, sols)) @test isapprox(mean(λs), Eλ; atol = 0.01) @@ -156,12 +156,12 @@ let alg = Coevolve() for vr_aggregator in (VR_FRM(), VR_Direct(), VR_DirectFW()) oprob = ODEProblem(f!, u0, tspan, p) jumps = hawkes_jump(u0, g, h) - jprob = JumpProblem(oprob, alg, jumps...; vr_aggregator, dep_graph = g, rng, use_vrj_bounds = false) + jprob = JumpProblem(oprob, alg, jumps...; vr_aggregator, dep_graph = g, use_vrj_bounds = false) @test length(jprob.variable_jumps) == 1 sols = Vector{ODESolution}(undef, Nsims) for n in 1:Nsims reset_history!(h) - sols[n] = solve(jprob, Tsit5()) + sols[n] = solve(jprob, Tsit5(); rng) end cols = length(u0) diff --git a/test/linearreaction_test.jl b/test/linearreaction_test.jl index 657982a6b..67e13b159 100644 --- a/test/linearreaction_test.jl +++ b/test/linearreaction_test.jl @@ -29,7 +29,7 @@ exactmeanval = exactmean(tf, rates) function runSSAs(jump_prob) Asamp = zeros(Int, Nsims) for i in 1:Nsims - sol = solve(jump_prob, SSAStepper()) + sol = solve(jump_prob, SSAStepper(); rng) Asamp[i] = sol[1, end] end mean(Asamp) @@ -52,7 +52,7 @@ function A_to_B_tuple(N, method) jumps = ((jump for jump in jumpvec)...,) jset = JumpSet((), jumps, nothing, nothing) prob = DiscreteProblem([A0, 0], (0.0, tf)) - jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), rng, + jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), namedpars...) jump_prob @@ -74,7 +74,7 @@ function A_to_B_vec(N, method) # convert jumpvec to tuple to send to JumpProblem... jset = JumpSet((), jumps, nothing, nothing) prob = DiscreteProblem([A0, 0], (0.0, tf)) - jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), rng, + jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), namedpars...) jump_prob @@ -92,7 +92,7 @@ function A_to_B_ma(N, method) majumps = MassActionJump(rates, reactstoch, netstoch) jset = JumpSet((), (), nothing, majumps) prob = DiscreteProblem([A0, 0], (0.0, tf)) - jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), rng, + jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), namedpars...) jump_prob @@ -126,7 +126,7 @@ function A_to_B_hybrid(N, method) majumps = MassActionJump(rates[1:switchidx], reactstoch, netstoch) jset = JumpSet((), jumps, nothing, majumps) prob = DiscreteProblem([A0, 0], (0.0, tf)) - jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), rng, + jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), namedpars...) jump_prob @@ -161,7 +161,7 @@ function A_to_B_hybrid_nojset(N, method) jumps = (constjumps..., majumps) prob = DiscreteProblem([A0, 0], (0.0, tf)) jump_prob = JumpProblem(prob, method, jumps...; save_positions = (false, false), - rng, namedpars...) + namedpars...) jump_prob end @@ -190,7 +190,7 @@ function A_to_B_hybrid_vecs(N, method) end jset = JumpSet((), jumpvec, nothing, majumps) prob = DiscreteProblem([A0, 0], (0.0, tf)) - jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), rng, + jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), namedpars...) jump_prob @@ -220,7 +220,7 @@ function A_to_B_hybrid_vecs_scalars(N, method) end jset = JumpSet((), jumpvec, nothing, majumps) prob = DiscreteProblem([A0, 0], (0.0, tf)) - jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), rng, + jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), namedpars...) jump_prob @@ -252,7 +252,7 @@ function A_to_B_hybrid_tups_scalars(N, method) jumps = ((maj for maj in majumpsv)..., (jump for jump in jumpvec)...) prob = DiscreteProblem([A0, 0], (0.0, tf)) jump_prob = JumpProblem(prob, method, jumps...; save_positions = (false, false), - rng, namedpars...) + namedpars...) jump_prob end @@ -282,7 +282,7 @@ function A_to_B_hybrid_tups(N, method) jumps = ((jump for jump in jumpvec)...,) jset = JumpSet((), jumps, nothing, majumps) prob = DiscreteProblem([A0, 0], (0.0, tf)) - jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), rng, + jump_prob = JumpProblem(prob, method, jset; save_positions = (false, false), namedpars...) jump_prob diff --git a/test/longtimes_test.jl b/test/longtimes_test.jl index 7c787b4f3..4a63422e1 100644 --- a/test/longtimes_test.jl +++ b/test/longtimes_test.jl @@ -11,6 +11,6 @@ u0 = [5] tspan = (0.0, 2e6) dt = tspan[2] / 1000 dprob = DiscreteProblem(u0, tspan, p) -jprob = JumpProblem(dprob, Direct(), maj, save_positions = (false, false), rng = rng) -sol = solve(jprob, SSAStepper(), saveat = tspan[1]:dt:tspan[2]) +jprob = JumpProblem(dprob, Direct(), maj, save_positions = (false, false)) +sol = solve(jprob, SSAStepper(); saveat = tspan[1]:dt:tspan[2], rng) @test length(unique(sol.u[(end - 10):end][:])) > 1 diff --git a/test/monte_carlo_test.jl b/test/monte_carlo_test.jl index c197df9d5..f675dc9fc 100644 --- a/test/monte_carlo_test.jl +++ b/test/monte_carlo_test.jl @@ -8,27 +8,27 @@ prob = SDEProblem(f, g, [1.0], (0.0, 1.0)) rate = (u, p, t) -> 200.0 affect! = integrator -> (integrator.u[1] = integrator.u[1] / 2) jump = VariableRateJump(rate, affect!, save_positions = (false, true)) -jump_prob = JumpProblem(prob, Direct(), jump; vr_aggregator = VR_FRM(), rng) +jump_prob = JumpProblem(prob, Direct(), jump; vr_aggregator = VR_FRM()) monte_prob = EnsembleProblem(jump_prob) -sol = solve(monte_prob, SRIW1(), EnsembleSerial(), trajectories = 3, - save_everystep = false, dt = 0.001, adaptive = false) +sol = solve(monte_prob, SRIW1(), EnsembleSerial(); trajectories = 3, + save_everystep = false, dt = 0.001, adaptive = false, rng) @test sol.u[1].t[2] != sol.u[2].t[2] != sol.u[3].t[2] -jump_prob = JumpProblem(prob, Direct(), jump; vr_aggregator = VR_Direct(), rng) +jump_prob = JumpProblem(prob, Direct(), jump; vr_aggregator = VR_Direct()) monte_prob = EnsembleProblem(jump_prob) -sol = solve(monte_prob, SRIW1(), EnsembleSerial(), trajectories = 3, - save_everystep = false, dt = 0.001, adaptive = false) +sol = solve(monte_prob, SRIW1(), EnsembleSerial(); trajectories = 3, + save_everystep = false, dt = 0.001, adaptive = false, rng) @test allunique(sol.u[1].t) -jump_prob = JumpProblem(prob, Direct(), jump; vr_aggregator = VR_DirectFW(), rng) +jump_prob = JumpProblem(prob, Direct(), jump; vr_aggregator = VR_DirectFW()) monte_prob = EnsembleProblem(jump_prob) -sol = solve(monte_prob, SRIW1(), EnsembleSerial(), trajectories = 3, - save_everystep = false, dt = 0.001, adaptive = false) +sol = solve(monte_prob, SRIW1(), EnsembleSerial(); trajectories = 3, + save_everystep = false, dt = 0.001, adaptive = false, rng) @test allunique(sol.u[1].t) jump = ConstantRateJump(rate, affect!) -jump_prob = JumpProblem(prob, Direct(), jump; save_positions = (true, false), rng) +jump_prob = JumpProblem(prob, Direct(), jump; save_positions = (true, false)) monte_prob = EnsembleProblem(jump_prob) -sol = solve(monte_prob, SRIW1(), EnsembleSerial(), trajectories = 3, - save_everystep = false, dt = 0.001, adaptive = false) +sol = solve(monte_prob, SRIW1(), EnsembleSerial(); trajectories = 3, + save_everystep = false, dt = 0.001, adaptive = false, rng) @test sol.u[1].t[2] != sol.u[2].t[2] != sol.u[3].t[2] diff --git a/test/regular_jumps.jl b/test/regular_jumps.jl index f4d009eae..99c3375c7 100644 --- a/test/regular_jumps.jl +++ b/test/regular_jumps.jl @@ -36,11 +36,11 @@ end u0 = [999.0, 10.0, 0.0] # S, I, R tspan = (0.0, 250.0) prob_disc = DiscreteProblem(u0, tspan, p) - jump_prob = JumpProblem(prob_disc, Direct(), jumps...; rng, save_positions = (false, false)) + jump_prob = JumpProblem(prob_disc, Direct(), jumps...; save_positions = (false, false)) # Solve with SSAStepper (save only at t_compare times) sol_direct = solve(EnsembleProblem(jump_prob), SSAStepper(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare) + trajectories = Nsims, saveat = t_compare, rng) # RegularJump formulation for SimpleTauLeaping regular_rate = (out, u, p, t) -> begin @@ -55,22 +55,22 @@ end dc[3] = counts[2] end rj = RegularJump(regular_rate, regular_c, 3) - jump_prob_tau = JumpProblem(prob_disc, PureLeaping(), rj; rng) + jump_prob_tau = JumpProblem(prob_disc, PureLeaping(), rj) # Solve with SimpleTauLeaping (save only at t_compare times) sol_simple = solve(EnsembleProblem(jump_prob_tau), SimpleTauLeaping(), EnsembleSerial(); - trajectories = Nsims, dt = 0.1, saveat = t_compare) + trajectories = Nsims, dt = 0.1, saveat = t_compare, rng) # MassActionJump formulation for SimpleExplicitTauLeaping reactant_stoich = [[1 => 1, 2 => 1], [2 => 1], Pair{Int, Int}[]] net_stoich = [[1 => -1, 2 => 1], [2 => -1, 3 => 1], [1 => 1]] param_idxs = [1, 2, 3] maj = MassActionJump(reactant_stoich, net_stoich; param_idxs) - jump_prob_maj = JumpProblem(prob_disc, PureLeaping(), maj; rng) + jump_prob_maj = JumpProblem(prob_disc, PureLeaping(), maj) # Solve with SimpleExplicitTauLeaping (save only at t_compare times) sol_adaptive = solve(EnsembleProblem(jump_prob_maj), SimpleExplicitTauLeaping(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare) + trajectories = Nsims, saveat = t_compare, rng) # Compute mean I trajectories via direct indexing (I is index 2 in SIR) mean_I_direct = compute_mean_at_saves(sol_direct, Nsims, npts, 2) @@ -101,11 +101,11 @@ end u0 = [999.0, 0.0, 10.0, 0.0] # S, E, I, R tspan = (0.0, 250.0) prob_disc = DiscreteProblem(u0, tspan, p) - jump_prob = JumpProblem(prob_disc, Direct(), jumps...; rng, save_positions = (false, false)) + jump_prob = JumpProblem(prob_disc, Direct(), jumps...; save_positions = (false, false)) # Solve with SSAStepper (save only at t_compare times) sol_direct = solve(EnsembleProblem(jump_prob), SSAStepper(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare) + trajectories = Nsims, saveat = t_compare, rng) # RegularJump formulation for SimpleTauLeaping regular_rate = (out, u, p, t) -> begin @@ -121,22 +121,22 @@ end dc[4] = counts[3] end rj = RegularJump(regular_rate, regular_c, 3) - jump_prob_tau = JumpProblem(prob_disc, PureLeaping(), rj; rng) + jump_prob_tau = JumpProblem(prob_disc, PureLeaping(), rj) # Solve with SimpleTauLeaping (save only at t_compare times) sol_simple = solve(EnsembleProblem(jump_prob_tau), SimpleTauLeaping(), EnsembleSerial(); - trajectories = Nsims, dt = 0.1, saveat = t_compare) + trajectories = Nsims, dt = 0.1, saveat = t_compare, rng) # MassActionJump formulation for SimpleExplicitTauLeaping reactant_stoich = [[1 => 1, 3 => 1], [2 => 1], [3 => 1]] net_stoich = [[1 => -1, 2 => 1], [2 => -1, 3 => 1], [3 => -1, 4 => 1]] param_idxs = [1, 2, 3] maj = MassActionJump(reactant_stoich, net_stoich; param_idxs) - jump_prob_maj = JumpProblem(prob_disc, PureLeaping(), maj; rng) + jump_prob_maj = JumpProblem(prob_disc, PureLeaping(), maj) # Solve with SimpleExplicitTauLeaping (save only at t_compare times) sol_adaptive = solve(EnsembleProblem(jump_prob_maj), SimpleExplicitTauLeaping(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare) + trajectories = Nsims, saveat = t_compare, rng) # Compute mean I trajectories via direct indexing (I is index 3 in SEIR) mean_I_direct = compute_mean_at_saves(sol_direct, Nsims, npts, 3) @@ -183,7 +183,7 @@ end maj = MassActionJump(rates, reactant_stoich, net_stoich) # Test PureLeaping JumpProblem creation - jp_pure = JumpProblem(prob, PureLeaping(), JumpSet(maj); rng) + 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 @@ -194,7 +194,7 @@ end affect!(integrator) = (integrator.u[1] -= 1; integrator.u[3] += 1) crj = ConstantRateJump(rate, affect!) - jp_pure_crj = JumpProblem(prob, PureLeaping(), JumpSet(crj); rng) + 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 @@ -204,7 +204,7 @@ end vaffect!(integrator) = (integrator.u[1] -= 1; integrator.u[3] += 1) vrj = VariableRateJump(vrate, vaffect!) - jp_pure_vrj = JumpProblem(prob, PureLeaping(), JumpSet(vrj); rng) + 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 @@ -224,7 +224,7 @@ end regj = RegularJump(rj_rate, rj_c, 1) - jp_pure_regj = JumpProblem(prob, PureLeaping(), JumpSet(regj); rng) + 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 @@ -232,7 +232,7 @@ end # 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; rng) + 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 @@ -243,14 +243,14 @@ end # Test spatial system error spatial_sys = CartesianGrid((2, 2)) hopping_consts = [1.0] - @test_throws ErrorException JumpProblem(prob, PureLeaping(), JumpSet(maj); rng, + @test_throws ErrorException JumpProblem(prob, PureLeaping(), JumpSet(maj); spatial_system = spatial_sys) - @test_throws ErrorException JumpProblem(prob, PureLeaping(), JumpSet(maj); rng, + @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); rng) + 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 @@ -266,23 +266,23 @@ end prob = DiscreteProblem(u0, tspan) # SSAStepper with save_positions=(false,false) + saveat: only saveat times stored - jp = JumpProblem(prob, Direct(), crj; rng, save_positions = (false, false)) - sol = solve(jp, SSAStepper(); saveat = 1.0) + jp = JumpProblem(prob, Direct(), crj; save_positions = (false, false)) + sol = solve(jp, SSAStepper(); saveat = 1.0, rng) @test sol.t == collect(0.0:1.0:10.0) # SSAStepper with default save_positions + saveat: jump times stored too - jp2 = JumpProblem(prob, Direct(), crj; rng) - sol2 = solve(jp2, SSAStepper(); saveat = 1.0) + jp2 = JumpProblem(prob, Direct(), crj) + sol2 = solve(jp2, SSAStepper(); saveat = 1.0, rng) @test length(sol2.t) > length(sol.t) # --- SimpleTauLeaping save_start/save_end/saveat tests --- regular_rate = (out, u, p, t) -> (out[1] = 1.0) regular_c = (dc, u, p, t, counts, mark) -> (dc[1] = counts[1]) rj = RegularJump(regular_rate, regular_c, 1) - jp_tau = JumpProblem(prob, PureLeaping(), rj; rng) + jp_tau = JumpProblem(prob, PureLeaping(), rj) # No saveat: stores every dt step (save_start=true, save_end=true by default) - sol_tau = solve(jp_tau, SimpleTauLeaping(); dt = 1.0) + sol_tau = solve(jp_tau, SimpleTauLeaping(); dt = 1.0, rng) @test sol_tau.t == collect(0.0:1.0:10.0) # saveat as Number: defaults save_start=true, save_end=true @@ -334,7 +334,7 @@ end reactant_stoich = [[1 => 1]] net_stoich = [[1 => -1]] maj = MassActionJump([0.1], reactant_stoich, net_stoich) - jp_explicit = JumpProblem(prob_decay, PureLeaping(), maj; rng) + jp_explicit = JumpProblem(prob_decay, PureLeaping(), maj) # saveat as Number: defaults save_start=true, save_end=true sol = solve(jp_explicit, SimpleExplicitTauLeaping(); saveat = 2.0) diff --git a/test/reversible_binding.jl b/test/reversible_binding.jl index f872d92a4..456a96d2a 100644 --- a/test/reversible_binding.jl +++ b/test/reversible_binding.jl @@ -20,10 +20,10 @@ tspan = (0.0, 5.0) prob = DiscreteProblem(u0, tspan, rates) majumps = MassActionJump(rates, reactstoch, netstoch) -function getmean(jprob, Nsims) +function getmean(jprob, Nsims; rng = nothing) Amean = 0 for i in 1:Nsims - sol = solve(jprob, SSAStepper()) + sol = isnothing(rng) ? solve(jprob, SSAStepper()) : solve(jprob, SSAStepper(); rng) Amean += sol[1, end] end Amean /= Nsims @@ -48,8 +48,7 @@ mastereq_mean = mastereqmean(u0, rates) algs = (JumpProcesses.JUMP_AGGREGATORS..., JumpProcesses.NullAggregator()) relative_tolerance = 0.01 for alg in algs - local jprob = JumpProblem(prob, alg, majumps, save_positions = (false, false), - rng = rng) - local Amean = getmean(jprob, Nsims) + local jprob = JumpProblem(prob, alg, majumps, save_positions = (false, false)) + local Amean = getmean(jprob, Nsims; rng) @test abs(Amean - mastereq_mean) / mastereq_mean < relative_tolerance end diff --git a/test/rng_kwarg_tests.jl b/test/rng_kwarg_tests.jl index a619edb76..a677a9118 100644 --- a/test/rng_kwarg_tests.jl +++ b/test/rng_kwarg_tests.jl @@ -2,191 +2,159 @@ using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq, Test using StableRNGs, Random # ========================================================================== -# Test that rng can be passed via JumpProblem kwargs OR solve kwargs, -# and that solve-level rng takes precedence over JumpProblem-level rng. -# -# Strategy: use different RNG *types* to verify which one the integrator -# receives. StableRNG is passed at one level and Xoshiro at another, -# then we check the type on the integrator. +# Test that rng/seed can be passed via solve/init kwargs for all pathways, +# and that JumpProblem(; rng=...) throws an error. # ========================================================================== # -------------------------------------------------------------------------- # Problem constructors # -------------------------------------------------------------------------- -function make_ssa_jump_prob(; kwargs...) +function make_ssa_jump_prob() j1 = ConstantRateJump((u, p, t) -> 10.0, integrator -> (integrator.u[1] += 1)) j2 = ConstantRateJump((u, p, t) -> 0.5 * u[1], integrator -> (integrator.u[1] -= 1)) dprob = DiscreteProblem([10], (0.0, 20.0)) - JumpProblem(dprob, Direct(), j1, j2; kwargs...) + JumpProblem(dprob, Direct(), j1, j2) end -function make_ode_vr_jump_prob(; kwargs...) +function make_ode_vr_jump_prob() f!(du, u, p, t) = (du[1] = -0.1 * u[1]; nothing) oprob = ODEProblem(f!, [100.0], (0.0, 10.0)) vrj = VariableRateJump((u, p, t) -> 0.5 * u[1], integrator -> (integrator.u[1] -= 1.0)) - JumpProblem(oprob, Direct(), vrj; kwargs...) + JumpProblem(oprob, Direct(), vrj) end -function make_sde_vr_jump_prob(; kwargs...) +function make_sde_vr_jump_prob() f!(du, u, p, t) = (du[1] = -0.1 * u[1]; nothing) g!(du, u, p, t) = (du[1] = 0.1 * u[1]; nothing) sprob = SDEProblem(f!, g!, [100.0], (0.0, 10.0)) vrj = VariableRateJump((u, p, t) -> 0.5 * u[1], integrator -> (integrator.u[1] -= 1.0)) - JumpProblem(sprob, Direct(), vrj; kwargs...) + JumpProblem(sprob, Direct(), vrj) end # ========================================================================== -# 1. SSAStepper: rng via JumpProblem +# 1. JumpProblem(; rng=...) throws ArgumentError # ========================================================================== -@testset "SSAStepper: rng via JumpProblem kwargs" begin - xrng = Xoshiro(42) - jprob = make_ssa_jump_prob(; rng = xrng) - integrator = init(jprob, SSAStepper()) - @test SciMLBase.get_rng(integrator) isa Xoshiro - sol = solve(jprob, SSAStepper()) - @test sol.retcode == ReturnCode.Success +@testset "JumpProblem(; rng=...) throws ArgumentError" begin + dprob = DiscreteProblem([10], (0.0, 10.0)) + j1 = ConstantRateJump((u, p, t) -> 1.0, integrator -> (integrator.u[1] += 1)) + @test_throws ArgumentError JumpProblem(dprob, Direct(), j1; rng = StableRNG(42)) end # ========================================================================== -# 2. SSAStepper: rng via solve overrides JumpProblem +# 2. SSAStepper: rng via solve/init # ========================================================================== -@testset "SSAStepper: solve rng overrides JumpProblem rng" begin - jprob = make_ssa_jump_prob(; rng = Xoshiro(42)) - integrator = init(jprob, SSAStepper(); rng = StableRNG(99)) - @test SciMLBase.get_rng(integrator) isa StableRNG +@testset "SSAStepper: rng via solve kwargs" begin + jprob = make_ssa_jump_prob() + integrator = init(jprob, SSAStepper(); rng = Xoshiro(42)) + @test SciMLBase.get_rng(integrator) isa Xoshiro + sol = solve(jprob, SSAStepper(); rng = Xoshiro(42)) + @test sol.retcode == ReturnCode.Success end # ========================================================================== -# 3. SSAStepper: reproducibility via JumpProblem rng +# 3. SSAStepper: reproducibility via solve rng # ========================================================================== -@testset "SSAStepper: JumpProblem rng reproducibility" begin - jprob1 = make_ssa_jump_prob(; rng = StableRNG(123)) - jprob2 = make_ssa_jump_prob(; rng = StableRNG(123)) - sol1 = solve(jprob1, SSAStepper()) - sol2 = solve(jprob2, SSAStepper()) +@testset "SSAStepper: solve rng reproducibility" begin + jprob = make_ssa_jump_prob() + sol1 = solve(jprob, SSAStepper(); rng = StableRNG(123)) + sol2 = solve(jprob, SSAStepper(); rng = StableRNG(123)) @test sol1.t == sol2.t @test sol1.u == sol2.u end # ========================================================================== -# 4. SSAStepper: solve rng overrides for reproducibility +# 4. SSAStepper: different seeds → different trajectories # ========================================================================== -@testset "SSAStepper: solve rng override reproducibility" begin - jprob = make_ssa_jump_prob(; rng = Xoshiro(1)) - sol1 = solve(jprob, SSAStepper(); rng = StableRNG(42)) - sol2 = solve(jprob, SSAStepper(); rng = StableRNG(42)) - @test sol1.t == sol2.t - @test sol1.u == sol2.u +@testset "SSAStepper: different seeds → different trajectories" begin + jprob = make_ssa_jump_prob() + sol1 = solve(jprob, SSAStepper(); rng = StableRNG(100)) + sol2 = solve(jprob, SSAStepper(); rng = StableRNG(200)) + sol3 = solve(jprob, SSAStepper(); rng = StableRNG(300)) + times = [sol1.t[2], sol2.t[2], sol3.t[2]] + @test allunique(times) end # ========================================================================== -# 5. ODE + VR: rng via JumpProblem +# 5. ODE + VR: rng via solve/init # ========================================================================== -@testset "ODE + VR: rng via JumpProblem kwargs" begin - jprob = make_ode_vr_jump_prob(; rng = Xoshiro(42)) - integrator = init(jprob, Tsit5()) +@testset "ODE + VR: rng via solve kwargs" begin + jprob = make_ode_vr_jump_prob() + integrator = init(jprob, Tsit5(); rng = Xoshiro(42)) @test SciMLBase.get_rng(integrator) isa Xoshiro end # ========================================================================== -# 6. ODE + VR: solve rng overrides JumpProblem rng +# 6. ODE + VR: reproducibility via solve rng # ========================================================================== -@testset "ODE + VR: solve rng overrides JumpProblem rng" begin - jprob = make_ode_vr_jump_prob(; rng = Xoshiro(42)) - integrator = init(jprob, Tsit5(); rng = StableRNG(99)) - @test SciMLBase.get_rng(integrator) isa StableRNG -end - -# ========================================================================== -# 7. ODE + VR: reproducibility via JumpProblem rng -# ========================================================================== -@testset "ODE + VR: JumpProblem rng reproducibility" begin - jprob1 = make_ode_vr_jump_prob(; rng = StableRNG(123)) - jprob2 = make_ode_vr_jump_prob(; rng = StableRNG(123)) - sol1 = solve(jprob1, Tsit5()) - sol2 = solve(jprob2, Tsit5()) +@testset "ODE + VR: solve rng reproducibility" begin + jprob = make_ode_vr_jump_prob() + sol1 = solve(jprob, Tsit5(); rng = StableRNG(123)) + sol2 = solve(jprob, Tsit5(); rng = StableRNG(123)) @test sol1.t ≈ sol2.t @test sol1.u[end] ≈ sol2.u[end] end # ========================================================================== -# 8. ODE + VR: solve rng overrides for reproducibility +# 7. ODE + VR: different seeds → different trajectories # ========================================================================== -@testset "ODE + VR: solve rng override reproducibility" begin - jprob = make_ode_vr_jump_prob(; rng = Xoshiro(1)) - sol1 = solve(jprob, Tsit5(); rng = StableRNG(42)) - sol2 = solve(jprob, Tsit5(); rng = StableRNG(42)) - @test sol1.t ≈ sol2.t - @test sol1.u[end] ≈ sol2.u[end] +@testset "ODE + VR: different seeds → different trajectories" begin + jprob = make_ode_vr_jump_prob() + sols = [solve(jprob, Tsit5(); rng = StableRNG(s)) for s in (100, 200, 300)] + finals = [s.u[end][1] for s in sols] + @test allunique(finals) end # ========================================================================== -# 9. SDE + VR: rng via JumpProblem +# 8. SDE + VR: rng via solve/init # ========================================================================== -@testset "SDE + VR: rng via JumpProblem kwargs" begin - jprob = make_sde_vr_jump_prob(; rng = Xoshiro(42)) - integrator = init(jprob, EM(); dt = 0.01) +@testset "SDE + VR: rng via solve kwargs" begin + jprob = make_sde_vr_jump_prob() + integrator = init(jprob, EM(); dt = 0.01, rng = Xoshiro(42)) @test SciMLBase.get_rng(integrator) isa Xoshiro end # ========================================================================== -# 10. SDE + VR: solve rng overrides JumpProblem rng -# ========================================================================== -@testset "SDE + VR: solve rng overrides JumpProblem rng" begin - jprob = make_sde_vr_jump_prob(; rng = Xoshiro(42)) - integrator = init(jprob, EM(); dt = 0.01, rng = StableRNG(99)) - @test SciMLBase.get_rng(integrator) isa StableRNG -end - -# ========================================================================== -# 11. SDE + VR: reproducibility via JumpProblem rng +# 9. SDE + VR: reproducibility via solve rng # ========================================================================== -@testset "SDE + VR: JumpProblem rng reproducibility" begin - jprob1 = make_sde_vr_jump_prob(; rng = StableRNG(123)) - jprob2 = make_sde_vr_jump_prob(; rng = StableRNG(123)) - sol1 = solve(jprob1, EM(); dt = 0.01, save_everystep = false) - sol2 = solve(jprob2, EM(); dt = 0.01, save_everystep = false) +@testset "SDE + VR: solve rng reproducibility" begin + jprob = make_sde_vr_jump_prob() + sol1 = solve(jprob, EM(); dt = 0.01, save_everystep = false, rng = StableRNG(123)) + sol2 = solve(jprob, EM(); dt = 0.01, save_everystep = false, rng = StableRNG(123)) @test sol1.u[end] ≈ sol2.u[end] end # ========================================================================== -# 12. Tau-leaping: rng via JumpProblem +# 10. SimpleTauLeaping: rng via solve kwargs # ========================================================================== -@testset "SimpleTauLeaping: rng via JumpProblem kwargs" begin +@testset "SimpleTauLeaping: rng via solve kwargs" begin rate(out, u, p, t) = (out .= max.(u, 0); nothing) c(du, u, p, t, counts, mark) = (du .= counts; nothing) rj = RegularJump(rate, c, 2) dprob = DiscreteProblem([100, 100], (0.0, 1.0)) - jprob = JumpProblem(dprob, PureLeaping(), rj; rng = StableRNG(42)) - sol1 = solve(jprob, SimpleTauLeaping(); dt = 0.01) - jprob2 = JumpProblem(dprob, PureLeaping(), rj; rng = StableRNG(42)) - sol2 = solve(jprob2, SimpleTauLeaping(); dt = 0.01) + jprob = JumpProblem(dprob, PureLeaping(), rj) + sol1 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(42)) + sol2 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(42)) @test sol1.u == sol2.u end # ========================================================================== -# 13. Tau-leaping: solve rng overrides JumpProblem rng +# 11. SimpleTauLeaping: different seeds → different trajectories # ========================================================================== -@testset "SimpleTauLeaping: solve rng overrides JumpProblem rng" begin +@testset "SimpleTauLeaping: different seeds → different trajectories" begin rate(out, u, p, t) = (out .= max.(u, 0); nothing) c(du, u, p, t, counts, mark) = (du .= counts; nothing) rj = RegularJump(rate, c, 2) dprob = DiscreteProblem([100, 100], (0.0, 1.0)) - # JumpProblem has Xoshiro, solve has StableRNG - jprob = JumpProblem(dprob, PureLeaping(), rj; rng = Xoshiro(1)) + jprob = JumpProblem(dprob, PureLeaping(), rj) sol1 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(42)) - sol2 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(42)) - @test sol1.u == sol2.u - # Different from using the JumpProblem rng - jprob2 = JumpProblem(dprob, PureLeaping(), rj; rng = Xoshiro(1)) - sol3 = solve(jprob2, SimpleTauLeaping(); dt = 0.01) - @test sol1.u != sol3.u + sol2 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(99)) + @test sol1.u != sol2.u end # ========================================================================== -# 14. has_rng / get_rng / set_rng! interface on SSAIntegrator +# 12. has_rng / get_rng / set_rng! interface on SSAIntegrator # ========================================================================== @testset "SSAIntegrator RNG interface" begin jprob = make_ssa_jump_prob() @@ -202,7 +170,7 @@ end end # ========================================================================== -# 15. No rng kwarg: uses default_rng (non-reproducible but functional) +# 13. No rng kwarg: uses default_rng (non-reproducible but functional) # ========================================================================== @testset "No rng kwarg: functional solve" begin @testset "SSAStepper" begin @@ -217,3 +185,21 @@ end @test sol.retcode == ReturnCode.Success end end + +# ========================================================================== +# 14. seed kwarg: creates Xoshiro from integer seed +# ========================================================================== +@testset "seed kwarg creates Xoshiro" begin + jprob = make_ssa_jump_prob() + integrator = init(jprob, SSAStepper(); seed = 42) + @test SciMLBase.get_rng(integrator) isa Xoshiro +end + +# ========================================================================== +# 15. rng takes priority over seed +# ========================================================================== +@testset "rng takes priority over seed" begin + jprob = make_ssa_jump_prob() + integrator = init(jprob, SSAStepper(); rng = StableRNG(42), seed = 99) + @test SciMLBase.get_rng(integrator) isa StableRNG +end diff --git a/test/save_positions.jl b/test/save_positions.jl index 13413b5b2..5d0a683f2 100644 --- a/test/save_positions.jl +++ b/test/save_positions.jl @@ -15,8 +15,8 @@ let jump = VariableRateJump((u, p, t) -> 0, (integrator) -> integrator.u[1] += 1; urate = (u, p, t) -> 1.0, rateinterval = (u, p, t) -> 5.0) jumpproblem = JumpProblem(dprob, alg, jump; dep_graph = [[1]], - save_positions = (false, true), rng) - sol = solve(jumpproblem, SSAStepper()) + save_positions = (false, true)) + sol = solve(jumpproblem, SSAStepper(); rng) @test sol.t == [0.0, 30.0] oprob = ODEProblem((du, u, p, t) -> 0, u0, tspan) @@ -26,8 +26,8 @@ let for vr_agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) jumpproblem = JumpProblem( oprob, alg, jump; vr_aggregator = vr_agg, dep_graph = [[1]], - save_positions = (false, true), rng) - sol = solve(jumpproblem, Tsit5(); save_everystep = false) + save_positions = (false, true)) + sol = solve(jumpproblem, Tsit5(); rng, save_everystep = false) @test sol.t == [0.0, 30.0] end end @@ -46,20 +46,20 @@ let # for pure jump problems dense = save_everystep vals = (true, true, true, false) for (sp, val) in zip(sps, vals) - jprob = JumpProblem(dprob, Direct(), crj; save_positions = sp, rng) - sol = solve(jprob, SSAStepper()) + jprob = JumpProblem(dprob, Direct(), crj; save_positions = sp) + sol = solve(jprob, SSAStepper(); rng) @test SciMLBase.isdenseplot(sol) == val end # for mixed problems sol.dense currently ignores save_positions oprob = ODEProblem((du, u, p, t) -> du[1] = 0.1, u0, tspan) for sp in sps - jprob = JumpProblem(oprob, Direct(), crj; save_positions = sp, rng) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(oprob, Direct(), crj; save_positions = sp) + sol = solve(jprob, Tsit5(); rng) @test sol.dense == true @test SciMLBase.isdenseplot(sol) == true - sol = solve(jprob, Tsit5(); dense = false) + sol = solve(jprob, Tsit5(); rng, dense = false) @test sol.dense == false @test SciMLBase.isdenseplot(sol) == false end diff --git a/test/saveat_regression.jl b/test/saveat_regression.jl index 03665d76a..108fde97e 100644 --- a/test/saveat_regression.jl +++ b/test/saveat_regression.jl @@ -10,12 +10,12 @@ maj = MassActionJump(rate_consts, reactant_stoich, net_stoich) n0 = [1, 1, 0] tspan = (0, 0.2) dprob = DiscreteProblem(n0, tspan) -jprob = JumpProblem(dprob, Direct(), maj, save_positions = (false, false), rng = rng) +jprob = JumpProblem(dprob, Direct(), maj, save_positions = (false, false)) ts = collect(0:0.002:tspan[2]) NA = zeros(length(ts)) Nsims = 10_000 -sol = JumpProcesses.solve(EnsembleProblem(jprob), SSAStepper(), saveat = ts, - trajectories = Nsims) +sol = JumpProcesses.solve(EnsembleProblem(jprob), SSAStepper(); saveat = ts, + trajectories = Nsims, rng) for i in 1:length(sol) NA .+= sol.u[i][1, :] @@ -26,10 +26,10 @@ for i in 1:length(ts) end NA = zeros(length(ts)) -jprob = JumpProblem(dprob, Direct(), maj; rng = rng) +jprob = JumpProblem(dprob, Direct(), maj) sol = nothing; GC.gc(); -sol = JumpProcesses.solve(EnsembleProblem(jprob), SSAStepper(), trajectories = Nsims) +sol = JumpProcesses.solve(EnsembleProblem(jprob), SSAStepper(); trajectories = Nsims, rng) for i in 1:Nsims for n in 1:length(ts) diff --git a/test/sir_model.jl b/test/sir_model.jl index e8cea455e..8f4cc1ccc 100644 --- a/test/sir_model.jl +++ b/test/sir_model.jl @@ -18,8 +18,8 @@ end jump2 = ConstantRateJump(rate, affect!) prob = DiscreteProblem([999.0, 1.0, 0.0], (0.0, 250.0)) -jump_prob = JumpProblem(prob, Direct(), jump, jump2; rng = rng) -integrator = init(jump_prob, FunctionMap()) +jump_prob = JumpProblem(prob, Direct(), jump, jump2) +integrator = init(jump_prob, FunctionMap(); rng) condition(u, t, integrator) = t == 100 function purge_affect!(integrator) @@ -27,8 +27,8 @@ function purge_affect!(integrator) reset_aggregated_jumps!(integrator) end cb = DiscreteCallback(condition, purge_affect!, save_positions = (false, false)) -sol = solve(jump_prob, FunctionMap(), callback = cb, tstops = [100]) -sol = solve(jump_prob, SSAStepper(), callback = cb, tstops = [100]) +sol = solve(jump_prob, FunctionMap(); callback = cb, tstops = [100], rng) +sol = solve(jump_prob, SSAStepper(); callback = cb, tstops = [100], rng) # test README example using the auto-solver selection runs let diff --git a/test/spatial/ABC.jl b/test/spatial/ABC.jl index 8d7230a74..3935ffca9 100644 --- a/test/spatial/ABC.jl +++ b/test/spatial/ABC.jl @@ -38,10 +38,10 @@ prob = DiscreteProblem(starting_state, tspan, rates) hopping_constants = [hopping_rate for i in starting_state] # algs = [NSM(), DirectCRDirect()] -function get_mean_end_state(jump_prob, Nsims) +function get_mean_end_state(jump_prob, Nsims; rng = nothing) end_state = zeros(size(jump_prob.prob.u0)) for i in 1:Nsims - sol = solve(jump_prob, SSAStepper()) + sol = solve(jump_prob, SSAStepper(); rng) end_state .+= sol.u[end] end end_state / Nsims @@ -52,19 +52,19 @@ grids = [CartesianGridRej(dims), Graphs.grid(dims)] jump_problems = JumpProblem[JumpProblem(prob, NSM(), majumps, hopping_constants = hopping_constants, spatial_system = grid, - save_positions = (false, false), rng = rng) + save_positions = (false, false)) for grid in grids] push!(jump_problems, JumpProblem(prob, DirectCRDirect(), majumps, hopping_constants = hopping_constants, - spatial_system = grids[1], save_positions = (false, false), rng = rng)) + spatial_system = grids[1], save_positions = (false, false))) # setup flattenned jump prob push!(jump_problems, JumpProblem(prob, NRM(), majumps, hopping_constants = hopping_constants, - spatial_system = grids[1], save_positions = (false, false), rng = rng)) + spatial_system = grids[1], save_positions = (false, false))) # test for spatial_jump_prob in jump_problems - solution = solve(spatial_jump_prob, SSAStepper()) - mean_end_state = get_mean_end_state(spatial_jump_prob, Nsims) + solution = solve(spatial_jump_prob, SSAStepper(); rng) + mean_end_state = get_mean_end_state(spatial_jump_prob, Nsims; rng) mean_end_state = reshape(mean_end_state, num_species, num_nodes) diff = sum(mean_end_state, dims = 2) - non_spatial_mean for (i, d) in enumerate(diff) diff --git a/test/spatial/diffusion.jl b/test/spatial/diffusion.jl index c014b5c52..e501ab07f 100644 --- a/test/spatial/diffusion.jl +++ b/test/spatial/diffusion.jl @@ -5,10 +5,10 @@ using Graphs using StableRNGs rng = StableRNG(12345) -function get_mean_sol(jump_prob, Nsims, saveat) - sol = solve(jump_prob, SSAStepper(), saveat = saveat).u +function get_mean_sol(jump_prob, Nsims, saveat; rng = nothing) + sol = solve(jump_prob, SSAStepper(); saveat, rng).u for i in 1:(Nsims - 1) - sol += solve(jump_prob, SSAStepper(), saveat = saveat).u + sol += solve(jump_prob, SSAStepper(); saveat, rng).u end sol / Nsims end @@ -66,24 +66,24 @@ grids = [CartesianGridRej(dims), Graphs.grid(dims)] jump_problems = JumpProblem[JumpProblem(prob, algs[2], majumps, hopping_constants = hopping_constants, spatial_system = grid, - save_positions = (false, false), rng = rng) + save_positions = (false, false)) for grid in grids] sizehint!(jump_problems, 15 + length(jump_problems)) # flattenned jump prob push!(jump_problems, JumpProblem(prob, NRM(), majumps, hopping_constants = hopping_constants, - spatial_system = grids[1], save_positions = (false, false), rng = rng)) + spatial_system = grids[1], save_positions = (false, false))) # hop rates of form D_s hop_constants = [hopping_rate] for alg in algs push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants = hop_constants, - spatial_system = grids[1], save_positions = (false, false), rng = rng)) + spatial_system = grids[1], save_positions = (false, false))) push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants = hop_constants, - spatial_system = grids[end], save_positions = (false, false), rng = rng)) + spatial_system = grids[end], save_positions = (false, false))) end # hop rates of form L_{s,i,j} @@ -96,10 +96,10 @@ end for alg in algs push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants = hop_constants, - spatial_system = grids[1], save_positions = (false, false), rng = rng)) + spatial_system = grids[1], save_positions = (false, false))) push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants = hop_constants, - spatial_system = grids[end], save_positions = (false, false), rng = rng)) + spatial_system = grids[end], save_positions = (false, false))) end # hop rates of form D_s * L_{i,j} @@ -112,11 +112,11 @@ for alg in algs push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants = Pair(species_hop_constants, site_hop_constants), - spatial_system = grids[1], save_positions = (false, false), rng = rng)) + spatial_system = grids[1], save_positions = (false, false))) push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants = Pair(species_hop_constants, site_hop_constants), - spatial_system = grids[end], save_positions = (false, false), rng = rng)) + spatial_system = grids[end], save_positions = (false, false))) end # hop rates of form D_{s,i} * L_{i,j} @@ -129,16 +129,16 @@ for alg in algs push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants = Pair(species_hop_constants, site_hop_constants), - spatial_system = grids[1], save_positions = (false, false), rng = rng)) + spatial_system = grids[1], save_positions = (false, false))) push!(jump_problems, JumpProblem(prob, alg, majumps, hopping_constants = Pair(species_hop_constants, site_hop_constants), - spatial_system = grids[end], save_positions = (false, false), rng = rng)) + spatial_system = grids[end], save_positions = (false, false))) end # testing for (j, spatial_jump_prob) in enumerate(jump_problems) - mean_sol = get_mean_sol(spatial_jump_prob, Nsims, tf / num_time_points) + mean_sol = get_mean_sol(spatial_jump_prob, Nsims, tf / num_time_points; rng) for (i, t) in enumerate(times) local diff = analytic_solution(t) - reshape(mean_sol[i], num_nodes, 1) @test abs(sum(diff[1:center_node]) / sum(analytic_solution(t)[1:center_node])) < @@ -165,7 +165,7 @@ tspan = (0.0, 10.0) prob = DiscreteProblem(starting_state, tspan) jp = JumpProblem(prob, NSM(), majumps, hopping_constants = hopping_constants, - spatial_system = grid, save_positions = (false, false), rng = rng) -sol = solve(jp, SSAStepper()) + spatial_system = grid, save_positions = (false, false)) +sol = solve(jp, SSAStepper(); rng) @test sol.u[end][1, 1] == sum(sol.u[end]) diff --git a/test/spatial/spatial_majump.jl b/test/spatial/spatial_majump.jl index 6afcd367c..6018b68fd 100644 --- a/test/spatial/spatial_majump.jl +++ b/test/spatial/spatial_majump.jl @@ -61,26 +61,26 @@ non_uniform_majumps = [non_uniform_majumps_1, non_uniform_majumps_2, non_uniform uniform_jump_problems = JumpProblem[JumpProblem(prob, NSM(), majump, hopping_constants = hopping_constants, spatial_system = grid, - save_positions = (false, false), rng = rng) + save_positions = (false, false)) for majump in uniform_majumps] # flattenned append!(uniform_jump_problems, JumpProblem[JumpProblem(prob, NRM(), majump, hopping_constants = hopping_constants, - spatial_system = grid, save_positions = (false, false), rng = rng) + spatial_system = grid, save_positions = (false, false)) for majump in uniform_majumps]) # non-uniform non_uniform_jump_problems = JumpProblem[JumpProblem(prob, NSM(), majump, hopping_constants = hopping_constants, spatial_system = grid, - save_positions = (false, false), rng = rng) + save_positions = (false, false)) for majump in non_uniform_majumps] # testing -function get_mean_end_state(jump_prob, Nsims) +function get_mean_end_state(jump_prob, Nsims; rng = nothing) end_state = zeros(size(jump_prob.prob.u0)) for i in 1:Nsims - sol = solve(jump_prob, SSAStepper()) + sol = solve(jump_prob, SSAStepper(); rng) end_state .+= sol.u[end] end end_state / Nsims @@ -106,8 +106,8 @@ ode_prob = ODEProblem(f, zeros(num_nodes), tspan) sol = solve(ode_prob, Tsit5()) for spatial_jump_prob in uniform_jump_problems - solution = solve(spatial_jump_prob, SSAStepper()) - mean_end_state = get_mean_end_state(spatial_jump_prob, Nsims) + solution = solve(spatial_jump_prob, SSAStepper(); rng) + mean_end_state = get_mean_end_state(spatial_jump_prob, Nsims; rng) mean_end_state = reshape(mean_end_state, num_nodes) diff = mean_end_state - sol.u[end] for (i, d) in enumerate(diff) @@ -122,8 +122,8 @@ end ode_prob = ODEProblem(f2, zeros(num_nodes), tspan) sol = solve(ode_prob, Tsit5()) -solution = solve(non_uniform_jump_problems[1], SSAStepper()) -mean_end_state = get_mean_end_state(non_uniform_jump_problems[1], Nsims) +solution = solve(non_uniform_jump_problems[1], SSAStepper(); rng) +mean_end_state = get_mean_end_state(non_uniform_jump_problems[1], Nsims; rng) mean_end_state = reshape(mean_end_state, num_nodes) diff = mean_end_state - sol.u[end] for (i, d) in enumerate(diff) @@ -135,8 +135,8 @@ f3(u, p, t) = L * u - diagm([0.0, 0.0, death_rate, 0.0, 0.0]) * u + ones(num_nod ode_prob = ODEProblem(f3, zeros(num_nodes), tspan) sol = solve(ode_prob, Tsit5()) -solution = solve(non_uniform_jump_problems[2], SSAStepper()) -mean_end_state = get_mean_end_state(non_uniform_jump_problems[2], Nsims) +solution = solve(non_uniform_jump_problems[2], SSAStepper(); rng) +mean_end_state = get_mean_end_state(non_uniform_jump_problems[2], Nsims; rng) mean_end_state = reshape(mean_end_state, num_nodes) diff = mean_end_state - sol.u[end] for (i, d) in enumerate(diff) @@ -150,8 +150,8 @@ end ode_prob = ODEProblem(f4, zeros(num_nodes), tspan) sol = solve(ode_prob, Tsit5()) -solution = solve(non_uniform_jump_problems[3], SSAStepper()) -mean_end_state = get_mean_end_state(non_uniform_jump_problems[3], Nsims) +solution = solve(non_uniform_jump_problems[3], SSAStepper(); rng) +mean_end_state = get_mean_end_state(non_uniform_jump_problems[3], Nsims; rng) mean_end_state = reshape(mean_end_state, num_nodes) diff = mean_end_state - sol.u[end] for (i, d) in enumerate(diff) diff --git a/test/splitcoupled.jl b/test/splitcoupled.jl index e871e2c0f..1030ec88f 100644 --- a/test/splitcoupled.jl +++ b/test/splitcoupled.jl @@ -12,16 +12,15 @@ jump1 = ConstantRateJump(rate, affect!) prob = DiscreteProblem([10], (0.0, 50.0)) prob_control = DiscreteProblem([10], (0.0, 50.0)) -jump_prob = JumpProblem(prob, Direct(), jump1; rng = rng) -jump_prob_control = JumpProblem(prob_control, Direct(), jump1; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump1) +jump_prob_control = JumpProblem(prob_control, Direct(), jump1) coupling_map = [(1, 1)] coupled_prob = SplitCoupledJumpProblem( - jump_prob, jump_prob_control, Direct(), coupling_map; - rng = rng) + jump_prob, jump_prob_control, Direct(), coupling_map) -@time sol = solve(coupled_prob, FunctionMap()) -@time solve(jump_prob, FunctionMap()) +@time sol = solve(coupled_prob, FunctionMap(); rng) +@time solve(jump_prob, FunctionMap(); rng) @test [s[1] - s[2] for s in sol.u] == zeros(length(sol.t)) # coupling two copies of the same process should give zero rate = (u, p, t) -> 1.0 @@ -42,34 +41,31 @@ end # Jump ODE to jump ODE prob = ODEProblem(f, [1.0], (0.0, 1.0)) prob_control = ODEProblem(f, [1.0], (0.0, 1.0)) -jump_prob = JumpProblem(prob, Direct(), jump1; rng = rng) -jump_prob_control = JumpProblem(prob_control, Direct(), jump2; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump1) +jump_prob_control = JumpProblem(prob_control, Direct(), jump2) coupled_prob = SplitCoupledJumpProblem( - jump_prob, jump_prob_control, Direct(), coupling_map; - rng = rng) -sol = solve(coupled_prob, Tsit5()) + jump_prob, jump_prob_control, Direct(), coupling_map) +sol = solve(coupled_prob, Tsit5(); rng) @test mean([abs(s[1] - s[2]) for s in sol.u]) <= 5.0 # Jump SDE to Jump SDE prob = SDEProblem(f, g, [1.0], (0.0, 1.0)) prob_control = SDEProblem(f, g, [1.0], (0.0, 1.0)) -jump_prob = JumpProblem(prob, Direct(), jump1; rng = rng) -jump_prob_control = JumpProblem(prob_control, Direct(), jump1; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump1) +jump_prob_control = JumpProblem(prob_control, Direct(), jump1) coupled_prob = SplitCoupledJumpProblem( - jump_prob, jump_prob_control, Direct(), coupling_map; - rng = rng) -sol = solve(coupled_prob, SRIW1()) + jump_prob, jump_prob_control, Direct(), coupling_map) +sol = solve(coupled_prob, SRIW1(); rng) @test mean([abs(s[1] - s[2]) for s in sol.u]) <= 5.0 # Jump SDE to Jump ODE prob = ODEProblem(f, [1.0], (0.0, 1.0)) prob_control = SDEProblem(f, g, [1.0], (0.0, 1.0)) -jump_prob = JumpProblem(prob, Direct(), jump1; rng = rng) -jump_prob_control = JumpProblem(prob_control, Direct(), jump1; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump1) +jump_prob_control = JumpProblem(prob_control, Direct(), jump1) coupled_prob = SplitCoupledJumpProblem( - jump_prob, jump_prob_control, Direct(), coupling_map; - rng = rng) -sol = solve(coupled_prob, SRIW1()) + jump_prob, jump_prob_control, Direct(), coupling_map) +sol = solve(coupled_prob, SRIW1(); rng) @test mean([abs(s[1] - s[2]) for s in sol.u]) <= 5.0 # Jump SDE to Discrete @@ -79,12 +75,11 @@ affect! = function (integrator) end prob = DiscreteProblem([1.0], (0.0, 1.0)) prob_control = SDEProblem(f, g, [1.0], (0.0, 1.0)) -jump_prob = JumpProblem(prob, Direct(), jump1; rng = rng) -jump_prob_control = JumpProblem(prob_control, Direct(), jump1; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump1) +jump_prob_control = JumpProblem(prob_control, Direct(), jump1) coupled_prob = SplitCoupledJumpProblem( - jump_prob, jump_prob_control, Direct(), coupling_map; - rng = rng) -sol = solve(coupled_prob, SRIW1()) + jump_prob, jump_prob_control, Direct(), coupling_map) +sol = solve(coupled_prob, SRIW1(); rng) # test mass action jumps coupled to ODE # 0 -> A (stochasic) and A -> 0 (ODE) @@ -96,13 +91,12 @@ f = function (du, u, p, t) du[1] = -1.0 * u[1] end odeprob = ODEProblem(f, [10.0], (0.0, 10.0)) -jump_prob = JumpProblem(odeprob, Direct(), majumps, save_positions = (false, false); - rng = rng) +jump_prob = JumpProblem(odeprob, Direct(), majumps, save_positions = (false, false)) Nsims = 8000 Amean = 0.0 for i in 1:Nsims global Amean - local sol = solve(jump_prob, Tsit5(), saveat = 10.0) + local sol = solve(jump_prob, Tsit5(); saveat = 10.0, rng) Amean += sol[1, end] end Amean /= Nsims diff --git a/test/ssa_tests.jl b/test/ssa_tests.jl index e82c50959..79bdbe15f 100644 --- a/test/ssa_tests.jl +++ b/test/ssa_tests.jl @@ -16,31 +16,30 @@ end jump2 = ConstantRateJump(rate, affect!) prob = DiscreteProblem([10.0], (0.0, 3.0)) -jump_prob = JumpProblem(prob, Direct(), jump, jump2; rng = rng) +jump_prob = JumpProblem(prob, Direct(), jump, jump2) -integrator = init(jump_prob, SSAStepper()) +integrator = init(jump_prob, SSAStepper(); rng) step!(integrator) integrator.u[1] # test different saving behaviors -sol = solve(jump_prob, SSAStepper()) +sol = solve(jump_prob, SSAStepper(); rng) @test SciMLBase.successful_retcode(sol) @test sol.t[begin] == 0.0 @test sol.t[end] == 3.0 -sol = solve(jump_prob, SSAStepper(), save_end = false) +sol = solve(jump_prob, SSAStepper(); save_end = false, rng) @test sol.t[begin] == 0.0 @test sol.t[end] < 3.0 -sol = solve(jump_prob, SSAStepper(), save_start = false) +sol = solve(jump_prob, SSAStepper(); save_start = false, rng) @test sol.t[begin] > 0.0 @test sol.t[end] == 3.0 -jump_prob = JumpProblem(prob, Direct(), jump, jump2, save_positions = (false, false); - rng = rng) -sol = solve(jump_prob, SSAStepper(), save_start = false, save_end = false) +jump_prob = JumpProblem(prob, Direct(), jump, jump2, save_positions = (false, false)) +sol = solve(jump_prob, SSAStepper(); save_start = false, save_end = false, rng) @test isempty(sol.t) && isempty(sol.u) -sol = solve(jump_prob, SSAStepper(), saveat = 0.0:0.1:2.9) +sol = solve(jump_prob, SSAStepper(); saveat = 0.0:0.1:2.9, rng) @test sol.t == collect(0.0:0.1:3.0) diff --git a/test/thread_safety.jl b/test/thread_safety.jl index 197ed72c6..3a7c2043a 100644 --- a/test/thread_safety.jl +++ b/test/thread_safety.jl @@ -9,10 +9,10 @@ params = (1.0, 2.0, 50.0) tspan = (0.0, 4.0) u0 = [5] dprob = DiscreteProblem(u0, tspan, params) -jprob = JumpProblem(dprob, Direct(), maj; rng = rng) -solve(EnsembleProblem(jprob), SSAStepper(), EnsembleThreads(); trajectories = 10) +jprob = JumpProblem(dprob, Direct(), maj) +solve(EnsembleProblem(jprob), SSAStepper(), EnsembleThreads(); trajectories = 10, rng) solve(EnsembleProblem(jprob; safetycopy = true), SSAStepper(), EnsembleThreads(); - trajectories = 10) + trajectories = 10, rng) # test for https://github.com/SciML/JumpProcesses.jl/issues/472 let diff --git a/test/variable_rate.jl b/test/variable_rate.jl index 2669b70e2..63a6567f6 100644 --- a/test/variable_rate.jl +++ b/test/variable_rate.jl @@ -30,15 +30,15 @@ f = function (du, u, p, t) end prob = ODEProblem(f, [0.2], (0.0, 10.0)) -jump_prob = JumpProblem(prob, jump, jump2; vr_aggregator = VR_FRM(), rng) -integrator = init(jump_prob, Tsit5()) -sol = solve(jump_prob, Tsit5()) -sol = solve(jump_prob, Rosenbrock23(autodiff = AutoFiniteDiff())) -sol = solve(jump_prob, Rosenbrock23()) - -jump_prob_gill = JumpProblem(prob, jump, jump2; vr_aggregator = VR_Direct(), rng) -integrator = init(jump_prob_gill, Tsit5()) -sol_gill = solve(jump_prob_gill, Tsit5()) +jump_prob = JumpProblem(prob, jump, jump2; vr_aggregator = VR_FRM()) +integrator = init(jump_prob, Tsit5(); rng) +sol = solve(jump_prob, Tsit5(); rng) +sol = solve(jump_prob, Rosenbrock23(autodiff = AutoFiniteDiff()); rng) +sol = solve(jump_prob, Rosenbrock23(); rng) + +jump_prob_gill = JumpProblem(prob, jump, jump2; vr_aggregator = VR_Direct()) +integrator = init(jump_prob_gill, Tsit5(); rng) +sol_gill = solve(jump_prob_gill, Tsit5(); rng) sol_gill = solve(jump_prob, Rosenbrock23(autodiff = AutoFiniteDiff())) sol_gill = solve(jump_prob, Rosenbrock23()) @test maximum([sol.u[i][2] for i in 1:length(sol)]) <= 1e-12 @@ -48,10 +48,10 @@ g = function (du, u, p, t) du[1] = u[1] end prob = SDEProblem(f, g, [0.2], (0.0, 10.0)) -jump_prob = JumpProblem(prob, jump, jump2; vr_aggregator = VR_FRM(), rng = rng) -sol = solve(jump_prob, SRIW1()) -jump_prob_gill = JumpProblem(prob, jump, jump2; vr_aggregator = VR_Direct(), rng = rng) -sol_gill = solve(jump_prob_gill, SRIW1()) +jump_prob = JumpProblem(prob, jump, jump2; vr_aggregator = VR_FRM()) +sol = solve(jump_prob, SRIW1(); rng) +jump_prob_gill = JumpProblem(prob, jump, jump2; vr_aggregator = VR_Direct()) +sol_gill = solve(jump_prob_gill, SRIW1(); rng) @test maximum([sol.u[i][2] for i in 1:length(sol)]) <= 1e-12 @test maximum([sol.u[i][3] for i in 1:length(sol)]) <= 1e-12 @@ -74,10 +74,10 @@ function affect_switch!(integrator) end jump_switch = VariableRateJump(rate_switch, affect_switch!) prob = SDEProblem(ff, gg, ones(2), (0.0, 1.0), 0, noise_rate_prototype = zeros(2, 2)) -jump_prob = JumpProblem(prob, jump_switch; vr_aggregator = VR_FRM(), rng) -jump_prob_gill = JumpProblem(prob, jump_switch; vr_aggregator = VR_Direct(), rng) -sol = solve(jump_prob, SRA1(), dt = 1.0) -sol_gill = solve(jump_prob_gill, SRA1(), dt = 1.0) +jump_prob = JumpProblem(prob, jump_switch; vr_aggregator = VR_FRM()) +jump_prob_gill = JumpProblem(prob, jump_switch; vr_aggregator = VR_Direct()) +sol = solve(jump_prob, SRA1(), dt = 1.0; rng) +sol_gill = solve(jump_prob_gill, SRA1(), dt = 1.0; rng) ## Some integration tests @@ -88,10 +88,10 @@ prob = ODEProblem(f2, [0.2], (0.0, 10.0)) rate2(u, p, t) = 2 affect2!(integrator) = (integrator.u[1] = integrator.u[1] / 2) jump = ConstantRateJump(rate2, affect2!) -jump_prob = JumpProblem(prob, jump; vr_aggregator = VR_FRM(), rng) -jump_prob_gill = JumpProblem(prob, jump; vr_aggregator = VR_Direct(), rng) -sol = solve(jump_prob, Tsit5()) -sol_gill = solve(jump_prob_gill, Tsit5()) +jump_prob = JumpProblem(prob, jump; vr_aggregator = VR_FRM()) +jump_prob_gill = JumpProblem(prob, jump; vr_aggregator = VR_Direct()) +sol = solve(jump_prob, Tsit5(); rng) +sol_gill = solve(jump_prob_gill, Tsit5(); rng) sol(4.0) sol.u[4] @@ -99,10 +99,10 @@ rate2b(u, p, t) = u[1] affect2b!(integrator) = (integrator.u[1] = integrator.u[1] / 2) jump = VariableRateJump(rate2b, affect2b!) jump2 = deepcopy(jump) -jump_prob = JumpProblem(prob, jump, jump2; vr_aggregator = VR_FRM(), rng) -jump_prob_gill = JumpProblem(prob, jump, jump2; vr_aggregator = VR_Direct(), rng) -sol = solve(jump_prob, Tsit5()) -sol_gill = solve(jump_prob_gill, Tsit5()) +jump_prob = JumpProblem(prob, jump, jump2; vr_aggregator = VR_FRM()) +jump_prob_gill = JumpProblem(prob, jump, jump2; vr_aggregator = VR_Direct()) +sol = solve(jump_prob, Tsit5(); rng) +sol_gill = solve(jump_prob_gill, Tsit5(); rng) sol(4.0) sol.u[4] @@ -110,10 +110,10 @@ function g2(du, u, p, t) du[1] = u[1] end prob = SDEProblem(f2, g2, [0.2], (0.0, 10.0)) -jump_prob = JumpProblem(prob, jump, jump2; vr_aggregator = VR_FRM(), rng) -jump_prob_gill = JumpProblem(prob, jump, jump2; vr_aggregator = VR_Direct(), rng) -sol = solve(jump_prob, SRIW1()) -sol_gill = solve(jump_prob_gill, SRIW1()) +jump_prob = JumpProblem(prob, jump, jump2; vr_aggregator = VR_FRM()) +jump_prob_gill = JumpProblem(prob, jump, jump2; vr_aggregator = VR_Direct()) +sol = solve(jump_prob, SRIW1(); rng) +sol_gill = solve(jump_prob_gill, SRIW1(); rng) sol(4.0) sol.u[4] @@ -129,10 +129,10 @@ function affect3!(integrator) integrator.u[4] = 1) end jump = VariableRateJump(rate3, affect3!) -jump_prob = JumpProblem(prob, jump; vr_aggregator = VR_FRM(), rng) -jump_prob_gill = JumpProblem(prob, jump; vr_aggregator = VR_Direct(), rng) -sol = solve(jump_prob, Tsit5()) -sol_gill = solve(jump_prob_gill, Tsit5()) +jump_prob = JumpProblem(prob, jump; vr_aggregator = VR_FRM()) +jump_prob_gill = JumpProblem(prob, jump; vr_aggregator = VR_Direct()) +sol = solve(jump_prob, Tsit5(); rng) +sol_gill = solve(jump_prob_gill, Tsit5(); rng) # test for https://discourse.julialang.org/t/differentialequations-jl-package-variable-rate-jumps-with-complex-variables/80366/2 function f4(dx, x, p, t) @@ -146,10 +146,10 @@ jump = VariableRateJump(rate4, affect4!) x₀ = 1.0 + 0.0im Δt = (0.0, 6.0) prob = ODEProblem(f4, [x₀], Δt) -jump_prob = JumpProblem(prob, jump; vr_aggregator = VR_FRM(), rng) -jump_prob_gill = JumpProblem(prob, jump; vr_aggregator = VR_Direct(), rng) -sol = solve(jump_prob, Tsit5()) -sol_gill = solve(jump_prob_gill, Tsit5()) +jump_prob = JumpProblem(prob, jump; vr_aggregator = VR_FRM()) +jump_prob_gill = JumpProblem(prob, jump; vr_aggregator = VR_Direct()) +sol = solve(jump_prob, Tsit5(); rng) +sol_gill = solve(jump_prob_gill, Tsit5(); rng) # Out of place test drift(x, p, t) = p * x @@ -158,7 +158,7 @@ affect!2(integrator) = (integrator.u ./= 2; nothing) x0 = rand(2) prob = ODEProblem(drift, x0, (0.0, 10.0), 2.0) jump = VariableRateJump(rate2c, affect!2) -jump_prob = JumpProblem(prob, Direct(), jump; vr_aggregator = VR_FRM(), rng) +jump_prob = JumpProblem(prob, Direct(), jump; vr_aggregator = VR_FRM()) # test to check lack of dependency graphs is caught in Coevolve for systems with non-maj # jumps @@ -256,12 +256,12 @@ let d_jump = VariableRateJump(d_rate, death!) ode_prob = ODEProblem(ode_fxn, u0, tspan, p) - sjm_prob = JumpProblem(ode_prob, b_jump, d_jump; vr_aggregator = VR_FRM(), rng) + sjm_prob = JumpProblem(ode_prob, b_jump, d_jump; vr_aggregator = VR_FRM()) # After callback initialize, integrator.u.jump_u should have unique thresholds # that differ between sequential solves (RNG advances each time). jump_u_old = zeros(length(sjm_prob.prob.u0.jump_u)) for i in 1:Nsims - integrator = init(sjm_prob, Tsit5(); saveat = tspan[2]) + integrator = init(sjm_prob, Tsit5(); saveat = tspan[2], rng) @test allunique(integrator.u.jump_u) @test integrator.u.jump_u != jump_u_old jump_u_old .= integrator.u.jump_u @@ -339,10 +339,10 @@ end # Function to run ensemble and compute statistics function run_ensemble(prob, alg, jumps...; vr_aggregator = VR_FRM(), Nsims = 8000) rng = StableRNG(12345) - jump_prob = JumpProblem(prob, Direct(), jumps...; vr_aggregator, rng) + jump_prob = JumpProblem(prob, Direct(), jumps...; vr_aggregator) total = 0.0 for i in 1:Nsims - sol = solve(jump_prob, alg; save_everystep = false) + sol = solve(jump_prob, alg; save_everystep = false, rng) total += sol.u[end][1] end return total / Nsims @@ -438,10 +438,10 @@ let jump_counts = zeros(Int, Nsims) p = [0.0, 0.0, 0] prob = ODEProblem(f, u0, tspan, p) - jump_prob = JumpProblem(prob, Direct(), birth_jump, death_jump; vr_aggregator, rng) + jump_prob = JumpProblem(prob, Direct(), birth_jump, death_jump; vr_aggregator) for i in 1:Nsims - sol = solve(jump_prob, Tsit5(); save_everystep = false) + sol = solve(jump_prob, Tsit5(); save_everystep = false, rng) jump_counts[i] = jump_prob.prob.p[3] jump_prob.prob.p[3] = 0 end @@ -496,8 +496,8 @@ let ] for agg in aggregators - local jprob = JumpProblem(prob, agg, maj, vrj; rng) - local sol = solve(jprob, Tsit5()) + local jprob = JumpProblem(prob, agg, maj, vrj) + local sol = solve(jprob, Tsit5(); rng) @test SciMLBase.successful_retcode(sol) # Verify conservation: total population should be conserved @test sol.u[end][1] + sol.u[end][2] ≈ u0[1] + u0[2] @@ -527,8 +527,8 @@ let prob = ODEProblem(f!, u0, tspan) # Test with Direct aggregator (most common case) - jprob = JumpProblem(prob, Direct(), crj, vrj; rng) - sol = solve(jprob, Tsit5()) + jprob = JumpProblem(prob, Direct(), crj, vrj) + sol = solve(jprob, Tsit5(); rng) @test SciMLBase.successful_retcode(sol) @test sol.u[end][1] + sol.u[end][2] ≈ u0[1] + u0[2] end @@ -559,8 +559,8 @@ let # Test RSSA and RSSACR specifically (the aggregators that had the bug) for agg in [RSSA(), RSSACR()] - local jprob = JumpProblem(prob, agg, maj, vrj1, vrj2; rng) - local sol = solve(jprob, Tsit5()) + local jprob = JumpProblem(prob, agg, maj, vrj1, vrj2) + local sol = solve(jprob, Tsit5(); rng) @test SciMLBase.successful_retcode(sol) @test sol.u[end][1] + sol.u[end][2] ≈ u0[1] + u0[2] end From a7de080958638d377afc28c35b6e16d8eb175527 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 06:29:10 -0500 Subject: [PATCH 04/52] Clean up code quality: revert unnecessary diffs, add set_rng! type checking - Revert <:Any back to named type params in ssajump.jl (minimize diff from master) - Revert solkwargs to simpler ternary form in problem.jl (rng removal made filter unnecessary) - Add type checking to SSAIntegrator set_rng! to match StochasticDiffEq pattern Co-Authored-By: Claude Opus 4.5 --- src/SSA_stepper.jl | 8 ++++++++ src/aggregators/ssajump.jl | 8 ++++---- src/problem.jl | 4 ++-- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index c4f72906d..913bcd749 100644 --- a/src/SSA_stepper.jl +++ b/src/SSA_stepper.jl @@ -118,6 +118,14 @@ end SciMLBase.has_rng(::SSAIntegrator) = true SciMLBase.get_rng(integrator::SSAIntegrator) = integrator.rng function SciMLBase.set_rng!(integrator::SSAIntegrator, rng) + R = typeof(integrator.rng) + if !isa(rng, R) + throw(ArgumentError( + "Cannot set RNG of type $(typeof(rng)) on an integrator " * + "whose RNG type parameter is $R. " * + "Construct a new integrator via `init(prob, alg; rng = your_rng)` instead." + )) + end integrator.rng = rng nothing end diff --git a/src/aggregators/ssajump.jl b/src/aggregators/ssajump.jl index 0e72b2348..fec06c185 100644 --- a/src/aggregators/ssajump.jl +++ b/src/aggregators/ssajump.jl @@ -53,8 +53,8 @@ end nothing end -@inline function concretize_affects!(p::AbstractSSAJumpAggregator{<:Any, <:Any, <:Any, F2}, - ::I) where {F2 <: Tuple, I <: SciMLBase.DEIntegrator} +@inline function concretize_affects!(p::AbstractSSAJumpAggregator{T, S, F1, F2}, + ::I) where {T, S, F1, F2 <: Tuple, I <: SciMLBase.DEIntegrator} nothing end @@ -86,8 +86,8 @@ function (p::AbstractSSAJumpAggregator)(integrator::I) where {I <: SciMLBase.DEI end function (p::AbstractSSAJumpAggregator{ - <:Any, <:Any, <:Any, F2})(integrator::SciMLBase.DEIntegrator) where - {F2 <: Union{Tuple, Nothing}} + T, S, F1, F2})(integrator::SciMLBase.DEIntegrator) where + {T, S, F1, F2 <: Union{Tuple, Nothing}} execute_jumps!(p, integrator, integrator.u, integrator.p, integrator.t, p.affects!) generate_jumps!(p, integrator, integrator.u, integrator.p, integrator.t) register_next_jump_time!(integrator, p, integrator.t) diff --git a/src/problem.jl b/src/problem.jl index 718179f7b..3ce038de5 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -309,7 +309,7 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS jump_cbs = CallbackSet(constant_jump_callback, variable_jump_callback) iip = isinplace_jump(prob, jumps.regular_jump) - solkwargs = make_kwarg(; filter(!isnothing, (; callback, tstops))...) + solkwargs = tstops === nothing ? make_kwarg(; callback) : make_kwarg(; callback, tstops) JumpProblem{iip, typeof(new_prob), typeof(aggregator), typeof(jump_cbs), typeof(disc_agg), typeof(crjs), typeof(cvrjs), typeof(jumps.regular_jump), @@ -355,7 +355,7 @@ function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; vrjs = jumps.variable_jumps iip = isinplace_jump(prob, jumps.regular_jump) - solkwargs = make_kwarg(; filter(!isnothing, (; callback, tstops))...) + solkwargs = tstops === nothing ? make_kwarg(; callback) : make_kwarg(; callback, tstops) JumpProblem{iip, typeof(prob), typeof(aggregator), typeof(jump_cbs), typeof(disc_agg), typeof(crjs), typeof(vrjs), typeof(jumps.regular_jump), From 0ded5246c803ba755afcbc229134be561e4fa173 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 06:38:07 -0500 Subject: [PATCH 05/52] Add tests for set_rng! type mismatch error on SSAIntegrator Co-Authored-By: Claude Opus 4.5 --- test/ensemble_problems.jl | 3 +++ test/rng_kwarg_tests.jl | 3 +++ 2 files changed, 6 insertions(+) diff --git a/test/ensemble_problems.jl b/test/ensemble_problems.jl index 7cfc7066f..27e446504 100644 --- a/test/ensemble_problems.jl +++ b/test/ensemble_problems.jl @@ -181,6 +181,9 @@ end new_rng = StableRNG(99) SciMLBase.set_rng!(integrator, new_rng) @test SciMLBase.get_rng(integrator) === new_rng + + # mismatched RNG type should throw + @test_throws ArgumentError SciMLBase.set_rng!(integrator, Random.Xoshiro(123)) end # ========================================================================== diff --git a/test/rng_kwarg_tests.jl b/test/rng_kwarg_tests.jl index a677a9118..010b96cef 100644 --- a/test/rng_kwarg_tests.jl +++ b/test/rng_kwarg_tests.jl @@ -167,6 +167,9 @@ end new_rng = StableRNG(99) SciMLBase.set_rng!(integrator, new_rng) @test SciMLBase.get_rng(integrator) === new_rng + + # mismatched RNG type should throw + @test_throws ArgumentError SciMLBase.set_rng!(integrator, Random.Xoshiro(123)) end # ========================================================================== From 20376a9ece795299d021d9f67b8f9d8b1e58c6f5 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 06:45:31 -0500 Subject: [PATCH 06/52] Remove unused StableRNG import and variable from GPU test The GPU tau-leaping kernel uses PassthroughRNG internally; the StableRNG was never consumed by the EnsembleGPUKernel path. Co-Authored-By: Claude Opus 4.5 --- test/gpu/regular_jumps.jl | 3 --- 1 file changed, 3 deletions(-) diff --git a/test/gpu/regular_jumps.jl b/test/gpu/regular_jumps.jl index b35276cf4..bb136dd18 100644 --- a/test/gpu/regular_jumps.jl +++ b/test/gpu/regular_jumps.jl @@ -1,9 +1,6 @@ using JumpProcesses, DiffEqBase using Test, LinearAlgebra, Statistics using KernelAbstractions, Adapt, CUDA -using StableRNGs -rng = StableRNG(12345) - Nsims = 100_000 # SIR model with influx From 0c32cceea42d1b7636f5270be5d8023525b381d2 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 06:59:52 -0500 Subject: [PATCH 07/52] Simplify test solve calls: remove unnecessary isnothing(rng) checks solve(; rng = nothing) falls through to default_rng() via resolve_rng, so the ternary guards are redundant. Co-Authored-By: Claude Opus 4.5 --- test/bimolerx_test.jl | 4 ++-- test/geneexpr_test.jl | 3 +-- test/reversible_binding.jl | 2 +- 3 files changed, 4 insertions(+), 5 deletions(-) diff --git a/test/bimolerx_test.jl b/test/bimolerx_test.jl index 832968b77..05917bf5a 100644 --- a/test/bimolerx_test.jl +++ b/test/bimolerx_test.jl @@ -59,9 +59,9 @@ function runSSAs(jump_prob; use_stepper = true, rng = nothing) Psamp = zeros(Int, Nsims) for i in 1:Nsims sol = if use_stepper - isnothing(rng) ? solve(jump_prob, SSAStepper()) : solve(jump_prob, SSAStepper(); rng) + solve(jump_prob, SSAStepper(); rng) else - isnothing(rng) ? solve(jump_prob) : solve(jump_prob; rng) + solve(jump_prob; rng) end Psamp[i] = sol[1, end] end diff --git a/test/geneexpr_test.jl b/test/geneexpr_test.jl index d26d9b8a2..5212e0710 100644 --- a/test/geneexpr_test.jl +++ b/test/geneexpr_test.jl @@ -26,10 +26,9 @@ function runSSAs(jump_prob; use_stepper = true, rng = nothing) Psamp = zeros(Int, Nsims) for i in 1:Nsims sol = if use_stepper - isnothing(rng) ? solve(jump_prob, SSAStepper()) : solve(jump_prob, SSAStepper(); rng) else - isnothing(rng) ? solve(jump_prob) : solve(jump_prob; rng) + solve(jump_prob; rng) end Psamp[i] = sol[3, end] end diff --git a/test/reversible_binding.jl b/test/reversible_binding.jl index 456a96d2a..d0e9b9a6d 100644 --- a/test/reversible_binding.jl +++ b/test/reversible_binding.jl @@ -23,7 +23,7 @@ majumps = MassActionJump(rates, reactstoch, netstoch) function getmean(jprob, Nsims; rng = nothing) Amean = 0 for i in 1:Nsims - sol = isnothing(rng) ? solve(jprob, SSAStepper()) : solve(jprob, SSAStepper(); rng) + sol = solve(jprob, SSAStepper(); rng) Amean += sol[1, end] end Amean /= Nsims From 33b86e9cb520947309e01890a4ce74b6d78eb95e Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 07:14:08 -0500 Subject: [PATCH 08/52] Use jump times instead of final state for trajectory uniqueness tests Comparing continuous-valued jump times is more robust than discrete final state values, which could collide by chance. Co-Authored-By: Claude Opus 4.5 --- test/ensemble_problems.jl | 8 ++------ test/rng_kwarg_tests.jl | 4 ++-- 2 files changed, 4 insertions(+), 8 deletions(-) diff --git a/test/ensemble_problems.jl b/test/ensemble_problems.jl index 27e446504..61c82e313 100644 --- a/test/ensemble_problems.jl +++ b/test/ensemble_problems.jl @@ -54,8 +54,6 @@ first_jump_time(traj) = traj.t[2] trajectories = 3, rng = StableRNG(12345)) times = [first_jump_time(sol.u[i]) for i in 1:3] @test allunique(times) - finals = [sol.u[i].u[end][1] for i in 1:3] - @test allunique(finals) end # EM() uses a fixed time grid so jump event times aren't directly visible @@ -87,8 +85,6 @@ end sols = [solve(jprob, Tsit5(); rng) for _ in 1:3] times = [first_jump_time(s) for s in sols] @test allunique(times) - finals = [s.u[end][1] for s in sols] - @test allunique(finals) end end @@ -161,8 +157,8 @@ end @testset "ODE + VR ($agg): different seeds → different trajectories" for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) jprob = make_vr_jump_prob(agg) sols = [solve(jprob, Tsit5(); rng = StableRNG(s)) for s in (100, 200, 300)] - finals = [s.u[end][1] for s in sols] - @test allunique(finals) + times = [first_jump_time(s) for s in sols] + @test allunique(times) end end diff --git a/test/rng_kwarg_tests.jl b/test/rng_kwarg_tests.jl index 010b96cef..09f6ab1ee 100644 --- a/test/rng_kwarg_tests.jl +++ b/test/rng_kwarg_tests.jl @@ -102,8 +102,8 @@ end @testset "ODE + VR: different seeds → different trajectories" begin jprob = make_ode_vr_jump_prob() sols = [solve(jprob, Tsit5(); rng = StableRNG(s)) for s in (100, 200, 300)] - finals = [s.u[end][1] for s in sols] - @test allunique(finals) + times = [s.t[2] for s in sols] + @test allunique(times) end # ========================================================================== From 2b6d3c6287279efdb66388f60a9bd2de0d3d4421 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 07:54:40 -0500 Subject: [PATCH 09/52] Strengthen VR_FRM threshold uniqueness test Check jump_u thresholds directly via init, and verify both event times and post-event thresholds from ensemble solve. Co-Authored-By: Claude Opus 4.5 --- test/ensemble_problems.jl | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/test/ensemble_problems.jl b/test/ensemble_problems.jl index 61c82e313..f08ea524c 100644 --- a/test/ensemble_problems.jl +++ b/test/ensemble_problems.jl @@ -187,18 +187,29 @@ end # # For VR_FRM, each trajectory's first jump time is determined by the initial # jump_u threshold (set to -randexp() by the VR_FRMEventCallback initialize). -# Distinct thresholds → distinct first event times, so we verify by checking -# that the second time point (first event) differs across serial trajectories. +# We verify both the thresholds (via init) and the resulting event times. # ========================================================================== @testset "VR_FRM: jump_u thresholds unique per trajectory (EnsembleSerial)" begin jprob = make_vr_jump_prob(VR_FRM()) + + # Check jump_u thresholds directly via init (callback sets them during initialization) + rng = StableRNG(12345) + thresholds = [begin + integrator = init(jprob, Tsit5(); rng) + integrator.u.jump_u[1] + end for _ in 1:3] + @test allunique(thresholds) + + # From a full ensemble solve, check both first event times and the + # post-event jump_u thresholds (u[3] is the post-event save where + # jump_u has been reset to a new -randexp() value). sol = solve(EnsembleProblem(jprob), Tsit5(), EnsembleSerial(); trajectories = 3, rng = StableRNG(12345)) - # The second time point is when the first variable-rate jump fires, - # directly reflecting the initial -randexp() threshold. event_times = [sol.u[i].t[2] for i in 1:3] @test allunique(event_times) + post_event_thresholds = [sol.u[i].u[3].jump_u[1] for i in 1:3] + @test allunique(post_event_thresholds) end # ========================================================================== From 44ced847b40dc25d07821139a750ea1cfab6fcdd Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 09:33:08 -0500 Subject: [PATCH 10/52] Move EnsembleThreads tests to thread_safety.jl, fix Project.toml - Remove StochasticDiffEq from [deps] (test-only, belongs in [extras]) - Move EnsembleThreads data race tests from ensemble_problems.jl to thread_safety.jl with stronger assertions (400 trajectories, allunique on first reaction times, all 3 VR aggregators, adaptive SDE solver) - Renumber remaining ensemble_problems.jl sections Co-Authored-By: Claude Opus 4.5 --- Project.toml | 1 - test/ensemble_problems.jl | 43 ++++----------------------------------- test/thread_safety.jl | 21 ++++++++++++++++++- 3 files changed, 24 insertions(+), 41 deletions(-) diff --git a/Project.toml b/Project.toml index 3a528ceb5..0d264333b 100644 --- a/Project.toml +++ b/Project.toml @@ -18,7 +18,6 @@ RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" -StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" [weakdeps] diff --git a/test/ensemble_problems.jl b/test/ensemble_problems.jl index aa619eabe..efff24f6a 100644 --- a/test/ensemble_problems.jl +++ b/test/ensemble_problems.jl @@ -90,42 +90,7 @@ end end # ========================================================================== -# 3. Threaded ensemble: no data race on the shared JumpProblem -# -# With integrator-owned RNGs, each thread's integrator gets its own -# default_rng(). We only assert completion here — uniqueness is tested -# via explicit rng kwarg in section 4. -# ========================================================================== - -@testset "EnsembleThreads: no data race" begin - @testset "SSAStepper" begin - jprob = make_ssa_jump_prob() - sol = solve(EnsembleProblem(jprob), SSAStepper(), EnsembleThreads(); - trajectories = 4) - @test length(sol) == 4 - end - - @testset "ODE + VR ($agg)" for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) - jprob = make_vr_jump_prob(agg) - sol = solve(EnsembleProblem(jprob), Tsit5(), EnsembleThreads(); - trajectories = 4, save_everystep = false) - @test length(sol) == 4 - end - - @testset "SDE + VR (VR_FRM): unique trajectories" begin - jprob = make_sde_vr_jump_prob(VR_FRM()) - # StochasticDiffEq generates per-trajectory seeds, so trajectories - # should be distinct. - sol = solve(EnsembleProblem(jprob), EM(), EnsembleThreads(); - trajectories = 4, dt = 0.01, save_everystep = false) - @test length(sol) == 4 - finals = [sol.u[i].u[end][1] for i in 1:4] - @test length(unique(finals)) > 1 - end -end - -# ========================================================================== -# 4. rng kwarg reproducibility: same rng seed → identical trajectory, +# 3. rng kwarg reproducibility: same rng seed → identical trajectory, # different rng seeds → different trajectories # ========================================================================== @@ -164,7 +129,7 @@ end end # ========================================================================== -# 5. has_rng / get_rng / set_rng! interface on SSAIntegrator +# 4. has_rng / get_rng / set_rng! interface on SSAIntegrator # ========================================================================== @testset "SSAIntegrator RNG interface" begin @@ -184,7 +149,7 @@ end end # ========================================================================== -# 6. Variable-rate: jump_u thresholds are unique per trajectory +# 5. Variable-rate: jump_u thresholds are unique per trajectory # # For VR_FRM, each trajectory's first jump time is determined by the initial # jump_u threshold (set to -randexp() by the VR_FRMEventCallback initialize). @@ -214,7 +179,7 @@ end end # ========================================================================== -# 7. JumpProblem rng kwarg forwarded to solver +# 6. JumpProblem rng kwarg throws ArgumentError # ========================================================================== @testset "JumpProblem rng kwarg throws ArgumentError" begin diff --git a/test/thread_safety.jl b/test/thread_safety.jl index 349346db4..2f80d5e87 100644 --- a/test/thread_safety.jl +++ b/test/thread_safety.jl @@ -1,5 +1,5 @@ using DiffEqBase, Test -using JumpProcesses, OrdinaryDiffEq +using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq using StableRNGs rng = StableRNG(12345) @@ -34,3 +34,22 @@ let @test allunique(firstrx_time) end end + +# SDE + variable-rate jumps with EnsembleThreads +let + f!(du, u, p, t) = (du[1] = -0.1 * u[1]; nothing) + g!(du, u, p, t) = (du[1] = 0.1 * u[1]; nothing) + sde_prob = SDEProblem(f!, g!, [100.0], (0.0, 10.0)) + vrj = VariableRateJump((u, p, t) -> 0.5 * u[1], + integrator -> (integrator.u[1] -= 1.0)) + + for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) + jump_prob = JumpProblem(sde_prob, Direct(), vrj; vr_aggregator = agg) + prob_func(prob, i, repeat) = deepcopy(prob) + prob = EnsembleProblem(jump_prob; prob_func) + sol = solve(prob, SRIW1(), EnsembleThreads(); + trajectories = 400, save_everystep = false) + firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol)] + @test allunique(firstrx_time) + end +end From 55e130b71d306a7a40742187cdb3e6f476021e47 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 11:54:25 -0500 Subject: [PATCH 11/52] Address PR review: strengthen tests, add seed coverage MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ensemble_problems.jl: check initial jump_u threshold (u[2]) instead of post-event (u[3]) now that u_modified!(true) produces init save - extended_jump_array_remake.jl: rework to solve → remake → init → compare jump_u, verifying fresh thresholds after remake - geneexpr_test.jl: remove unnecessary isnothing(rng) branch - linearreaction_test.jl: make rng const for type stability - variable_rate.jl: switch getmean from seed integer to StableRNG - rng_kwarg_tests.jl: add SDE+VR different seeds test, SDE+VR no-rng functional test, seed kwarg reproducibility and different-seeds tests for all pathways (SSAStepper, ODE+VR, SDE+VR), use first_jump_time helper for robust event time extraction Co-Authored-By: Claude Opus 4.5 --- test/ensemble_problems.jl | 8 +-- test/extended_jump_array_remake.jl | 46 +++++++----- test/geneexpr_test.jl | 6 +- test/linearreaction_test.jl | 2 +- test/rng_kwarg_tests.jl | 111 ++++++++++++++++++++++++++--- test/variable_rate.jl | 10 +-- 6 files changed, 139 insertions(+), 44 deletions(-) diff --git a/test/ensemble_problems.jl b/test/ensemble_problems.jl index efff24f6a..da11b321c 100644 --- a/test/ensemble_problems.jl +++ b/test/ensemble_problems.jl @@ -168,14 +168,14 @@ end @test allunique(thresholds) # From a full ensemble solve, check both first event times and the - # post-event jump_u thresholds (u[3] is the post-event save where - # jump_u has been reset to a new -randexp() value). + # initial jump_u thresholds (u[2] is the initialization save where + # jump_u has been set to -randexp() by the callback). sol = solve(EnsembleProblem(jprob), Tsit5(), EnsembleSerial(); trajectories = 3, rng = StableRNG(12345)) event_times = [first_jump_time(sol.u[i]) for i in 1:3] @test allunique(event_times) - post_event_thresholds = [sol.u[i].u[3].jump_u[1] for i in 1:3] - @test allunique(post_event_thresholds) + init_thresholds = [sol.u[i].u[2].jump_u[1] for i in 1:3] + @test allunique(init_thresholds) end # ========================================================================== diff --git a/test/extended_jump_array_remake.jl b/test/extended_jump_array_remake.jl index 224064ae1..7a32c64d7 100644 --- a/test/extended_jump_array_remake.jl +++ b/test/extended_jump_array_remake.jl @@ -5,8 +5,6 @@ using JumpProcesses, OrdinaryDiffEq, Test, SymbolicIndexingInterface using StableRNGs @testset "remake JumpProblem with VariableRateJumps (ExtendedJumpArray)" begin - rng = StableRNG(12345) - # Setup: Create an ODEProblem with SymbolCache for symbolic indexing f(du, u, p, t) = (du .= 0; nothing) g = ODEFunction(f; sys = SymbolCache([:X, :Y], [:k1, :k2], :t)) @@ -23,15 +21,19 @@ using StableRNGs @test jprob.prob.u0 isa ExtendedJumpArray @test jprob.prob.u0.u == [10.0, 5.0] + # Solve original problem and capture jump_u after initialization + orig_integrator = init(jprob, Tsit5(); rng = StableRNG(42)) + orig_jump_u = copy(orig_integrator.u.jump_u) + @testset "remake with numeric Vector{Float64}" begin prob2 = remake(jprob; u0 = [20.0, 10.0]) @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u == [20.0, 10.0] - # jump_u is zeroed at construction; callback initializes it at solve time @test all(iszero, prob2.prob.u0.jump_u) - # After init the callback should set jump_u to non-zero thresholds - integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + # After init, callback sets fresh jump_u thresholds (different RNG seed) + integrator = init(prob2, Tsit5(); rng = StableRNG(99)) @test any(!iszero, integrator.u.jump_u) + @test integrator.u.jump_u != orig_jump_u end @testset "remake with ExtendedJumpArray (no resample)" begin @@ -48,34 +50,34 @@ using StableRNGs end @testset "remake with Symbol pairs" begin - # This was the FAILING case - should work after fix prob2 = remake(jprob; u0 = [:X => 25.0]) @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u[1] == 25.0 - # jump_u zeroed at construction, initialized by callback at solve time @test all(iszero, prob2.prob.u0.jump_u) - integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + # After init, callback sets fresh jump_u thresholds (different RNG seed) + integrator = init(prob2, Tsit5(); rng = StableRNG(99)) @test any(!iszero, integrator.u.jump_u) + @test integrator.u.jump_u != orig_jump_u end @testset "remake with multiple Symbol pairs" begin prob2 = remake(jprob; u0 = [:X => 35.0, :Y => 15.0]) @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u == [35.0, 15.0] - # jump_u zeroed at construction, initialized by callback at solve time @test all(iszero, prob2.prob.u0.jump_u) - integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + integrator = init(prob2, Tsit5(); rng = StableRNG(99)) @test any(!iszero, integrator.u.jump_u) + @test integrator.u.jump_u != orig_jump_u end @testset "remake with Dict" begin prob2 = remake(jprob; u0 = Dict(:X => 40.0)) @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u[1] == 40.0 - # jump_u zeroed at construction, initialized by callback at solve time @test all(iszero, prob2.prob.u0.jump_u) - integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + integrator = init(prob2, Tsit5(); rng = StableRNG(99)) @test any(!iszero, integrator.u.jump_u) + @test integrator.u.jump_u != orig_jump_u end @testset "remake with parameters only (u0 unchanged)" begin @@ -93,21 +95,27 @@ using StableRNGs @test prob2.prob.u0 isa ExtendedJumpArray @test prob2.prob.u0.u[1] == 50.0 @test prob2.prob.p[1] == 3.0 - # jump_u zeroed at construction, initialized by callback at solve time @test all(iszero, prob2.prob.u0.jump_u) - integrator = init(prob2, Tsit5(); rng = StableRNG(42)) + integrator = init(prob2, Tsit5(); rng = StableRNG(99)) @test any(!iszero, integrator.u.jump_u) + @test integrator.u.jump_u != orig_jump_u end @testset "remake preserves problem solvability" begin - # Ensure remade problems can actually be solved + # Solve original, then remake and solve again — jump_u should differ + sol1 = solve(jprob, Tsit5(); rng = StableRNG(42)) + @test SciMLBase.successful_retcode(sol1) + prob2 = remake(jprob; u0 = [5.0, 2.0]) - sol = solve(prob2, Tsit5()) - @test SciMLBase.successful_retcode(sol) + sol2 = solve(prob2, Tsit5(); rng = StableRNG(99)) + @test SciMLBase.successful_retcode(sol2) + # Different RNG seeds → different jump_u thresholds after init + @test sol1.u[2].jump_u != sol2.u[2].jump_u - # With symbolic map (after fix) + # With symbolic map prob3 = remake(jprob; u0 = [:X => 8.0]) - sol3 = solve(prob3, Tsit5()) + sol3 = solve(prob3, Tsit5(); rng = StableRNG(77)) @test SciMLBase.successful_retcode(sol3) + @test sol1.u[2].jump_u != sol3.u[2].jump_u end end diff --git a/test/geneexpr_test.jl b/test/geneexpr_test.jl index 5212e0710..45e9a02b8 100644 --- a/test/geneexpr_test.jl +++ b/test/geneexpr_test.jl @@ -38,11 +38,7 @@ end function runSSAs_ode(vrjprob; rng = nothing) Psamp = zeros(Float64, Nsims) tsave = vrjprob.prob.tspan[2] - integrator = if isnothing(rng) - init(vrjprob, Tsit5(); saveat = tsave) - else - init(vrjprob, Tsit5(); saveat = tsave, rng) - end + integrator = init(vrjprob, Tsit5(); saveat = tsave, rng) solve!(integrator) Psamp[1] = integrator.sol[3, end] for i in 2:Nsims diff --git a/test/linearreaction_test.jl b/test/linearreaction_test.jl index 67e13b159..7883883c3 100644 --- a/test/linearreaction_test.jl +++ b/test/linearreaction_test.jl @@ -3,7 +3,7 @@ using DiffEqBase, JumpProcesses, Statistics using Test using StableRNGs -rng = StableRNG(12345) +const rng = StableRNG(12345) # using BenchmarkTools # dobenchmark = true diff --git a/test/rng_kwarg_tests.jl b/test/rng_kwarg_tests.jl index 09f6ab1ee..f3284f3bf 100644 --- a/test/rng_kwarg_tests.jl +++ b/test/rng_kwarg_tests.jl @@ -33,6 +33,9 @@ function make_sde_vr_jump_prob() JumpProblem(sprob, Direct(), vrj) end +# Helpers +first_jump_time(traj) = traj.t[findfirst(>(traj.t[1]), traj.t)] + # ========================================================================== # 1. JumpProblem(; rng=...) throws ArgumentError # ========================================================================== @@ -102,7 +105,7 @@ end @testset "ODE + VR: different seeds → different trajectories" begin jprob = make_ode_vr_jump_prob() sols = [solve(jprob, Tsit5(); rng = StableRNG(s)) for s in (100, 200, 300)] - times = [s.t[2] for s in sols] + times = [first_jump_time(s) for s in sols] @test allunique(times) end @@ -126,7 +129,18 @@ end end # ========================================================================== -# 10. SimpleTauLeaping: rng via solve kwargs +# 10. SDE + VR: different seeds → different trajectories +# ========================================================================== +@testset "SDE + VR: different seeds → different trajectories" begin + jprob = make_sde_vr_jump_prob() + sols = [solve(jprob, SRIW1(); save_everystep = false, + rng = StableRNG(s)) for s in (100, 200, 300)] + times = [first_jump_time(s) for s in sols] + @test allunique(times) +end + +# ========================================================================== +# 11. SimpleTauLeaping: rng via solve kwargs # ========================================================================== @testset "SimpleTauLeaping: rng via solve kwargs" begin rate(out, u, p, t) = (out .= max.(u, 0); nothing) @@ -140,7 +154,7 @@ end end # ========================================================================== -# 11. SimpleTauLeaping: different seeds → different trajectories +# 12. SimpleTauLeaping: different seeds → different trajectories # ========================================================================== @testset "SimpleTauLeaping: different seeds → different trajectories" begin rate(out, u, p, t) = (out .= max.(u, 0); nothing) @@ -154,7 +168,7 @@ end end # ========================================================================== -# 12. has_rng / get_rng / set_rng! interface on SSAIntegrator +# 13. has_rng / get_rng / set_rng! interface on SSAIntegrator # ========================================================================== @testset "SSAIntegrator RNG interface" begin jprob = make_ssa_jump_prob() @@ -173,7 +187,7 @@ end end # ========================================================================== -# 13. No rng kwarg: uses default_rng (non-reproducible but functional) +# 14. No rng kwarg: uses default_rng (non-reproducible but functional) # ========================================================================== @testset "No rng kwarg: functional solve" begin @testset "SSAStepper" begin @@ -187,19 +201,96 @@ end sol = solve(jprob, Tsit5()) @test sol.retcode == ReturnCode.Success end + + @testset "SDE + VR" begin + jprob = make_sde_vr_jump_prob() + sol = solve(jprob, EM(); dt = 0.01) + @test sol.retcode == ReturnCode.Success + end end # ========================================================================== -# 14. seed kwarg: creates Xoshiro from integer seed +# 15. seed kwarg: creates Xoshiro from integer seed # ========================================================================== @testset "seed kwarg creates Xoshiro" begin - jprob = make_ssa_jump_prob() - integrator = init(jprob, SSAStepper(); seed = 42) - @test SciMLBase.get_rng(integrator) isa Xoshiro + @testset "SSAStepper" begin + jprob = make_ssa_jump_prob() + integrator = init(jprob, SSAStepper(); seed = 42) + @test SciMLBase.get_rng(integrator) isa Xoshiro + end + + @testset "ODE + VR" begin + jprob = make_ode_vr_jump_prob() + integrator = init(jprob, Tsit5(); seed = 42) + @test SciMLBase.get_rng(integrator) isa Xoshiro + end + + @testset "SDE + VR" begin + jprob = make_sde_vr_jump_prob() + integrator = init(jprob, EM(); dt = 0.01, seed = 42) + @test SciMLBase.get_rng(integrator) isa Xoshiro + end +end + +# ========================================================================== +# 16. seed kwarg reproducibility: same seed → same trajectory +# ========================================================================== +@testset "seed kwarg reproducibility" begin + @testset "SSAStepper" begin + jprob = make_ssa_jump_prob() + sol1 = solve(jprob, SSAStepper(); seed = 42) + sol2 = solve(jprob, SSAStepper(); seed = 42) + @test sol1.t == sol2.t + @test sol1.u == sol2.u + end + + @testset "ODE + VR" begin + jprob = make_ode_vr_jump_prob() + sol1 = solve(jprob, Tsit5(); seed = 42) + sol2 = solve(jprob, Tsit5(); seed = 42) + @test sol1.t ≈ sol2.t + @test sol1.u[end] ≈ sol2.u[end] + end + + @testset "SDE + VR" begin + jprob = make_sde_vr_jump_prob() + sol1 = solve(jprob, EM(); dt = 0.01, save_everystep = false, seed = 42) + sol2 = solve(jprob, EM(); dt = 0.01, save_everystep = false, seed = 42) + @test sol1.u[end] ≈ sol2.u[end] + end +end + +# ========================================================================== +# 17. seed kwarg: different seeds → different trajectories +# ========================================================================== +@testset "seed kwarg: different seeds → different trajectories" begin + @testset "SSAStepper" begin + jprob = make_ssa_jump_prob() + sol1 = solve(jprob, SSAStepper(); seed = 100) + sol2 = solve(jprob, SSAStepper(); seed = 200) + sol3 = solve(jprob, SSAStepper(); seed = 300) + times = [sol1.t[2], sol2.t[2], sol3.t[2]] + @test allunique(times) + end + + @testset "ODE + VR" begin + jprob = make_ode_vr_jump_prob() + sols = [solve(jprob, Tsit5(); seed = s) for s in (100, 200, 300)] + times = [first_jump_time(s) for s in sols] + @test allunique(times) + end + + @testset "SDE + VR" begin + jprob = make_sde_vr_jump_prob() + sols = [solve(jprob, SRIW1(); save_everystep = false, + seed = s) for s in (100, 200, 300)] + times = [first_jump_time(s) for s in sols] + @test allunique(times) + end end # ========================================================================== -# 15. rng takes priority over seed +# 18. rng takes priority over seed # ========================================================================== @testset "rng takes priority over seed" begin jprob = make_ssa_jump_prob() diff --git a/test/variable_rate.jl b/test/variable_rate.jl index 63a6567f6..195479ec1 100644 --- a/test/variable_rate.jl +++ b/test/variable_rate.jl @@ -273,9 +273,9 @@ end # https://github.com/SciML/JumpProcesses.jl/issues/320 # note that even with the seeded StableRNG this test is not # deterministic for some reason. -function getmean(Nsims, prob, alg, tsave, seed) +function getmean(Nsims, prob, alg, tsave, rng) umean = zeros(length(tsave)) - integrator = init(prob, alg; saveat = tsave, seed) + integrator = init(prob, alg; saveat = tsave, rng) solve!(integrator) for j in eachindex(umean) umean[j] += integrator.sol.u[j][1] @@ -292,7 +292,7 @@ function getmean(Nsims, prob, alg, tsave, seed) end let - seed = 12345 + rng_seed = 12345 b = 2.0 d = 1.0 n0 = 1.0 @@ -327,9 +327,9 @@ let sjm_prob = JumpProblem(ode_prob, b_jump, d_jump; vr_aggregator) for alg in (Tsit5(), Rodas5P(linsolve = QRFactorization())) - umean = getmean(Nsims, sjm_prob, alg, tsave, seed) + umean = getmean(Nsims, sjm_prob, alg, tsave, StableRNG(rng_seed)) @test all(abs.(umean .- n.(tsave)) .< 0.05 * n.(tsave)) - seed += Nsims + rng_seed += Nsims end end end From 96864366137e68c4406fade5f85af108c9f8ba32 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 14:21:14 -0500 Subject: [PATCH 12/52] Address review round 3: fix docs example, add tau-leaping seed tests, strengthen thread_safety assertions - docs/advanced_point_process.md: move rng from JumpProblem to solve - rng_kwarg_tests.jl: add seed kwarg tests for SimpleTauLeaping (tests 13-14), renumber 15-20 - thread_safety.jl: SSAStepper EnsembleThreads now runs 400 trajectories and asserts allunique first jump times Co-Authored-By: Claude Opus 4.5 --- .../applications/advanced_point_process.md | 4 +- test/rng_kwarg_tests.jl | 40 ++++++++++++++++--- test/thread_safety.jl | 16 ++++++-- 3 files changed, 49 insertions(+), 11 deletions(-) diff --git a/docs/src/applications/advanced_point_process.md b/docs/src/applications/advanced_point_process.md index 976770017..87e397d86 100644 --- a/docs/src/applications/advanced_point_process.md +++ b/docs/src/applications/advanced_point_process.md @@ -387,10 +387,10 @@ function Base.rand(rng::AbstractRNG, out = Array{History, 1}(undef, n) p = params(pp) dprob = DiscreteProblem([0], tspan, p) - jprob = JumpProblem(dprob, Coevolve(), jumps...; dep_graph = pp.g, save_positions, rng) + jprob = JumpProblem(dprob, Coevolve(), jumps...; dep_graph = pp.g, save_positions) for i in 1:n params!(pp, p) - solve(jprob, SSAStepper()) + solve(jprob, SSAStepper(); rng) out[i] = deepcopy(p.h) end return out diff --git a/test/rng_kwarg_tests.jl b/test/rng_kwarg_tests.jl index f3284f3bf..a237cecdc 100644 --- a/test/rng_kwarg_tests.jl +++ b/test/rng_kwarg_tests.jl @@ -168,7 +168,35 @@ end end # ========================================================================== -# 13. has_rng / get_rng / set_rng! interface on SSAIntegrator +# 13. SimpleTauLeaping: seed kwarg reproducibility +# ========================================================================== +@testset "SimpleTauLeaping: seed kwarg reproducibility" begin + rate(out, u, p, t) = (out .= max.(u, 0); nothing) + c(du, u, p, t, counts, mark) = (du .= counts; nothing) + rj = RegularJump(rate, c, 2) + dprob = DiscreteProblem([100, 100], (0.0, 1.0)) + jprob = JumpProblem(dprob, PureLeaping(), rj) + sol1 = solve(jprob, SimpleTauLeaping(); dt = 0.01, seed = 42) + sol2 = solve(jprob, SimpleTauLeaping(); dt = 0.01, seed = 42) + @test sol1.u == sol2.u +end + +# ========================================================================== +# 14. SimpleTauLeaping: seed different seeds → different trajectories +# ========================================================================== +@testset "SimpleTauLeaping: seed different seeds → different trajectories" begin + rate(out, u, p, t) = (out .= max.(u, 0); nothing) + c(du, u, p, t, counts, mark) = (du .= counts; nothing) + rj = RegularJump(rate, c, 2) + dprob = DiscreteProblem([100, 100], (0.0, 1.0)) + jprob = JumpProblem(dprob, PureLeaping(), rj) + sol1 = solve(jprob, SimpleTauLeaping(); dt = 0.01, seed = 42) + sol2 = solve(jprob, SimpleTauLeaping(); dt = 0.01, seed = 99) + @test sol1.u != sol2.u +end + +# ========================================================================== +# 15. has_rng / get_rng / set_rng! interface on SSAIntegrator # ========================================================================== @testset "SSAIntegrator RNG interface" begin jprob = make_ssa_jump_prob() @@ -187,7 +215,7 @@ end end # ========================================================================== -# 14. No rng kwarg: uses default_rng (non-reproducible but functional) +# 16. No rng kwarg: uses default_rng (non-reproducible but functional) # ========================================================================== @testset "No rng kwarg: functional solve" begin @testset "SSAStepper" begin @@ -210,7 +238,7 @@ end end # ========================================================================== -# 15. seed kwarg: creates Xoshiro from integer seed +# 17. seed kwarg: creates Xoshiro from integer seed # ========================================================================== @testset "seed kwarg creates Xoshiro" begin @testset "SSAStepper" begin @@ -233,7 +261,7 @@ end end # ========================================================================== -# 16. seed kwarg reproducibility: same seed → same trajectory +# 18. seed kwarg reproducibility: same seed → same trajectory # ========================================================================== @testset "seed kwarg reproducibility" begin @testset "SSAStepper" begin @@ -261,7 +289,7 @@ end end # ========================================================================== -# 17. seed kwarg: different seeds → different trajectories +# 19. seed kwarg: different seeds → different trajectories # ========================================================================== @testset "seed kwarg: different seeds → different trajectories" begin @testset "SSAStepper" begin @@ -290,7 +318,7 @@ end end # ========================================================================== -# 18. rng takes priority over seed +# 20. rng takes priority over seed # ========================================================================== @testset "rng takes priority over seed" begin jprob = make_ssa_jump_prob() diff --git a/test/thread_safety.jl b/test/thread_safety.jl index 2f80d5e87..4c4cfcfdc 100644 --- a/test/thread_safety.jl +++ b/test/thread_safety.jl @@ -10,9 +10,19 @@ tspan = (0.0, 4.0) u0 = [5] dprob = DiscreteProblem(u0, tspan, params) jprob = JumpProblem(dprob, Direct(), maj) -solve(EnsembleProblem(jprob), SSAStepper(), EnsembleThreads(); trajectories = 10, rng) -solve(EnsembleProblem(jprob; safetycopy = true), SSAStepper(), EnsembleThreads(); - trajectories = 10, rng) + +# Verify threaded solves complete and produce distinct trajectories +sol = solve(EnsembleProblem(jprob), SSAStepper(), EnsembleThreads(); + trajectories = 400, rng) +@test length(sol) == 400 +firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol)] +@test allunique(firstrx_time) + +sol2 = solve(EnsembleProblem(jprob; safetycopy = true), SSAStepper(), EnsembleThreads(); + trajectories = 400, rng) +@test length(sol2) == 400 +firstrx_time2 = [sol2.u[i].t[findfirst(>(sol2.u[i].t[1]), sol2.u[i].t)] for i in 1:length(sol2)] +@test allunique(firstrx_time2) # test for https://github.com/SciML/JumpProcesses.jl/issues/472 let From ec3b6610227a1cd4f78c1b6f30f62f64a0da6d0a Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 24 Feb 2026 14:41:49 -0500 Subject: [PATCH 13/52] Add dedicated threaded CI for thread_safety tests - Create ThreadSafety.yml workflow: runs thread_safety.jl with --threads=4 - Add ThreadSafety group to Tests.yml matrix (serial) and runtests.jl - Move thread_safety.jl from InterfaceII to ThreadSafety group - Fix shared rng data race: SSAStepper ensemble tests now use task-local default_rng() instead of a shared StableRNG across threads Co-Authored-By: Claude Opus 4.5 --- .github/workflows/Tests.yml | 1 + .github/workflows/ThreadSafety.yml | 65 ++++++++++++++++++++++++++++++ test/runtests.jl | 5 ++- test/thread_safety.jl | 13 +++--- 4 files changed, 78 insertions(+), 6 deletions(-) create mode 100644 .github/workflows/ThreadSafety.yml diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index 3f223b3de..dd8ce04f5 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -29,6 +29,7 @@ jobs: group: - InterfaceI - InterfaceII + - ThreadSafety - QA exclude: - version: "pre" diff --git a/.github/workflows/ThreadSafety.yml b/.github/workflows/ThreadSafety.yml new file mode 100644 index 000000000..b82eda2db --- /dev/null +++ b/.github/workflows/ThreadSafety.yml @@ -0,0 +1,65 @@ +name: "Thread Safety Tests" + +on: + push: + branches: + - master + paths-ignore: + - 'docs/**' + pull_request: + branches: + - master + paths-ignore: + - 'docs/**' + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: ${{ github.ref_name != github.event.repository.default_branch || github.ref != 'refs/tags/v*' }} + +jobs: + thread-safety: + name: "Thread Safety" + strategy: + fail-fast: false + matrix: + version: + - "1" + - "lts" + - "pre" + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + + - name: "Setup Julia ${{ matrix.version }}" + uses: julia-actions/setup-julia@v2 + with: + version: "${{ matrix.version }}" + + - name: "Cache Julia packages" + uses: julia-actions/cache@v2 + with: + token: "${{ secrets.GITHUB_TOKEN }}" + + - name: "Build package" + uses: julia-actions/julia-buildpkg@v1 + + - name: "Run thread safety tests (4 threads)" + run: | + julia --threads=4 --code-coverage=user --check-bounds=yes --compiled-modules=yes \ + --project=@. --color=yes -e ' + using Pkg + Pkg.test() + ' + env: + GROUP: ThreadSafety + + - name: "Process Coverage" + uses: julia-actions/julia-processcoverage@v1 + + - name: "Report Coverage" + uses: codecov/codecov-action@v5 + if: ${{ github.event.pull_request.head.repo.full_name == github.repository }} + with: + files: lcov.info + token: "${{ secrets.CODECOV_TOKEN }}" + fail_ci_if_error: false diff --git a/test/runtests.jl b/test/runtests.jl index be0236c3b..e4b097c41 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -41,7 +41,6 @@ end @time @safetestset "Save_positions test" begin include("save_positions.jl") end @time @safetestset "RNG kwarg tests" begin include("rng_kwarg_tests.jl") end @time @safetestset "Ensemble Uniqueness test" begin include("ensemble_uniqueness.jl") end - @time @safetestset "Thread Safety test" begin include("thread_safety.jl") end @time @safetestset "Ensemble Problem Tests" begin include("ensemble_problems.jl") end @time @safetestset "A + B <--> C" begin include("reversible_binding.jl") end @time @safetestset "Remake tests" begin include("remake_test.jl") end @@ -58,6 +57,10 @@ end @time @safetestset "Pure diffusion" begin include("spatial/diffusion.jl") end end + if GROUP == "All" || GROUP == "ThreadSafety" + @time @safetestset "Thread Safety test (threaded)" begin include("thread_safety.jl") end + end + if GROUP == "CUDA" activate_gpu_env() @time @safetestset "GPU Tau Leaping test" begin include("gpu/regular_jumps.jl") end diff --git a/test/thread_safety.jl b/test/thread_safety.jl index 4c4cfcfdc..9f7c44918 100644 --- a/test/thread_safety.jl +++ b/test/thread_safety.jl @@ -1,7 +1,5 @@ using DiffEqBase, Test using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq -using StableRNGs -rng = StableRNG(12345) sr = [1.0, 2.0, 50.0] maj = MassActionJump(sr, [[1 => 1], [1 => 1], [0 => 1]], [[1 => 1], [1 => -1], [1 => 1]]) @@ -11,15 +9,20 @@ u0 = [5] dprob = DiscreteProblem(u0, tspan, params) jprob = JumpProblem(dprob, Direct(), maj) -# Verify threaded solves complete and produce distinct trajectories +# Verify threaded solves complete and produce distinct trajectories. +# NOTE: We intentionally do NOT pass `rng` here. In threaded ensembles, passing a +# shared rng object via `solve(...; rng=...)` does not yet provide correct +# per-trajectory stream handling. Until SciMLBase's ensemble RNG updates land +# (master rng -> per-trajectory rng), correctness in threaded contexts relies on +# task-local `Random.default_rng()`. sol = solve(EnsembleProblem(jprob), SSAStepper(), EnsembleThreads(); - trajectories = 400, rng) + trajectories = 400) @test length(sol) == 400 firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol)] @test allunique(firstrx_time) sol2 = solve(EnsembleProblem(jprob; safetycopy = true), SSAStepper(), EnsembleThreads(); - trajectories = 400, rng) + trajectories = 400) @test length(sol2) == 400 firstrx_time2 = [sol2.u[i].t[findfirst(>(sol2.u[i].t[1]), sol2.u[i].t)] for i in 1:length(sol2)] @test allunique(firstrx_time2) From 201e70e8410d4a2cf480ecbe2551845eae3e0213 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Wed, 25 Feb 2026 13:03:12 -0500 Subject: [PATCH 14/52] Add supports_solve_rng trait for JumpProblem solve paths Co-Authored-By: Claude Opus 4.5 --- src/SSA_stepper.jl | 1 + src/simple_regular_solve.jl | 3 +++ src/solve.jl | 6 ++++++ 3 files changed, 10 insertions(+) diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index 913bcd749..a3e4096ef 100644 --- a/src/SSA_stepper.jl +++ b/src/SSA_stepper.jl @@ -61,6 +61,7 @@ for details. """ struct SSAStepper <: DiffEqBase.DEAlgorithm end SciMLBase.allows_late_binding_tstops(::SSAStepper) = true +SciMLBase.supports_solve_rng(::JumpProblem, ::SSAStepper) = true """ $(TYPEDEF) diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index 7c8df9fb3..f71351efc 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -6,6 +6,9 @@ end SimpleExplicitTauLeaping(; epsilon = 0.05) = SimpleExplicitTauLeaping(epsilon) +SciMLBase.supports_solve_rng(::JumpProblem, ::SimpleTauLeaping) = true +SciMLBase.supports_solve_rng(::JumpProblem, ::SimpleExplicitTauLeaping) = true + 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 \ diff --git a/src/solve.jl b/src/solve.jl index db58e3e46..ffe3177dc 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -15,6 +15,9 @@ function resolve_rng(rng, seed) end end +SciMLBase.supports_solve_rng(jprob::JumpProblem, alg::DiffEqBase.DEAlgorithm) = + SciMLBase.supports_solve_rng(jprob.prob, alg) + function DiffEqBase.__solve(jump_prob::DiffEqBase.AbstractJumpProblem{P}, alg::DiffEqBase.DEAlgorithm; merge_callbacks = true, kwargs...) where {P} @@ -38,6 +41,9 @@ function DiffEqBase.__solve(jump_prob::DiffEqBase.AbstractJumpProblem{P}, integrator.sol end +SciMLBase.supports_solve_rng(jprob::JumpProblem, ::Nothing) = + jprob.prob isa DiffEqBase.DiscreteProblem + # if passed a JumpProblem over a DiscreteProblem, and no aggregator is selected use # SSAStepper function DiffEqBase.__solve(jump_prob::DiffEqBase.AbstractJumpProblem{P}; From 613611e505b39b621737337030d71aac7578c9ce Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 28 Feb 2026 07:24:12 -0500 Subject: [PATCH 15/52] Update package versions in Project.toml --- Project.toml | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index 0d264333b..35ea0c114 100644 --- a/Project.toml +++ b/Project.toml @@ -45,18 +45,18 @@ KernelAbstractions = "0.9" LinearAlgebra = "1" LinearSolve = "3" OrdinaryDiffEq = "6" -OrdinaryDiffEqCore = "3.9" +OrdinaryDiffEqCore = "3.11" Pkg = "1" PoissonRandom = "0.4" Random = "1" RecursiveArrayTools = "3.35" Reexport = "1.2" SafeTestsets = "0.1" -SciMLBase = "2.144" +SciMLBase = "2.147" StableRNGs = "1" StaticArrays = "1.9.8" Statistics = "1" -StochasticDiffEq = "6.82" +StochasticDiffEq = "6.95" SymbolicIndexingInterface = "0.3.36" Test = "1" julia = "1.10" From a434959ece71f37fe1393737453af3dcf026b6b1 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Thu, 26 Feb 2026 17:00:48 -0500 Subject: [PATCH 16/52] Add rescale_rates_on_update field to MassActionJump to prevent double-scaling MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit When MTKBase/Catalyst constructs MassActionJumps, the symbolic rate expressions already include stoichiometric scaling (e.g. k/3! for 3X → Y), so they pass scale_rates=false. However, update_parameters! defaulted to scale_rates=true, causing every subsequent parameter update path (reset_aggregated_jumps!, remake, finalize_parameters_hook!) to double-scale the rates. The fix stores the intended scaling behavior on the struct as rescale_rates_on_update, which update_parameters! now reads as its default. This field is propagated through all merge/combine paths and validated for consistency when merging multiple MassActionJumps. Also fixes SpatialMassActionJump conversion constructor to default scale_rates=false since it receives already-scaled rates. Co-Authored-By: Claude Opus 4.5 --- src/jumps.jl | 94 +++++++++------ src/spatial/spatial_massaction_jump.jl | 4 +- test/runtests.jl | 1 + test/scale_rates_field_test.jl | 156 +++++++++++++++++++++++++ 4 files changed, 216 insertions(+), 39 deletions(-) create mode 100644 test/scale_rates_field_test.jl diff --git a/src/jumps.jl b/src/jumps.jl index 63704ef92..07185d954 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -333,10 +333,13 @@ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump net_stoch::U """Parameter mapping functor to identify reaction rate constants with parameters in `p` vectors.""" param_mapper::V + """Whether `update_parameters!` should apply stoichiometric scaling to rates.""" + rescale_rates_on_update::Bool function MassActionJump{T, S, U, V}(rates::T, rs_in::S, ns::U, pmapper::V, scale_rates::Bool, useiszero::Bool, - nocopy::Bool) where {T <: AbstractVector, S, U, V} + nocopy::Bool, + rescale_rates_on_update::Bool = scale_rates) where {T <: AbstractVector, S, U, V} sr = nocopy ? rates : copy(rates) rs = nocopy ? rs_in : copy(rs_in) for i in eachindex(rs) @@ -348,14 +351,15 @@ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump if scale_rates && !isempty(sr) scalerates!(sr, rs) end - new(sr, rs, ns, pmapper) + new(sr, rs, ns, pmapper, rescale_rates_on_update) end function MassActionJump{Nothing, Vector{S}, Vector{U}, V}(::Nothing, rs_in::Vector{S}, ns::Vector{U}, pmapper::V, scale_rates::Bool, useiszero::Bool, - nocopy::Bool) where {S <: AbstractVector, + nocopy::Bool, + rescale_rates_on_update::Bool = scale_rates) where {S <: AbstractVector, U <: AbstractVector, V} rs = nocopy ? rs_in : copy(rs_in) for i in eachindex(rs) @@ -363,46 +367,51 @@ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump rs[i] = typeof(rs[i])() end end - new(nothing, rs, ns, pmapper) + new(nothing, rs, ns, pmapper, rescale_rates_on_update) end function MassActionJump{T, S, U, V}(rate::T, rs_in::S, ns::U, pmapper::V, scale_rates::Bool, useiszero::Bool, - nocopy::Bool) where {T <: Number, S, U, V} + nocopy::Bool, + rescale_rates_on_update::Bool = scale_rates) where {T <: Number, S, U, V} rs = rs_in if useiszero && (length(rs) == 1) && iszero(rs[1][1]) rs = typeof(rs)() end sr = scale_rates ? scalerate(rate, rs) : rate - new(sr, rs, ns, pmapper) + new(sr, rs, ns, pmapper, rescale_rates_on_update) end function MassActionJump{Nothing, S, U, V}(::Nothing, rs_in::S, ns::U, pmapper::V, scale_rates::Bool, useiszero::Bool, - nocopy::Bool) where {S, U, V} + nocopy::Bool, + rescale_rates_on_update::Bool = scale_rates) where {S, U, V} rs = rs_in if useiszero && (length(rs) == 1) && iszero(rs[1][1]) rs = typeof(rs)() end - new(nothing, rs, ns, pmapper) + new(nothing, rs, ns, pmapper, rescale_rates_on_update) end end function MassActionJump(usr::T, rs::S, ns::U, pmapper::V; scale_rates = true, - useiszero = true, nocopy = false) where {T, S, U, V} - MassActionJump{T, S, U, V}(usr, rs, ns, pmapper, scale_rates, useiszero, nocopy) + useiszero = true, nocopy = false, + rescale_rates_on_update = scale_rates) where {T, S, U, V} + MassActionJump{T, S, U, V}(usr, rs, ns, pmapper, scale_rates, useiszero, nocopy, + rescale_rates_on_update) end function MassActionJump(usr::T, rs, ns; scale_rates = true, useiszero = true, - nocopy = false) where {T <: AbstractVector} - MassActionJump(usr, rs, ns, nothing; scale_rates = scale_rates, useiszero = useiszero, - nocopy = nocopy) + nocopy = false, rescale_rates_on_update = scale_rates) where {T <: AbstractVector} + MassActionJump(usr, rs, ns, nothing; scale_rates, useiszero, nocopy, + rescale_rates_on_update) end function MassActionJump(usr::T, rs, ns; scale_rates = true, useiszero = true, - nocopy = false) where {T <: Number} - MassActionJump(usr, rs, ns, nothing; scale_rates = scale_rates, useiszero = useiszero, - nocopy = nocopy) + nocopy = false, rescale_rates_on_update = scale_rates) where {T <: Number} + MassActionJump(usr, rs, ns, nothing; scale_rates, useiszero, nocopy, + rescale_rates_on_update) end # with parameter indices or mapping, multiple jump case function MassActionJump(rs, ns; param_idxs = nothing, param_mapper = nothing, - scale_rates = true, useiszero = true, nocopy = false) + scale_rates = true, useiszero = true, nocopy = false, + rescale_rates_on_update = scale_rates) if param_mapper === nothing (param_idxs === nothing) && error("If no parameter indices are given via param_idxs, an explicit parameter mapping must be passed in via param_mapper.") @@ -413,8 +422,8 @@ function MassActionJump(rs, ns; param_idxs = nothing, param_mapper = nothing, pmapper = param_mapper end - MassActionJump(nothing, nocopy ? rs : copy(rs), ns, pmapper; scale_rates = scale_rates, - useiszero = useiszero, nocopy = true) + MassActionJump(nothing, nocopy ? rs : copy(rs), ns, pmapper; scale_rates, useiszero, + nocopy = true, rescale_rates_on_update) end using_params(maj::MassActionJump{T, S, U, Nothing}) where {T, S, U} = false @@ -470,18 +479,19 @@ function Base.merge(pmap1::MassActionJumpParamMapper{Int}, end """ -update_parameters!(maj::MassActionJump, newparams; scale_rates=true) +update_parameters!(maj::MassActionJump, newparams; scale_rates=maj.rescale_rates_on_update) Updates the passed in MassActionJump with the parameter values in `newparams`. Notes: - Requires the jump to have been constructed with a user-passed `param_idxs` or `param_mapper`. - - `scale_rates=true` will scale the parameter representing the jump rate by an - appropriate combinatoric factor. i.e for 3A --> B at rate k it will scale - k --> k/3!. + - `scale_rates` defaults to `maj.rescale_rates_on_update`, which itself defaults to the + `scale_rates` value used when constructing the jump. When `true`, the parameter + representing the jump rate will be scaled by an appropriate combinatoric factor, i.e. + for 3A --> B at rate k it will scale k --> k/3!. """ -function update_parameters!(maj::MassActionJump, newparams; scale_rates = true, kwargs...) +function update_parameters!(maj::MassActionJump, newparams; scale_rates = maj.rescale_rates_on_update, kwargs...) (maj.param_mapper === nothing) && error("MassActionJumps must be constructed with param_idxs or a param_mapper to be updateable.") maj.param_mapper(maj, newparams; scale_rates, kwargs) @@ -556,8 +566,12 @@ function JumpSet(vjs, cjs, rj, majv::Vector{T}) where {T <: MassActionJump} error("JumpSets do not accept empty mass action jump collections; use \"nothing\" instead.") end + sr_val = majv[1].rescale_rates_on_update + if !all(m -> m.rescale_rates_on_update == sr_val, majv) + error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") + end maj = setup_majump_to_merge(majv[1].scaled_rates, majv[1].reactant_stoch, - majv[1].net_stoch, majv[1].param_mapper) + majv[1].net_stoch, majv[1].param_mapper, sr_val) for i in 2:length(majv) massaction_jump_combine(maj, majv[i]) end @@ -618,35 +632,35 @@ end # functionality to merge two mass action jumps together function check_majump_type(maj::MassActionJump{S, T, U, V}) where {S <: Number, T, U, V} setup_majump_to_merge(maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, - maj.param_mapper) + maj.param_mapper, maj.rescale_rates_on_update) end function check_majump_type(maj::MassActionJump{Nothing, T, U, V}) where {T, U, V} setup_majump_to_merge(maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, - maj.param_mapper) + maj.param_mapper, maj.rescale_rates_on_update) end # if given containers of rates and stoichiometry directly create a jump -function setup_majump_to_merge(sr::T, rs::AbstractVector{S}, ns::AbstractVector{U}, - pmapper) where {T <: AbstractVector, S <: AbstractArray, +function setup_majump_to_merge(sr::T, rs::AbstractVector{S}, ns::AbstractVector{U}, pmapper, + rescale_rates_on_update::Bool) where {T <: AbstractVector, S <: AbstractArray, U <: AbstractArray} - MassActionJump(sr, rs, ns, pmapper; scale_rates = false) + MassActionJump(sr, rs, ns, pmapper; scale_rates = false, rescale_rates_on_update) end # if just given the data for one jump (and not in a container) wrap in a vector -function setup_majump_to_merge(sr::S, rs::T, ns::U, - pmapper) where {S <: Number, T <: AbstractArray, +function setup_majump_to_merge(sr::S, rs::T, ns::U, pmapper, + rescale_rates_on_update::Bool) where {S <: Number, T <: AbstractArray, U <: AbstractArray} MassActionJump([sr], [rs], [ns], (pmapper === nothing) ? pmapper : to_collection(pmapper); - scale_rates = false) + scale_rates = false, rescale_rates_on_update) end # if no rate field setup yet -function setup_majump_to_merge(::Nothing, rs::T, ns::U, - pmapper) where {T <: AbstractArray, U <: AbstractArray} +function setup_majump_to_merge(::Nothing, rs::T, ns::U, pmapper, + rescale_rates_on_update::Bool) where {T <: AbstractArray, U <: AbstractArray} MassActionJump(nothing, [rs], [ns], (pmapper === nothing) ? pmapper : to_collection(pmapper); - scale_rates = false) + scale_rates = false, rescale_rates_on_update) end # when given a collection of reactions to add to maj @@ -697,10 +711,12 @@ function majump_merge!(maj::MassActionJump{T, S, U, V}, sr::T, rs::S, ns::U, (param_mapper === nothing) || error("Error, trying to merge a MassActionJump with a parameter mapping to one without a parameter mapping.") return MassActionJump(rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], - param_mapper; scale_rates = false) + param_mapper; scale_rates = false, + rescale_rates_on_update = maj.rescale_rates_on_update) else return MassActionJump(rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], - merge(maj.param_mapper, param_mapper); scale_rates = false) + merge(maj.param_mapper, param_mapper); scale_rates = false, + rescale_rates_on_update = maj.rescale_rates_on_update) end end @@ -708,6 +724,8 @@ massaction_jump_combine(maj1::MassActionJump, maj2::Nothing) = maj1 massaction_jump_combine(maj1::Nothing, maj2::MassActionJump) = maj2 massaction_jump_combine(maj1::Nothing, maj2::Nothing) = maj1 function massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) + (maj1.rescale_rates_on_update == maj2.rescale_rates_on_update) || + error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") majump_merge!(maj1, maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, maj2.param_mapper) end diff --git a/src/spatial/spatial_massaction_jump.jl b/src/spatial/spatial_massaction_jump.jl index f1a4fea55..c9af9d1e2 100644 --- a/src/spatial/spatial_massaction_jump.jl +++ b/src/spatial/spatial_massaction_jump.jl @@ -82,7 +82,9 @@ function SpatialMassActionJump(urates::A, rs, ns; scale_rates = true, useiszero useiszero = useiszero, nocopy = nocopy) end -function SpatialMassActionJump(ma_jumps::MassActionJump{T, S, U, V}; scale_rates = true, +# scale_rates defaults to false since ma_jumps.scaled_rates are already scaled; +# passing true would double-scale. +function SpatialMassActionJump(ma_jumps::MassActionJump{T, S, U, V}; scale_rates = false, useiszero = true, nocopy = false) where {T, S, U, V} SpatialMassActionJump(ma_jumps.scaled_rates, ma_jumps.reactant_stoch, ma_jumps.net_stoch, ma_jumps.param_mapper; diff --git a/test/runtests.jl b/test/runtests.jl index e4b097c41..6df569d28 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -31,6 +31,7 @@ end @time @safetestset "Mass Action Jump Tests; Nonlinear Rx Model" begin include("bimolerx_test.jl") end @time @safetestset "Mass Action Jump Tests; Special Cases" begin include("degenerate_rx_cases.jl") end @time @safetestset "Mass Action Jump Tests; Floating Point Inputs" begin include("fp_unknowns.jl") end + @time @safetestset "scale_rates Field Tests" begin include("scale_rates_field_test.jl") end @time @safetestset "Direct allocations test" begin include("allocations.jl") end @time @safetestset "Bracketing Tests" begin include("bracketing.jl") end @time @safetestset "Composition-Rejection Table Tests" begin include("table_test.jl") end diff --git a/test/scale_rates_field_test.jl b/test/scale_rates_field_test.jl new file mode 100644 index 000000000..f07ce8d5f --- /dev/null +++ b/test/scale_rates_field_test.jl @@ -0,0 +1,156 @@ +using JumpProcesses, OrdinaryDiffEq, Test + +# Reaction: 3X → Y (third-order, factorial(3) = 6) +# reactant_stoch: species 1 consumed with stoichiometry 3 +# net_stoch: species 1 loses 3, species 2 gains 1 +reactant_stoch = [[1 => 3]] +net_stoch = [[1 => -3, 2 => 1]] + +# Custom mapper: returns pre-scaled rates (like MTKBase's JumpSysMajParamMapper) +struct PreScaledMapper + param_idxs::Vector{Int} + reactant_stoch::Vector{Vector{Pair{Int, Int}}} +end +function (m::PreScaledMapper)(params) + rates = [params[i] for i in m.param_idxs] + JumpProcesses.scalerates!(rates, m.reactant_stoch) + rates +end +function (m::PreScaledMapper)(maj::MassActionJump, newparams; scale_rates, kwargs...) + for i in 1:JumpProcesses.get_num_majumps(maj) + maj.scaled_rates[i] = newparams[m.param_idxs[i]] + end + JumpProcesses.scalerates!(maj.scaled_rates, maj.reactant_stoch) + nothing +end +JumpProcesses.to_collection(m::PreScaledMapper) = m +function Base.merge!(m1::PreScaledMapper, m2::PreScaledMapper) + append!(m1.param_idxs, m2.param_idxs) + append!(m1.reactant_stoch, m2.reactant_stoch) +end + +# Test 1: rescale_rates_on_update field is stored correctly +@testset "rescale_rates_on_update field storage" begin + # Default: scale_rates = true → rescale_rates_on_update = true + maj = MassActionJump([6.0], reactant_stoch, net_stoch) + @test maj.rescale_rates_on_update == true + @test maj.scaled_rates[1] ≈ 1.0 # 6.0 / 3! = 1.0 + + # Explicit: scale_rates = false → rescale_rates_on_update = false + maj = MassActionJump([6.0], reactant_stoch, net_stoch; scale_rates = false) + @test maj.rescale_rates_on_update == false + @test maj.scaled_rates[1] ≈ 6.0 # no scaling + + # Parameterized + maj = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) + @test maj.rescale_rates_on_update == true +end + +# Test 2: update_parameters! respects stored rescale_rates_on_update +@testset "update_parameters! respects rescale_rates_on_update" begin + # With param_idxs and scale_rates = true (default) — built-in mapper path + p = [6.0] + maj = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) + dprob = DiscreteProblem([100, 0], (0.0, 1.0), p) + jprob = JumpProblem(dprob, Direct(), maj) + @test jprob.massaction_jump.rescale_rates_on_update == true + @test jprob.massaction_jump.scaled_rates[1] ≈ 1.0 # 6.0 / 3! + + # update_parameters! should scale (rescale_rates_on_update = true from struct) + JumpProcesses.update_parameters!(jprob.massaction_jump, [12.0]) + @test jprob.massaction_jump.scaled_rates[1] ≈ 2.0 # 12.0 / 3! +end + +# Test 3: Custom pre-scaled mapper with scale_rates = false — the bug reproducer +@testset "Pre-scaled mapper with scale_rates=false (bug reproducer)" begin + mapper = PreScaledMapper([1], reactant_stoch) + k = 6.0 + p = [k] + + maj = MassActionJump(reactant_stoch, net_stoch; param_mapper = mapper, scale_rates = false) + dprob = DiscreteProblem([100, 0], (0.0, 100.0), p) + jprob = JumpProblem(dprob, Direct(), maj; scale_rates = false) + + expected_scaled = k / factorial(3) # 1.0 + @test jprob.massaction_jump.scaled_rates[1] ≈ expected_scaled + @test jprob.massaction_jump.rescale_rates_on_update == false + + # Test reset_aggregated_jumps! does NOT double-scale + integ = init(jprob, SSAStepper()) + integ.p[1] = 18.0 + reset_aggregated_jumps!(integ) + new_expected = 18.0 / factorial(3) # 3.0, NOT 3.0/6 = 0.5 + @test jprob.massaction_jump.scaled_rates[1] ≈ new_expected + + # Test remake does NOT double-scale + jprob2 = remake(jprob; p = [24.0]) + remake_expected = 24.0 / factorial(3) # 4.0, NOT 4.0/6 ≈ 0.667 + @test jprob2.massaction_jump.scaled_rates[1] ≈ remake_expected + + # Test remake round-trip + jprob3 = remake(jprob2; p = [k]) + @test jprob3.massaction_jump.scaled_rates[1] ≈ expected_scaled +end + +# Test 4: Callback parameter changes with built-in mapper (rescale_rates_on_update = true) +@testset "Callback with built-in mapper" begin + p = [6.0] + maj = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) + dprob = DiscreteProblem([100, 0], (0.0, 2000.0), p) + jprob = JumpProblem(dprob, Direct(), maj; save_positions = (false, false)) + + condit(u, t, integrator) = t == 1000.0 + function affect!(integrator) + integrator.p[1] = 24.0 + reset_aggregated_jumps!(integrator) + end + cb = DiscreteCallback(condit, affect!) + sol = solve(jprob, SSAStepper(); tstops = [1000.0], callback = cb) + @test jprob.massaction_jump.scaled_rates[1] ≈ 4.0 # 24.0 / 3! +end + +# Test 5: rescale_rates_on_update propagated through JumpSet merge and JumpProblem varargs +# Use explicit-rate MAJs since the JumpSet vector merge path doesn't support +# parameterized (Nothing-rated) MAJs. +@testset "rescale_rates_on_update propagated through merge paths" begin + reactant_stoch2 = [[2 => 3]] + net_stoch2 = [[2 => -3, 1 => 1]] + + # --- JumpSet merge path --- + + # Two MAJs with matching rescale_rates_on_update = false (rates already scaled) + maj1 = MassActionJump([1.0], reactant_stoch, net_stoch; scale_rates = false) + maj2 = MassActionJump([2.0], reactant_stoch2, net_stoch2; scale_rates = false) + jset = JumpSet(; massaction_jumps = [maj1, maj2]) + @test jset.massaction_jump.rescale_rates_on_update == false + + # Two MAJs with matching rescale_rates_on_update = true (rates get scaled) + maj3 = MassActionJump([6.0], reactant_stoch, net_stoch) # scale_rates = true (default) + maj4 = MassActionJump([12.0], reactant_stoch2, net_stoch2) + jset2 = JumpSet(; massaction_jumps = [maj3, maj4]) + @test jset2.massaction_jump.rescale_rates_on_update == true + + # Mismatched rescale_rates_on_update via JumpSet — should error + maj_true = MassActionJump([6.0], reactant_stoch, net_stoch) # rescale = true + maj_false = MassActionJump([1.0], reactant_stoch2, net_stoch2; scale_rates = false) # rescale = false + @test_throws ErrorException JumpSet(; massaction_jumps = [maj_true, maj_false]) + + # --- JumpProblem varargs path (split_jumps → massaction_jump_combine) --- + + dprob = DiscreteProblem([100, 100], (0.0, 1.0), [1.0, 1.0]) + + # Two MAJs with matching rescale_rates_on_update = false via JumpProblem varargs + maj_f1 = MassActionJump([1.0], reactant_stoch, net_stoch; scale_rates = false) + maj_f2 = MassActionJump([2.0], reactant_stoch2, net_stoch2; scale_rates = false) + jprob_f = JumpProblem(dprob, Direct(), maj_f1, maj_f2; scale_rates = false) + @test jprob_f.massaction_jump.rescale_rates_on_update == false + + # Two MAJs with matching rescale_rates_on_update = true via JumpProblem varargs + maj_t1 = MassActionJump([6.0], reactant_stoch, net_stoch) + maj_t2 = MassActionJump([12.0], reactant_stoch2, net_stoch2) + jprob_t = JumpProblem(dprob, Direct(), maj_t1, maj_t2) + @test jprob_t.massaction_jump.rescale_rates_on_update == true + + # Mismatched rescale_rates_on_update via JumpProblem varargs — should error + @test_throws ErrorException JumpProblem(dprob, Direct(), maj_true, maj_false) +end From 83352c59934099fb3a84193ef96e40e9a4eca1a6 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Thu, 26 Feb 2026 18:20:47 -0500 Subject: [PATCH 17/52] Fix PreScaledMapper test to actually catch double-scaling regression The mapper now conditionally applies scalerates! based on scale_rates, matching MTKBase JumpSysMajParamMapper behavior. Previously it ignored scale_rates entirely, so the test would pass even with the old buggy default of scale_rates=true. Co-Authored-By: Claude Opus 4.5 --- test/scale_rates_field_test.jl | 11 +++++++++-- 1 file changed, 9 insertions(+), 2 deletions(-) diff --git a/test/scale_rates_field_test.jl b/test/scale_rates_field_test.jl index f07ce8d5f..3f06ed888 100644 --- a/test/scale_rates_field_test.jl +++ b/test/scale_rates_field_test.jl @@ -6,7 +6,13 @@ using JumpProcesses, OrdinaryDiffEq, Test reactant_stoch = [[1 => 3]] net_stoch = [[1 => -3, 2 => 1]] -# Custom mapper: returns pre-scaled rates (like MTKBase's JumpSysMajParamMapper) +# Custom mapper mimicking MTKBase's JumpSysMajParamMapper. +# The initial rates callable always pre-scales (simulating symbolic expressions that +# already contain the combinatoric factor, e.g. k/3! for 3X → Y). +# The update callable also pre-scales, then conditionally applies scalerates! based on +# scale_rates — matching MTKBase's behavior. This is what makes the test an actual +# regression test: with the old default of scale_rates=true, the second scalerates! +# would double-scale. struct PreScaledMapper param_idxs::Vector{Int} reactant_stoch::Vector{Vector{Pair{Int, Int}}} @@ -20,7 +26,8 @@ function (m::PreScaledMapper)(maj::MassActionJump, newparams; scale_rates, kwarg for i in 1:JumpProcesses.get_num_majumps(maj) maj.scaled_rates[i] = newparams[m.param_idxs[i]] end - JumpProcesses.scalerates!(maj.scaled_rates, maj.reactant_stoch) + JumpProcesses.scalerates!(maj.scaled_rates, m.reactant_stoch) + scale_rates && JumpProcesses.scalerates!(maj.scaled_rates, maj.reactant_stoch) nothing end JumpProcesses.to_collection(m::PreScaledMapper) = m From 4a52abe0448f26a0fdc21a1c116e7d2ce801f69a Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Thu, 26 Feb 2026 21:11:45 -0500 Subject: [PATCH 18/52] Update project version to 9.23 --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 35ea0c114..1c59737f5 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "JumpProcesses" uuid = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5" authors = ["Chris Rackauckas "] -version = "9.22.2" +version = "9.23" [deps] ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9" From d1d6beb928d1f66cd0be8d5a56164a83eb9c4dc3 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 28 Feb 2026 12:31:43 -0500 Subject: [PATCH 19/52] add thread-safety warning to JumpProblem docstring Document that JumpProblem contains mutable state and is not thread-safe for concurrent solve calls. Note that EnsembleProblem handles isolation automatically via SciMLBase per-task deepcopy and per-trajectory RNG. Co-Authored-By: Claude Opus 4.5 --- src/problem.jl | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/src/problem.jl b/src/problem.jl index 3ce038de5..56d58da62 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -67,6 +67,15 @@ Please see the [tutorial page](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) in the DifferentialEquations.jl [docs](https://docs.sciml.ai/JumpProcesses/stable/) for usage examples and commonly asked questions. + +!!! warning "Thread Safety" + `JumpProblem` contains mutable state (aggregator data, callbacks) and is **not + thread-safe**. A single `JumpProblem` instance must not be solved concurrently from + multiple threads or tasks without first creating independent copies via `deepcopy`. + When running ensemble simulations via `EnsembleProblem`, this is handled automatically + — the `SciMLBase` ensemble layer provides per-task isolation and per-trajectory RNG + seeding. This warning only applies to manually parallelized `solve` calls outside the + ensemble interface. """ mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggregator}, J1, J2, J3, J4, K} <: DiffEqBase.AbstractJumpProblem{P, J} From a07a527ec48e5875fdbc8378c889b7e5f07593ca Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sun, 1 Mar 2026 17:21:51 -0500 Subject: [PATCH 20/52] Temporarily use SciMLBase fork branch for ensemble_rng_redesign Co-Authored-By: Claude Opus 4.5 --- Project.toml | 3 +++ 1 file changed, 3 insertions(+) diff --git a/Project.toml b/Project.toml index 63f3ea262..242da5aee 100644 --- a/Project.toml +++ b/Project.toml @@ -61,6 +61,9 @@ SymbolicIndexingInterface = "0.3.36" Test = "1" julia = "1.10" +[sources] +SciMLBase = {url = "https://github.com/isaacsas/SciMLBase.jl.git", rev = "ensemble_rng_redesign"} + [extras] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" From 12e9866f7a18b0db4946461967985a0d7288d6f6 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 3 Mar 2026 15:21:57 -0500 Subject: [PATCH 21/52] Immutable MassActionJump refactor: move working rates to aggregation structs Move mutable rate working-copy (maj_rates) out of MassActionJump and into all 10 aggregation structs + RxRates. Rates are now lazily materialized via fill_scaled_rates! during initialize! instead of eagerly in the JumpProblem constructor. This fixes aliasing bugs where multiple JumpProblems sharing the same MAJ could silently corrupt each others rates. Key changes: - Add maj_rates field to all aggregation structs and RxRates - Add fill_scaled_rates! dispatcher for parameterized vs explicit MAJs - Update evalrxrate to 4-arg signature taking explicit maj_rates - Propagate maj_rates through calculate_jump_rate, rejectrx, bracketing - Remove update_parameters!, finalize_parameters_hook! - Simplify reset_aggregated_jumps! (remove update_jump_params kwarg) - JumpProblem constructor keeps original parameterized MAJ - Add merge guards for mapper-backed MAJs - Add SpatialMassActionJump guard for parameterized MAJs - Remove dead 1-arg mapper callables - Update all tests for lazy rate materialization pattern Co-Authored-By: Claude Opus 4.5 --- src/aggregators/aggregated_api.jl | 46 +++++-------- src/aggregators/bracketing.jl | 9 +-- src/aggregators/ccnrm.jl | 10 ++- src/aggregators/coevolve.jl | 10 ++- src/aggregators/direct.jl | 13 ++-- src/aggregators/directcr.jl | 8 ++- src/aggregators/frm.jl | 10 ++- src/aggregators/nrm.jl | 10 ++- src/aggregators/rdirect.jl | 7 +- src/aggregators/rssa.jl | 7 +- src/aggregators/rssacr.jl | 11 +-- src/aggregators/sortingdirect.jl | 5 +- src/aggregators/ssajump.jl | 18 +++-- src/jumps.jl | 76 ++++++++++++--------- src/massaction_rates.jl | 10 +-- src/problem.jl | 46 ++----------- src/simple_regular_solve.jl | 8 ++- src/spatial/directcrdirect.jl | 1 + src/spatial/flatten.jl | 4 +- src/spatial/nsm.jl | 1 + src/spatial/reaction_rates.jl | 15 ++-- src/spatial/spatial_massaction_jump.jl | 5 ++ test/jprob_symbol_indexing.jl | 43 ++++++++---- test/regular_jumps.jl | 6 +- test/scale_rates_field_test.jl | 95 ++++++++++++++++++-------- test/ssa_callback_test.jl | 13 ++-- 26 files changed, 283 insertions(+), 204 deletions(-) diff --git a/src/aggregators/aggregated_api.jl b/src/aggregators/aggregated_api.jl index a65c4ad82..ec3376d0e 100644 --- a/src/aggregators/aggregated_api.jl +++ b/src/aggregators/aggregated_api.jl @@ -1,55 +1,41 @@ """ - reset_aggregated_jumps!(integrator, uprev = nothing; update_jump_params=true) + reset_aggregated_jumps!(integrator, uprev = nothing) Reset the state of jump processes and associated solvers following a change -in parameters or such. - -Notes - - - `update_jump_params=true` will recalculate the rates stored within any - MassActionJump that was built from the parameter vector. If the parameter - vector is unchanged, this can safely be set to false to improve performance. +in parameters or such. Rate updates are handled automatically by `initialize!` +via `fill_scaled_rates!`. """ -function reset_aggregated_jumps!(integrator, uprev = nothing; update_jump_params = true, - kwargs...) - reset_aggregated_jumps!(integrator, uprev, integrator.opts.callback, - update_jump_params = update_jump_params, kwargs...) +function reset_aggregated_jumps!(integrator, uprev = nothing; kwargs...) + if haskey(kwargs, :update_jump_params) + throw(ArgumentError("`update_jump_params` keyword argument has been removed. " * + "Rate updates are now handled automatically by `initialize!` " * + "via `fill_scaled_rates!`.")) + end + reset_aggregated_jumps!(integrator, uprev, integrator.opts.callback) nothing end -function reset_aggregated_jumps!(integrator, uprev, callback::Nothing; - update_jump_params = true, kwargs...) +function reset_aggregated_jumps!(integrator, uprev, callback::Nothing) nothing end -function reset_aggregated_jumps!(integrator, uprev, callback::CallbackSet; - update_jump_params = true, kwargs...) +function reset_aggregated_jumps!(integrator, uprev, callback::CallbackSet) if !isempty(callback.discrete_callbacks) - reset_aggregated_jumps!(integrator, uprev, callback.discrete_callbacks..., - update_jump_params = update_jump_params, kwargs...) + reset_aggregated_jumps!(integrator, uprev, callback.discrete_callbacks...) end nothing end -function reset_aggregated_jumps!(integrator, uprev, cb::DiscreteCallback, cbs...; - update_jump_params = true, kwargs...) +function reset_aggregated_jumps!(integrator, uprev, cb::DiscreteCallback, cbs...) if cb.condition isa AbstractSSAJumpAggregator - maj = cb.condition.ma_jumps - update_jump_params && using_params(maj) && - update_parameters!(cb.condition.ma_jumps, integrator.p; kwargs...) cb.condition(cb, integrator.u, integrator.t, integrator) end - reset_aggregated_jumps!(integrator, uprev, cbs...; - update_jump_params = update_jump_params, kwargs...) + reset_aggregated_jumps!(integrator, uprev, cbs...) nothing end -function reset_aggregated_jumps!(integrator, uprev, cb::DiscreteCallback; - update_jump_params = true, kwargs...) +function reset_aggregated_jumps!(integrator, uprev, cb::DiscreteCallback) if cb.condition isa AbstractSSAJumpAggregator - maj = cb.condition.ma_jumps - update_jump_params && using_params(maj) && - update_parameters!(cb.condition.ma_jumps, integrator.p; kwargs...) cb.condition(cb, integrator.u, integrator.t, integrator) end nothing diff --git a/src/aggregators/bracketing.jl b/src/aggregators/bracketing.jl index 18bb9ece3..b66092c46 100644 --- a/src/aggregators/bracketing.jl +++ b/src/aggregators/bracketing.jl @@ -47,8 +47,8 @@ end @inline get_spec_brackets(bd, i, u::AbstractVector) = get_spec_brackets(bd, i, u[i]) # Get propensity brackets of massaction jump k. -@inline function get_majump_brackets(ulow, uhigh, k, majumps) - evalrxrate(ulow, k, majumps), evalrxrate(uhigh, k, majumps) +@inline function get_majump_brackets(ulow, uhigh, k, majumps, maj_rates) + evalrxrate(ulow, k, majumps, maj_rates), evalrxrate(uhigh, k, majumps, maj_rates) end # for constant rate jumps we must check the ordering of the bracket values @@ -66,7 +66,7 @@ get brackets for the rate of reaction rx by first checking if the reaction is a ma_jumps = p.ma_jumps num_majumps = get_num_majumps(ma_jumps) if rx <= num_majumps - return get_majump_brackets(p.ulow, p.uhigh, rx, ma_jumps) + return get_majump_brackets(p.ulow, p.uhigh, rx, ma_jumps, p.maj_rates) else @inbounds return get_cjump_brackets(p.ulow, p.uhigh, p.rates[rx - num_majumps], params, t) @@ -108,8 +108,9 @@ function set_bracketing!(p::AbstractSSAJumpAggregator, u, params, t) majumps = p.ma_jumps crlow = p.cur_rate_low crhigh = p.cur_rate_high + maj_rates = p.maj_rates @inbounds for k in 1:get_num_majumps(majumps) - crlow[k], crhigh[k] = get_majump_brackets(p.ulow, p.uhigh, k, majumps) + crlow[k], crhigh[k] = get_majump_brackets(p.ulow, p.uhigh, k, majumps, maj_rates) sum_rate += crhigh[k] end diff --git a/src/aggregators/ccnrm.jl b/src/aggregators/ccnrm.jl index 720868c50..4788df884 100644 --- a/src/aggregators/ccnrm.jl +++ b/src/aggregators/ccnrm.jl @@ -12,6 +12,7 @@ mutable struct CCNRMJumpAggregation{T, S, F1, F2, DEPGR, PT} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} @@ -21,6 +22,7 @@ end function CCNRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), num_specs, dep_graph = nothing, kwargs...) where {T, S, F1, F2} @@ -46,7 +48,7 @@ function CCNRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, affecttype = F2 <: Tuple ? F2 : Any CCNRMJumpAggregation{T, S, F1, affecttype, typeof(dg), typeof(ptt)}( nj, nj, njt, et, - crs, sr, maj, + crs, sr, maj, maj_rates, rs, affs!, sps, dg, ptt) end @@ -67,6 +69,7 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::CCNRMJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) rng = get_rng(integrator) initialize_rates_and_times!(p, u, params, t, rng) generate_jumps!(p, integrator, u, params, t) @@ -115,7 +118,7 @@ function update_dependent_rates!(p::CCNRMJumpAggregation, u, params, t, rng) # update the jump rate @inbounds cur_rates[rx] = calculate_jump_rate(ma_jumps, num_majumps, rates, u, - params, t, rx) + params, t, rx, p.maj_rates) # Calculate new jump times for dependent jumps if rx != p.next_jump && oldrate > zero(oldrate) @@ -141,8 +144,9 @@ function initialize_rates_and_times!(p::CCNRMJumpAggregation, u, params, t, rng) majumps = p.ma_jumps cur_rates = p.cur_rates pttdata = Vector{typeof(t)}(undef, length(cur_rates)) + maj_rates = p.maj_rates @inbounds for i in 1:get_num_majumps(majumps) - cur_rates[i] = evalrxrate(u, i, majumps) + cur_rates[i] = evalrxrate(u, i, majumps, maj_rates) pttdata[i] = t + randexp(rng) / cur_rates[i] end diff --git a/src/aggregators/coevolve.jl b/src/aggregators/coevolve.jl index ab5dc5f4f..651a7012d 100644 --- a/src/aggregators/coevolve.jl +++ b/src/aggregators/coevolve.jl @@ -10,6 +10,7 @@ mutable struct CoevolveJumpAggregation{T, S, F1, F2, GR, PQ} <: cur_rates::Vector{T} # the last computed upper bound for each rate sum_rate::Nothing # not used ma_jumps::S # MassActionJumps + maj_rates::Vector{T} # working copy of mass action jump rates rates::F1 # vector of rate functions affects!::F2 # vector of affect functions for VariableRateJumps save_positions::Tuple{Bool, Bool} # tuple for whether to save the jumps before and/or after event @@ -24,6 +25,7 @@ end function CoevolveJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::Nothing, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), u::U, dep_graph = nothing, lrates, urates, rateintervals, haslratevec, cur_lrates::Vector{T}) where {T, S, F1, F2, U} @@ -49,7 +51,7 @@ function CoevolveJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::Not pq = MutableBinaryMinHeap{T}() affecttype = F2 <: Tuple ? F2 : Any CoevolveJumpAggregation{T, S, F1, affecttype, typeof(dg), - typeof(pq)}(nj, nj, njt, et, crs, sr, maj, + typeof(pq)}(nj, nj, njt, et, crs, sr, maj, maj_rates, rs, affs!, sps, dg, pq, lrates, urates, rateintervals, haslratevec, cur_lrates) @@ -135,19 +137,21 @@ function aggregate(aggregator::Coevolve, u, p, t, end_time, constant_jumps, num_jumps = get_num_majumps(ma_jumps) + nrjs cur_rates = Vector{typeof(t)}(undef, num_jumps) + maj_rates = Vector{typeof(t)}(undef, get_num_majumps(ma_jumps)) cur_lrates = zeros(typeof(t), nvrjs) sum_rate = nothing next_jump = 0 next_jump_time = typemax(t) CoevolveJumpAggregation(next_jump, next_jump_time, end_time, cur_rates, sum_rate, ma_jumps, rates, affects!, save_positions; - u, dep_graph, lrates, urates, rateintervals, haslratevec, + maj_rates, u, dep_graph, lrates, urates, rateintervals, haslratevec, cur_lrates) end # set up a new simulation and calculate the first jump / jump time function initialize!(p::CoevolveJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) rng = get_rng(integrator) fill_rates_and_get_times!(p, u, params, t, rng) generate_jumps!(p, integrator, u, params, t) @@ -239,7 +243,7 @@ function update_dependent_rates!(p::CoevolveJumpAggregation, u, params, t, rng) end @inline function get_ma_urate(p::CoevolveJumpAggregation, i, u, params, t) - return evalrxrate(u, i, p.ma_jumps) + return evalrxrate(u, i, p.ma_jumps, p.maj_rates) end @inline function get_urate(p::CoevolveJumpAggregation, uidx, u, params, t) diff --git a/src/aggregators/direct.jl b/src/aggregators/direct.jl index b51e18f7c..f976a4cb0 100644 --- a/src/aggregators/direct.jl +++ b/src/aggregators/direct.jl @@ -7,16 +7,18 @@ mutable struct DirectJumpAggregation{T, S, F1, F2} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} end function DirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), kwargs...) where {T, S, F1, F2} affecttype = F2 <: Tuple ? F2 : Any - DirectJumpAggregation{T, S, F1, affecttype}(nj, nj, njt, et, crs, sr, maj, rs, - affs!, sps) + DirectJumpAggregation{T, S, F1, affecttype}(nj, nj, njt, et, crs, sr, maj, maj_rates, + rs, affs!, sps) end ############################# Required Functions ############################# @@ -46,6 +48,7 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::DirectJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) generate_jumps!(p, integrator, u, params, t) nothing end @@ -77,9 +80,10 @@ function time_to_next_jump(p::DirectJumpAggregation{T, S, F1}, u, params, # mass action rates majumps = p.ma_jumps + maj_rates = p.maj_rates idx = get_num_majumps(majumps) @inbounds for i in 1:idx - new_rate = evalrxrate(u, i, majumps) + new_rate = evalrxrate(u, i, majumps, maj_rates) cur_rates[i] = add_fast(new_rate, prev_rate) prev_rate = cur_rates[i] end @@ -119,9 +123,10 @@ function time_to_next_jump(p::DirectJumpAggregation{T, S, F1}, u, params, # mass action rates majumps = p.ma_jumps + maj_rates = p.maj_rates idx = get_num_majumps(majumps) @inbounds for i in 1:idx - new_rate = evalrxrate(u, i, majumps) + new_rate = evalrxrate(u, i, majumps, maj_rates) cur_rates[i] = add_fast(new_rate, prev_rate) prev_rate = cur_rates[i] end diff --git a/src/aggregators/directcr.jl b/src/aggregators/directcr.jl index 636300bb3..ae4ebf877 100644 --- a/src/aggregators/directcr.jl +++ b/src/aggregators/directcr.jl @@ -20,6 +20,7 @@ mutable struct DirectCRJumpAggregation{T, S, F1, F2, DEPGR, U <: PriorityTable, cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} @@ -32,6 +33,7 @@ end function DirectCRJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), num_specs, dep_graph = nothing, minrate = convert(T, MINJUMPRATE), maxrate = convert(T, Inf), @@ -64,7 +66,7 @@ function DirectCRJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, affecttype = F2 <: Tuple ? F2 : Any DirectCRJumpAggregation{T, S, F1, affecttype, typeof(dg), typeof(rt), typeof(ratetogroup)}(nj, nj, njt, et, crs, sr, maj, - rs, affs!, sps, dg, + maj_rates, rs, affs!, sps, dg, minrate, maxrate, rt, ratetogroup) end @@ -86,6 +88,7 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::DirectCRJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) # initialize rates fill_rates_and_sum!(p, u, params, t) @@ -134,7 +137,8 @@ function update_dependent_rates!(p::DirectCRJumpAggregation, u, params, t) oldrate = cur_rates[rx] # update rate - cur_rates[rx] = calculate_jump_rate(ma_jumps, num_majumps, rates, u, params, t, rx) + cur_rates[rx] = calculate_jump_rate(ma_jumps, num_majumps, rates, u, params, t, rx, + p.maj_rates) # update table update!(rt, rx, oldrate, cur_rates[rx]) diff --git a/src/aggregators/frm.jl b/src/aggregators/frm.jl index 0db30113b..128b93dba 100644 --- a/src/aggregators/frm.jl +++ b/src/aggregators/frm.jl @@ -7,16 +7,18 @@ mutable struct FRMJumpAggregation{T, S, F1, F2} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} end function FRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), kwargs...) where {T, S, F1, F2} affecttype = F2 <: Tuple ? F2 : Any - FRMJumpAggregation{T, S, F1, affecttype}(nj, nj, njt, et, crs, sr, maj, rs, - affs!, sps) + FRMJumpAggregation{T, S, F1, affecttype}(nj, nj, njt, et, crs, sr, maj, maj_rates, + rs, affs!, sps) end ############################# Required Functions ############################# @@ -48,6 +50,7 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::FRMJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) generate_jumps!(p, integrator, u, params, t) nothing end @@ -83,8 +86,9 @@ function next_ma_jump(p::FRMJumpAggregation, u, params, t, rng) ttnj = typemax(typeof(t)) nextrx = zero(Int) majumps = p.ma_jumps + maj_rates = p.maj_rates @inbounds for i in 1:get_num_majumps(majumps) - p.cur_rates[i] = evalrxrate(u, i, majumps) + p.cur_rates[i] = evalrxrate(u, i, majumps, maj_rates) dt = randexp(rng) / p.cur_rates[i] if dt < ttnj ttnj = dt diff --git a/src/aggregators/nrm.jl b/src/aggregators/nrm.jl index 934540448..ef1a6b74d 100644 --- a/src/aggregators/nrm.jl +++ b/src/aggregators/nrm.jl @@ -10,6 +10,7 @@ mutable struct NRMJumpAggregation{T, S, F1, F2, DEPGR, PQ} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} @@ -19,6 +20,7 @@ end function NRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), num_specs, dep_graph = nothing, kwargs...) where {T, S, F1, F2} @@ -40,7 +42,7 @@ function NRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, affecttype = F2 <: Tuple ? F2 : Any NRMJumpAggregation{T, S, F1, affecttype, typeof(dg), typeof(pq)}(nj, nj, njt, et, - crs, sr, maj, + crs, sr, maj, maj_rates, rs, affs!, sps, dg, pq) end @@ -61,6 +63,7 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::NRMJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) rng = get_rng(integrator) fill_rates_and_get_times!(p, u, params, t, rng) generate_jumps!(p, integrator, u, params, t) @@ -98,7 +101,7 @@ function update_dependent_rates!(p::NRMJumpAggregation, u, params, t, rng) # update the jump rate @inbounds cur_rates[rx] = calculate_jump_rate(ma_jumps, num_majumps, rates, u, - params, t, rx) + params, t, rx, p.maj_rates) # calculate new jump times for dependent jumps if rx != p.next_jump && oldrate > zero(oldrate) @@ -125,8 +128,9 @@ function fill_rates_and_get_times!(p::NRMJumpAggregation, u, params, t, rng) majumps = p.ma_jumps cur_rates = p.cur_rates pqdata = Vector{typeof(t)}(undef, length(cur_rates)) + maj_rates = p.maj_rates @inbounds for i in 1:get_num_majumps(majumps) - cur_rates[i] = evalrxrate(u, i, majumps) + cur_rates[i] = evalrxrate(u, i, majumps, maj_rates) pqdata[i] = t + randexp(rng) / cur_rates[i] end diff --git a/src/aggregators/rdirect.jl b/src/aggregators/rdirect.jl index eea283160..8d8b34c39 100644 --- a/src/aggregators/rdirect.jl +++ b/src/aggregators/rdirect.jl @@ -11,6 +11,7 @@ mutable struct RDirectJumpAggregation{T, S, F1, F2, DEPGR} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} @@ -22,6 +23,7 @@ end function RDirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), num_specs, counter_threshold = length(crs), dep_graph = nothing, kwargs...) where {T, S, F1, F2} @@ -42,7 +44,7 @@ function RDirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, m max_rate = maximum(crs) affecttype = F2 <: Tuple ? F2 : Any return RDirectJumpAggregation{T, S, F1, affecttype, typeof(dg)}(nj, nj, njt, et, - crs, sr, maj, rs, + crs, sr, maj, maj_rates, rs, affs!, sps, dg, max_rate, 0, counter_threshold) @@ -65,6 +67,7 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::RDirectJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) fill_rates_and_sum!(p, u, params, t) p.max_rate = maximum(p.cur_rates) generate_jumps!(p, integrator, u, params, t) @@ -118,7 +121,7 @@ function update_dependent_rates!(p::RDirectJumpAggregation, u, params, t) @inbounds for rx in dep_rxs @inbounds new_rate = calculate_jump_rate( ma_jumps, num_majumps, rates, u, params, t, - rx) + rx, p.maj_rates) sum_rate += new_rate - cur_rates[rx] if new_rate > p.max_rate p.max_rate = new_rate diff --git a/src/aggregators/rssa.jl b/src/aggregators/rssa.jl index f86e0b444..1e19259e0 100644 --- a/src/aggregators/rssa.jl +++ b/src/aggregators/rssa.jl @@ -14,6 +14,7 @@ mutable struct RSSAJumpAggregation{T, S, F1, F2, VJMAP, JVMAP, BD, U} <: cur_rate_high::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} @@ -26,6 +27,7 @@ end function RSSAJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), u::U, vartojumps_map = nothing, jumptovars_map = nothing, bracket_data = nothing, kwargs...) where {T, S, F1, F2, U} @@ -64,7 +66,7 @@ function RSSAJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, affecttype = F2 <: Tuple ? F2 : Any RSSAJumpAggregation{T, S, F1, affecttype, typeof(vtoj_map), typeof(jtov_map), typeof(bd), U}(nj, nj, njt, et, crl_bnds, - crh_bnds, sr, maj, rs, affs!, sps, + crh_bnds, sr, maj, maj_rates, rs, affs!, sps, vtoj_map, jtov_map, bd, ulow, uhigh) end @@ -86,6 +88,7 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::RSSAJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) set_bracketing!(p, u, params, t) generate_jumps!(p, integrator, u, params, t) nothing @@ -120,7 +123,7 @@ function generate_jumps!(p::RSSAJumpAggregation, integrator, u, params, t) end rerl += randexp(rng) @inbounds while rejectrx(ma_jumps, num_majumps, rates, cur_rate_high, - cur_rate_low, rng, u, jidx, params, t) + cur_rate_low, rng, u, jidx, params, t, p.maj_rates) # sample candidate reaction r = rand(rng) * sum_rate jidx = linear_search(cur_rate_high, r) diff --git a/src/aggregators/rssacr.jl b/src/aggregators/rssacr.jl index 9cc47c3b8..887c718ea 100644 --- a/src/aggregators/rssacr.jl +++ b/src/aggregators/rssacr.jl @@ -15,6 +15,7 @@ mutable struct RSSACRJumpAggregation{F, S, F1, F2, U, VJMAP, JVMAP, BD, cur_rate_high::Vector{F} sum_rate::F ma_jumps::S + maj_rates::Vector{F} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} @@ -30,8 +31,9 @@ mutable struct RSSACRJumpAggregation{F, S, F1, F2, U, VJMAP, JVMAP, BD, end function RSSACRJumpAggregation(nj::Int, njt::F, et::F, crs::Vector{F}, sum_rate::F, maj::S, - rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; u::U, - vartojumps_map = nothing, jumptovars_map = nothing, + rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{F}(undef, get_num_majumps(maj)), + u::U, vartojumps_map = nothing, jumptovars_map = nothing, bracket_data = nothing, minrate = convert(F, MINJUMPRATE), maxrate = convert(F, Inf), kwargs...) where {F, S, F1, F2, U} @@ -82,7 +84,7 @@ function RSSACRJumpAggregation(nj::Int, njt::F, et::F, crs::Vector{F}, sum_rate: RSSACRJumpAggregation{typeof(njt), S, F1, affecttype, U, typeof(vtoj_map), typeof(jtov_map), typeof(bd), typeof(rt), typeof(ratetogroup)}(nj, nj, njt, et, crl_bnds, crh_bnds, - sum_rate, maj, rs, affs!, sps, vtoj_map, + sum_rate, maj, maj_rates, rs, affs!, sps, vtoj_map, jtov_map, bd, ulow, uhigh, minrate, maxrate, rt, ratetogroup) end @@ -103,6 +105,7 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::RSSACRJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) set_bracketing!(p, u, params, t) # setup PriorityTable @@ -145,7 +148,7 @@ function generate_jumps!(p::RSSACRJumpAggregation, integrator, u, params, t) end rerl += randexp(rng) while rejectrx(ma_jumps, num_majumps, rates, cur_rate_high, cur_rate_low, rng, u, jidx, - params, t) + params, t, p.maj_rates) # sample candidate reaction jidx = sample(rt, cur_rate_high, rng) rerl += randexp(rng) diff --git a/src/aggregators/sortingdirect.jl b/src/aggregators/sortingdirect.jl index c20ea1ae4..6555a757f 100644 --- a/src/aggregators/sortingdirect.jl +++ b/src/aggregators/sortingdirect.jl @@ -11,6 +11,7 @@ mutable struct SortingDirectJumpAggregation{T, S, F1, F2, DEPGR} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} @@ -21,6 +22,7 @@ end function SortingDirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T, maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), num_specs, dep_graph = nothing, kwargs...) where {T, S, F1, F2} @@ -42,7 +44,7 @@ function SortingDirectJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr jtoidx = collect(1:length(crs)) affecttype = F2 <: Tuple ? F2 : Any SortingDirectJumpAggregation{T, S, F1, affecttype, typeof(dg)}(nj, nj, njt, et, - crs, sr, maj, rs, + crs, sr, maj, maj_rates, rs, affs!, sps, dg, jtoidx, zero(Int)) @@ -65,6 +67,7 @@ end # set up a new simulation and calculate the first jump / jump time function initialize!(p::SortingDirectJumpAggregation, integrator, u, params, t) p.end_time = integrator.sol.prob.tspan[2] + fill_scaled_rates!(p.maj_rates, p.ma_jumps, params) fill_rates_and_sum!(p, u, params, t) generate_jumps!(p, integrator, u, params, t) nothing diff --git a/src/aggregators/ssajump.jl b/src/aggregators/ssajump.jl index fec06c185..c68979d7f 100644 --- a/src/aggregators/ssajump.jl +++ b/src/aggregators/ssajump.jl @@ -127,12 +127,13 @@ function build_jump_aggregation(jump_agg_type, u, p, t, end_time, ma_jumps, rate # current jump rates, allows mass action rates and constant jumps cur_rates = Vector{typeof(t)}(undef, get_num_majumps(majumps) + length(rates)) + maj_rates = Vector{typeof(t)}(undef, get_num_majumps(majumps)) sum_rate = zero(typeof(t)) next_jump = 0 next_jump_time = typemax(typeof(t)) jump_agg_type(next_jump, next_jump_time, end_time, cur_rates, sum_rate, - majumps, rates, affects!, save_positions; kwargs...) + majumps, rates, affects!, save_positions; maj_rates, kwargs...) end """ @@ -146,8 +147,9 @@ function fill_rates_and_sum!(p::AbstractSSAJumpAggregator, u, params, t) # mass action jumps majumps = p.ma_jumps cur_rates = p.cur_rates + maj_rates = p.maj_rates @inbounds for i in 1:get_num_majumps(majumps) - cur_rates[i] = evalrxrate(u, i, majumps) + cur_rates[i] = evalrxrate(u, i, majumps, maj_rates) sum_rate += cur_rates[i] end @@ -169,9 +171,10 @@ end Recalculate the rate for the jump with index `rx`. """ -@inline function calculate_jump_rate(ma_jumps, num_majumps, rates, u, params, t, rx) +@inline function calculate_jump_rate(ma_jumps, num_majumps, rates, u, params, t, rx, + maj_rates) if rx <= num_majumps - return evalrxrate(u, rx, ma_jumps) + return evalrxrate(u, rx, ma_jumps, maj_rates) else @inbounds return rates[rx - num_majumps](u, params, t) end @@ -194,7 +197,7 @@ function update_dependent_rates!(p::AbstractSSAJumpAggregator, u, params, t) @inbounds for rx in dep_rxs sum_rate -= cur_rates[rx] @inbounds cur_rates[rx] = calculate_jump_rate(p.ma_jumps, num_majumps, p.rates, u, - params, t, rx) + params, t, rx, p.maj_rates) sum_rate += cur_rates[rx] end @@ -296,7 +299,7 @@ Perform rejection sampling test (used in RSSA methods). """ @inline function rejectrx( ma_jumps, num_majumps, rates, cur_rate_high, cur_rate_low, rng, u, - jidx, params, t) + jidx, params, t, maj_rates) # rejection test @inbounds r2 = rand(rng) * cur_rate_high[jidx] @inbounds crlow = cur_rate_low[jidx] @@ -305,7 +308,8 @@ Perform rejection sampling test (used in RSSA methods). return false else # calculate actual propensity, split up for type stability - crate = calculate_jump_rate(ma_jumps, num_majumps, rates, u, params, t, jidx) + crate = calculate_jump_rate(ma_jumps, num_majumps, rates, u, params, t, jidx, + maj_rates) if crate > zero(crate) && r2 <= crate return false end diff --git a/src/jumps.jl b/src/jumps.jl index 07185d954..c27038dd8 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -325,7 +325,13 @@ jprob = JumpProblem(prob, Direct(), maj) """ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump - """The (scaled) reaction rate constants.""" + """ + The (scaled) reaction rate constants. When stored within a `JumpProblem`, this vector is + shared across remade problems. Users are responsible for maintaining consistent + stoichiometric scaling and thread safety if mutating directly. In general, prefer using + the parameter index or mapping form for rates that need to change, rather than mutating + `scaled_rates` directly. + """ scaled_rates::T """The reactant stoichiometry vectors.""" reactant_stoch::S @@ -333,7 +339,7 @@ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump net_stoch::U """Parameter mapping functor to identify reaction rate constants with parameters in `p` vectors.""" param_mapper::V - """Whether `update_parameters!` should apply stoichiometric scaling to rates.""" + """Whether the built-in mapper callable should apply stoichiometric scaling to rates.""" rescale_rates_on_update::Bool function MassActionJump{T, S, U, V}(rates::T, rs_in::S, ns::U, pmapper::V, @@ -436,25 +442,18 @@ struct MassActionJumpParamMapper{U} param_idxs::U end -# create the initial parameter vector for use in a MassActionJump -# Note these are unscaled -function (ratemap::MassActionJumpParamMapper{U})(params) where {U <: AbstractArray} - [params[pidx] for pidx in ratemap.param_idxs] -end - -# Note this is unscaled -function (ratemap::MassActionJumpParamMapper{U})(params) where {U <: Int} - params[ratemap.param_idxs] -end - -# update a maj with parameter vectors -function (ratemap::MassActionJumpParamMapper{U})(maj::MassActionJump, newparams; - scale_rates, - kwargs...) where {U <: AbstractArray} - for i in 1:get_num_majumps(maj) - maj.scaled_rates[i] = newparams[ratemap.param_idxs[i]] +# In-place mapper callable API: `(mapper)(dest, maj, params)` +# +# Fill `dest` with the scaled rates for the mass action jump given `params`. +# Custom mappers (e.g. MTK's JumpSysMajParamMapper) should implement this +# 3-arg callable. Use `maj.rescale_rates_on_update` to decide whether to +# apply stoichiometric scaling via `scalerates!`. +function (mapper::MassActionJumpParamMapper{U})(dest::AbstractVector, + maj::MassActionJump, params) where {U <: AbstractArray} + @inbounds for i in eachindex(dest) + dest[i] = params[mapper.param_idxs[i]] end - scale_rates && scalerates!(maj.scaled_rates, maj.reactant_stoch) + maj.rescale_rates_on_update && scalerates!(dest, maj.reactant_stoch) nothing end @@ -479,24 +478,27 @@ function Base.merge(pmap1::MassActionJumpParamMapper{Int}, end """ -update_parameters!(maj::MassActionJump, newparams; scale_rates=maj.rescale_rates_on_update) - -Updates the passed in MassActionJump with the parameter values in `newparams`. + fill_scaled_rates!(dest::AbstractVector, maj::MassActionJump, params) -Notes: +Fill `dest` with the current scaled rates for the mass action jump. Dispatches on the +`scaled_rates` type parameter of `maj`: - - Requires the jump to have been constructed with a user-passed `param_idxs` or `param_mapper`. - - `scale_rates` defaults to `maj.rescale_rates_on_update`, which itself defaults to the - `scale_rates` value used when constructing the jump. When `true`, the parameter - representing the jump rate will be scaled by an appropriate combinatoric factor, i.e. - for 3A --> B at rate k it will scale k --> k/3!. + - `T = Nothing` (parameterized MAJ): delegates to `maj.param_mapper(dest, maj, params)`. + - `T <: AbstractVector` (non-parameterized MAJ): copies `maj.scaled_rates` into `dest`. """ -function update_parameters!(maj::MassActionJump, newparams; scale_rates = maj.rescale_rates_on_update, kwargs...) - (maj.param_mapper === nothing) && - error("MassActionJumps must be constructed with param_idxs or a param_mapper to be updateable.") - maj.param_mapper(maj, newparams; scale_rates, kwargs) +function fill_scaled_rates!(dest::AbstractVector, maj::MassActionJump{Nothing}, params) + maj.param_mapper(dest, maj, params) + nothing +end + +function fill_scaled_rates!(dest::AbstractVector, maj::MassActionJump{T}, + params) where {T <: AbstractVector} + dest .= maj.scaled_rates + nothing end +fill_scaled_rates!(dest, maj::Nothing, params) = nothing + """ $(TYPEDEF) @@ -566,6 +568,11 @@ function JumpSet(vjs, cjs, rj, majv::Vector{T}) where {T <: MassActionJump} error("JumpSets do not accept empty mass action jump collections; use \"nothing\" instead.") end + if any(using_params, majv) + error("Cannot merge MassActionJumps backed by parameter mappers. " * + "Construct a single MassActionJump with all reactions instead.") + end + sr_val = majv[1].rescale_rates_on_update if !all(m -> m.rescale_rates_on_update == sr_val, majv) error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") @@ -724,6 +731,9 @@ massaction_jump_combine(maj1::MassActionJump, maj2::Nothing) = maj1 massaction_jump_combine(maj1::Nothing, maj2::MassActionJump) = maj2 massaction_jump_combine(maj1::Nothing, maj2::Nothing) = maj1 function massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) + (using_params(maj1) || using_params(maj2)) && + error("Cannot merge MassActionJumps backed by parameter mappers. " * + "Construct a single MassActionJump with all reactions instead.") (maj1.rescale_rates_on_update == maj2.rescale_rates_on_update) || error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") majump_merge!(maj1, maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, diff --git a/src/massaction_rates.jl b/src/massaction_rates.jl index 06d4dfa8c..31e639f61 100644 --- a/src/massaction_rates.jl +++ b/src/massaction_rates.jl @@ -4,7 +4,7 @@ ############################################################################### @inline function evalrxrate(speciesvec::AbstractVector{T}, rxidx, - majump::MassActionJump{U})::R where {T <: Integer, R, U <: AbstractVector{R}} + majump::MassActionJump, maj_rates::AbstractVector{R})::R where {T <: Integer, R} val = one(T) @inbounds for specstoch in majump.reactant_stoch[rxidx] specpop = speciesvec[specstoch[1]] @@ -15,11 +15,11 @@ end end - @inbounds return val * majump.scaled_rates[rxidx] + @inbounds return val * maj_rates[rxidx] end @inline function evalrxrate(speciesvec::AbstractVector{T}, rxidx, - majump::MassActionJump{U})::R where {T <: Real, R, U <: AbstractVector{R}} + majump::MassActionJump, maj_rates::AbstractVector{R})::R where {T <: Real, R} val = one(T) @inbounds for specstoch in majump.reactant_stoch[rxidx] specpop = speciesvec[specstoch[1]] @@ -29,11 +29,11 @@ end val *= specpop end # we need to check the smallest rate law term is positive - # i.e. for an order k reaction: x - k + 1 > 0 + # i.e. for an order k reaction: x - k + 1 > 0 (specpop <= 0) && return zero(R) end - @inbounds return val * majump.scaled_rates[rxidx] + @inbounds return val * maj_rates[rxidx] end @inline function executerx!(speciesvec::AbstractVector{T}, rxidx::S, diff --git a/src/problem.jl b/src/problem.jl index 56d58da62..924b218fe 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -157,10 +157,6 @@ function DiffEqBase.remake(jprob::JumpProblem; u0 = missing, p = missing, newprob = DiffEqBase.remake(prob; u0, p, interpret_symbolicmap, use_defaults, kwargs...) end - # if the parameters were changed we must remake the MassActionJump too - if (p !== missing) && using_params(jprob.massaction_jump) - update_parameters!(jprob.massaction_jump, newprob.p; kwargs...) - end else ((u0 !== missing) || (p !== missing) || (:tspan ∈ keys(kwargs))) && error("If remaking a JumpProblem you can not pass both prob and any of u0, p, or tspan.") @@ -169,11 +165,6 @@ function DiffEqBase.remake(jprob::JumpProblem; u0 = missing, p = missing, # when passing a new wrapped problem directly we require u0 has the correct type (typeof(newprob.u0) == typeof(jprob.prob.u0)) || error("The new u0 within the passed prob does not have the same type as the existing u0. Please pass a u0 of type $(typeof(jprob.prob.u0)).") - - # we can't know if p was changed, so we must remake the MassActionJump - if using_params(jprob.massaction_jump) - update_parameters!(jprob.massaction_jump, newprob.p; kwargs...) - end end T(newprob, jprob.aggregator, jprob.discrete_jump_aggregation, jprob.jump_callback, @@ -181,14 +172,6 @@ function DiffEqBase.remake(jprob::JumpProblem; u0 = missing, p = missing, jprob.massaction_jump, jprob.kwargs) end -# for updating parameters in JumpProblems to update MassActionJumps -function SII.finalize_parameters_hook!(prob::JumpProblem, p) - if using_params(prob.massaction_jump) - update_parameters!(prob.massaction_jump, SII.parameter_values(prob)) - end - nothing -end - DiffEqBase.isinplace(::JumpProblem{iip}) where {iip} = iip JumpProblem(prob::JumpProblem) = prob @@ -255,17 +238,9 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS throw(ArgumentError("`rng` is no longer a keyword argument for `JumpProblem`. Pass `rng` to `solve` or `init` instead, e.g. `solve(jprob, SSAStepper(); rng = my_rng)`.")) end - # 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 + # keep the original MAJ (including {Nothing} parameterized MAJs); + # fill_scaled_rates! in each aggregator's initialize! handles rate setup + maj = jumps.massaction_jump ## Spatial jumps handling if spatial_system !== nothing && hopping_constants !== nothing @@ -330,7 +305,6 @@ end function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; save_positions = prob isa DiffEqBase.AbstractDiscreteProblem ? (false, true) : (true, true), - scale_rates = true, useiszero = true, spatial_system = nothing, hopping_constants = nothing, callback = nothing, tstops = nothing, kwargs...) @@ -342,17 +316,9 @@ function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; (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 + # keep the original MAJ (including {Nothing} parameterized MAJs); + # fill_scaled_rates! in the tau-leaping solver handles rate setup + maj = jumps.massaction_jump # For PureLeaping, all jumps are handled by the tau-leaping solver # No discrete jump aggregation or variable rate callbacks are created diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index f71351efc..78b704288 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -262,10 +262,10 @@ function compute_tau( end # Function to generate a mass action rate function -function massaction_rate(maj, numjumps) +function massaction_rate(maj, maj_rates, numjumps) return (out, u, p, t) -> begin for j in 1:numjumps - out[j] = evalrxrate(u, j, maj) + out[j] = evalrxrate(u, j, maj, maj_rates) end end end @@ -355,9 +355,11 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; maj = jump_prob.massaction_jump numjumps = get_num_majumps(maj) + maj_rates = Vector{typeof(tspan[2])}(undef, numjumps) + fill_scaled_rates!(maj_rates, maj, prob.p) rj = jump_prob.regular_jump # Extract rates - rate = rj !== nothing ? rj.rate : massaction_rate(maj, numjumps) + rate = rj !== nothing ? rj.rate : massaction_rate(maj, maj_rates, numjumps) c = rj !== nothing ? rj.c : nothing u0 = copy(prob.u0) p = prob.p diff --git a/src/spatial/directcrdirect.jl b/src/spatial/directcrdirect.jl index 0961ec25b..04e77334c 100644 --- a/src/spatial/directcrdirect.jl +++ b/src/spatial/directcrdirect.jl @@ -142,6 +142,7 @@ function fill_rates_and_get_times!( (; spatial_system, rx_rates, hop_rates, site_rates, rt) = aggregation u = integrator.u + fill_scaled_rates!(rx_rates.maj_rates, rx_rates.ma_jumps, integrator.p) reset!(rx_rates) reset!(hop_rates) site_rates .= zero(typeof(t)) diff --git a/src/spatial/flatten.jl b/src/spatial/flatten.jl index aed6279a2..2b0dbc672 100644 --- a/src/spatial/flatten.jl +++ b/src/spatial/flatten.jl @@ -15,7 +15,9 @@ function flatten(ma_jump, prob::DiscreteProblem, spatial_system, hopping_constan netstoch = ma_jump.net_stoch reactstoch = ma_jump.reactant_stoch rx_rates = if isa(ma_jump, MassActionJump) - ma_jump.scaled_rates + rates_init = Vector{typeof(tspan[1])}(undef, get_num_majumps(ma_jump)) + fill_scaled_rates!(rates_init, ma_jump, prob.p) + rates_init elseif isa(ma_jump, SpatialMassActionJump) num_nodes = num_sites(spatial_system) if isnothing(ma_jump.uniform_rates) && isnothing(ma_jump.spatial_rates) diff --git a/src/spatial/nsm.jl b/src/spatial/nsm.jl index acf6b56eb..e661cefa3 100644 --- a/src/spatial/nsm.jl +++ b/src/spatial/nsm.jl @@ -129,6 +129,7 @@ function fill_rates_and_get_times!(aggregation::NSMJumpAggregation, integrator, rng = get_rng(integrator) u = integrator.u + fill_scaled_rates!(rx_rates.maj_rates, rx_rates.ma_jumps, integrator.p) reset!(rx_rates) reset!(hop_rates) diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl index 737cc5c98..c06858519 100644 --- a/src/spatial/reaction_rates.jl +++ b/src/spatial/reaction_rates.jl @@ -12,6 +12,9 @@ struct RxRates{F, M} "AbstractMassActionJump" ma_jumps::M + + "working copy of scaled mass action jump rates" + maj_rates::Vector{F} end """ @@ -22,7 +25,8 @@ initializes RxRates with zero rates function RxRates(num_sites::Int, ma_jumps::M) where {M} numrxjumps = get_num_majumps(ma_jumps) rates = zeros(Float64, numrxjumps, num_sites) - RxRates{Float64, M}(rates, vec(sum(rates, dims = 1)), ma_jumps) + maj_rates = Vector{Float64}(undef, numrxjumps) + RxRates{Float64, M}(rates, vec(sum(rates, dims = 1)), ma_jumps, maj_rates) end num_rxs(rx_rates::RxRates) = get_num_majumps(rx_rates.ma_jumps) @@ -55,8 +59,9 @@ update rates of all reactions in rxs at site function update_rx_rates!(rx_rates::RxRates, rxs, u::AbstractMatrix, integrator, site) ma_jumps = rx_rates.ma_jumps + maj_rates = rx_rates.maj_rates @inbounds for rx in rxs - rate = eval_massaction_rate(u, rx, ma_jumps, site) + rate = eval_massaction_rate(u, rx, ma_jumps, site, maj_rates) set_rx_rate_at_site!(rx_rates, site, rx, rate) end end @@ -90,9 +95,9 @@ function Base.show(io::IO, ::MIME"text/plain", rx_rates::RxRates) println(io, "RxRates with $num_rxs reactions and $num_sites sites") end -function eval_massaction_rate(u, rx, ma_jumps::M, site) where {M <: SpatialMassActionJump} +function eval_massaction_rate(u, rx, ma_jumps::M, site, maj_rates) where {M <: SpatialMassActionJump} evalrxrate(u, rx, ma_jumps, site) end -function eval_massaction_rate(u, rx, ma_jumps::M, site) where {M <: MassActionJump} - evalrxrate((@view u[:, site]), rx, ma_jumps) +function eval_massaction_rate(u, rx, ma_jumps::M, site, maj_rates) where {M <: MassActionJump} + evalrxrate((@view u[:, site]), rx, ma_jumps, maj_rates) end diff --git a/src/spatial/spatial_massaction_jump.jl b/src/spatial/spatial_massaction_jump.jl index c9af9d1e2..52bcc7929 100644 --- a/src/spatial/spatial_massaction_jump.jl +++ b/src/spatial/spatial_massaction_jump.jl @@ -82,6 +82,11 @@ function SpatialMassActionJump(urates::A, rs, ns; scale_rates = true, useiszero useiszero = useiszero, nocopy = nocopy) end +function SpatialMassActionJump(ma_jumps::MassActionJump{Nothing}; kwargs...) + error("Cannot construct SpatialMassActionJump from a parameter-mapped MassActionJump. " * + "Provide explicit rate vectors or use a SpatialMassActionJump directly.") +end + # scale_rates defaults to false since ma_jumps.scaled_rates are already scaled; # passing true would double-scale. function SpatialMassActionJump(ma_jumps::MassActionJump{T, S, U, V}; scale_rates = false, diff --git a/test/jprob_symbol_indexing.jl b/test/jprob_symbol_indexing.jl index 4cad845cb..9419c3c51 100644 --- a/test/jprob_symbol_indexing.jl +++ b/test/jprob_symbol_indexing.jl @@ -24,16 +24,25 @@ jprob = JumpProblem(dprob, Direct(), crj1, crj2, maj) jprob[:a] = 20 @test jprob[:a] == 20 -# test mass action jumps update with parameter mutation in problems -@test jprob.massaction_jump.scaled_rates[1] == 1.0 +# test parameterized MAJ stores nothing for scaled_rates +@test jprob.massaction_jump.scaled_rates === nothing + +# test mass action jump rates update correctly after init +integ = init(jprob, SSAStepper()) +@test jprob.discrete_jump_aggregation.maj_rates == [1.0, 2.0] + +# test parameter mutation updates problem params jprob.ps[:p1] = 3.0 @test jprob.ps[:p1] == 3.0 -@test jprob.massaction_jump.scaled_rates[1] == 3.0 +# rates update after re-initialization +integ = init(jprob, SSAStepper()) +@test jprob.discrete_jump_aggregation.maj_rates[1] == 3.0 p1setter = setp(jprob, [:p1, :p2]) p1setter(jprob, [4.0, 10.0]) @test jprob.ps[:p1] == 4.0 @test jprob.ps[:p2] == 10.0 -@test jprob.massaction_jump.scaled_rates == [4.0, 10.0] +integ = init(jprob, SSAStepper()) +@test jprob.discrete_jump_aggregation.maj_rates == [4.0, 10.0] # integrator tests # note that `setu` is not currently supported as `set_u!` is not implemented for SSAStepper @@ -44,9 +53,8 @@ integ[[:b, :a]] = [40, 5] @test getp(integ, :p2)(integ) == 10.0 setp(integ, :p2)(integ, 15.0) @test getp(integ, :p2)(integ) == 15.0 -@test jprob.massaction_jump.scaled_rates[2] == 10.0 # jump rate not updated reset_aggregated_jumps!(integ) -@test jprob.massaction_jump.scaled_rates[2] == 15.0 # jump rate now updated +@test jprob.discrete_jump_aggregation.maj_rates[2] == 15.0 # jump rate now updated # remake tests dprob = DiscreteProblem(g, [0, 10], (0.0, 10.0), [1.0, 2.0]) @@ -54,19 +62,22 @@ jprob = JumpProblem(dprob, Direct(), crj1, crj2, maj) jprob = remake(jprob; u0 = [:a => -10, :b => 100], p = [:p2 => 3.5, :p1 => 0.5]) @test jprob.prob.u0 == [-10, 100] @test jprob.prob.p == [0.5, 3.5] -@test jprob.massaction_jump.scaled_rates == [0.5, 3.5] +integ = init(jprob, SSAStepper()) +@test jprob.discrete_jump_aggregation.maj_rates == [0.5, 3.5] jprob = remake(jprob; u0 = [:b => 10], p = [:p2 => 4.5]) @test jprob.prob.u0 == [-10, 10] @test jprob.prob.p == [0.5, 4.5] -@test jprob.massaction_jump.scaled_rates == [0.5, 4.5] +integ = init(jprob, SSAStepper()) +@test jprob.discrete_jump_aggregation.maj_rates == [0.5, 4.5] -# test updating problems via regular indexing still updates the mass action jump +# test updating problems via regular indexing dprob = DiscreteProblem(g, [0, 10], (0.0, 10.0), [1.0, 2.0]) jprob = JumpProblem(dprob, Direct(), crj1, crj2, maj) -@test jprob.massaction_jump.scaled_rates[1] == 1.0 +@test jprob.massaction_jump.scaled_rates === nothing jprob.ps[1] = 3.0 @test jprob.ps[1] == 3.0 -@test jprob.massaction_jump.scaled_rates[1] == 3.0 +integ = init(jprob, SSAStepper()) +@test jprob.discrete_jump_aggregation.maj_rates[1] == 3.0 # test updating integrators via regular indexing dprob = DiscreteProblem(g, [0, 10], (0.0, 10.0), [1.0, 2.0]) @@ -76,12 +87,12 @@ integ.u .= [40, 5] @test getu(integ, [1, 2])(integ) == [40, 5] @test getp(integ, 2)(integ) == 2.0 @test integ.p[2] == 2.0 -@test jprob.massaction_jump.scaled_rates[2] == 2.0 +@test jprob.discrete_jump_aggregation.maj_rates[2] == 2.0 setp(integ, 2)(integ, 15.0) @test integ.p[2] == 15.0 @test getp(integ, 2)(integ) == 15.0 reset_aggregated_jumps!(integ) -@test jprob.massaction_jump.scaled_rates[2] == 15.0 # jump rate now updated +@test jprob.discrete_jump_aggregation.maj_rates[2] == 15.0 # jump rate now updated # remake tests for regular indexing dprob = DiscreteProblem(g, [0, 10], (0.0, 10.0), [1.0, 2.0]) @@ -89,8 +100,10 @@ jprob = JumpProblem(dprob, Direct(), crj1, crj2, maj) jprob = remake(jprob; u0 = [-10, 100], p = [0.5, 3.5]) @test jprob.prob.u0 == [-10, 100] @test jprob.prob.p == [0.5, 3.5] -@test jprob.massaction_jump.scaled_rates == [0.5, 3.5] +integ = init(jprob, SSAStepper()) +@test jprob.discrete_jump_aggregation.maj_rates == [0.5, 3.5] jprob = remake(jprob; u0 = [2 => 10], p = [2 => 4.5]) @test jprob.prob.u0 == [-10, 10] @test jprob.prob.p == [0.5, 4.5] -@test jprob.massaction_jump.scaled_rates == [0.5, 4.5] +integ = init(jprob, SSAStepper()) +@test jprob.discrete_jump_aggregation.maj_rates == [0.5, 4.5] diff --git a/test/regular_jumps.jl b/test/regular_jumps.jl index 99c3375c7..6cbc28040 100644 --- a/test/regular_jumps.jl +++ b/test/regular_jumps.jl @@ -251,8 +251,10 @@ end # 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 + @test jp_params.massaction_jump.scaled_rates === nothing + dest = zeros(2) + JumpProcesses.fill_scaled_rates!(dest, jp_params.massaction_jump, p) + @test dest == [p[1], p[2] / 2] end # Test that saveat/save_start/save_end control which times are stored in solutions diff --git a/test/scale_rates_field_test.jl b/test/scale_rates_field_test.jl index 3f06ed888..a5c961070 100644 --- a/test/scale_rates_field_test.jl +++ b/test/scale_rates_field_test.jl @@ -7,12 +7,10 @@ reactant_stoch = [[1 => 3]] net_stoch = [[1 => -3, 2 => 1]] # Custom mapper mimicking MTKBase's JumpSysMajParamMapper. -# The initial rates callable always pre-scales (simulating symbolic expressions that -# already contain the combinatoric factor, e.g. k/3! for 3X → Y). -# The update callable also pre-scales, then conditionally applies scalerates! based on -# scale_rates — matching MTKBase's behavior. This is what makes the test an actual -# regression test: with the old default of scale_rates=true, the second scalerates! -# would double-scale. +# The 1-arg callable pre-scales rates (simulating symbolic expressions that already +# contain the combinatoric factor, e.g. k/3! for 3X → Y). +# The 3-arg in-place callable also pre-scales, then conditionally applies standard +# stoichiometric scaling based on maj.rescale_rates_on_update. struct PreScaledMapper param_idxs::Vector{Int} reactant_stoch::Vector{Vector{Pair{Int, Int}}} @@ -22,12 +20,12 @@ function (m::PreScaledMapper)(params) JumpProcesses.scalerates!(rates, m.reactant_stoch) rates end -function (m::PreScaledMapper)(maj::MassActionJump, newparams; scale_rates, kwargs...) - for i in 1:JumpProcesses.get_num_majumps(maj) - maj.scaled_rates[i] = newparams[m.param_idxs[i]] +function (m::PreScaledMapper)(dest::AbstractVector, maj::MassActionJump, params) + @inbounds for i in eachindex(dest) + dest[i] = params[m.param_idxs[i]] end - JumpProcesses.scalerates!(maj.scaled_rates, m.reactant_stoch) - scale_rates && JumpProcesses.scalerates!(maj.scaled_rates, maj.reactant_stoch) + JumpProcesses.scalerates!(dest, m.reactant_stoch) + maj.rescale_rates_on_update && JumpProcesses.scalerates!(dest, maj.reactant_stoch) nothing end JumpProcesses.to_collection(m::PreScaledMapper) = m @@ -53,19 +51,25 @@ end @test maj.rescale_rates_on_update == true end -# Test 2: update_parameters! respects stored rescale_rates_on_update -@testset "update_parameters! respects rescale_rates_on_update" begin +# Test 2: fill_scaled_rates! respects rescale_rates_on_update +@testset "fill_scaled_rates! respects rescale_rates_on_update" begin # With param_idxs and scale_rates = true (default) — built-in mapper path p = [6.0] maj = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) - dprob = DiscreteProblem([100, 0], (0.0, 1.0), p) - jprob = JumpProblem(dprob, Direct(), maj) - @test jprob.massaction_jump.rescale_rates_on_update == true - @test jprob.massaction_jump.scaled_rates[1] ≈ 1.0 # 6.0 / 3! + @test maj.rescale_rates_on_update == true + + dest = zeros(1) + JumpProcesses.fill_scaled_rates!(dest, maj, p) + @test dest[1] ≈ 1.0 # 6.0 / 3! - # update_parameters! should scale (rescale_rates_on_update = true from struct) - JumpProcesses.update_parameters!(jprob.massaction_jump, [12.0]) - @test jprob.massaction_jump.scaled_rates[1] ≈ 2.0 # 12.0 / 3! + # fill_scaled_rates! with new params + JumpProcesses.fill_scaled_rates!(dest, maj, [12.0]) + @test dest[1] ≈ 2.0 # 12.0 / 3! + + # Non-parameterized MAJ: fill_scaled_rates! copies stored rates + maj_explicit = MassActionJump([6.0], reactant_stoch, net_stoch) + JumpProcesses.fill_scaled_rates!(dest, maj_explicit, p) + @test dest[1] ≈ 1.0 # copies maj.scaled_rates which is 6.0/3! = 1.0 end # Test 3: Custom pre-scaled mapper with scale_rates = false — the bug reproducer @@ -79,24 +83,31 @@ end jprob = JumpProblem(dprob, Direct(), maj; scale_rates = false) expected_scaled = k / factorial(3) # 1.0 - @test jprob.massaction_jump.scaled_rates[1] ≈ expected_scaled + + # parameterized MAJ stores nothing for scaled_rates + @test jprob.massaction_jump.scaled_rates === nothing @test jprob.massaction_jump.rescale_rates_on_update == false - # Test reset_aggregated_jumps! does NOT double-scale + # rates are materialized after init triggers initialize! integ = init(jprob, SSAStepper()) + @test jprob.discrete_jump_aggregation.maj_rates[1] ≈ expected_scaled + + # Test reset_aggregated_jumps! does NOT double-scale integ.p[1] = 18.0 reset_aggregated_jumps!(integ) new_expected = 18.0 / factorial(3) # 3.0, NOT 3.0/6 = 0.5 - @test jprob.massaction_jump.scaled_rates[1] ≈ new_expected + @test jprob.discrete_jump_aggregation.maj_rates[1] ≈ new_expected - # Test remake does NOT double-scale + # Test remake + init materializes rates correctly jprob2 = remake(jprob; p = [24.0]) - remake_expected = 24.0 / factorial(3) # 4.0, NOT 4.0/6 ≈ 0.667 - @test jprob2.massaction_jump.scaled_rates[1] ≈ remake_expected + init(jprob2, SSAStepper()) + remake_expected = 24.0 / factorial(3) # 4.0 + @test jprob2.discrete_jump_aggregation.maj_rates[1] ≈ remake_expected # Test remake round-trip jprob3 = remake(jprob2; p = [k]) - @test jprob3.massaction_jump.scaled_rates[1] ≈ expected_scaled + init(jprob3, SSAStepper()) + @test jprob3.discrete_jump_aggregation.maj_rates[1] ≈ expected_scaled end # Test 4: Callback parameter changes with built-in mapper (rescale_rates_on_update = true) @@ -113,7 +124,7 @@ end end cb = DiscreteCallback(condit, affect!) sol = solve(jprob, SSAStepper(); tstops = [1000.0], callback = cb) - @test jprob.massaction_jump.scaled_rates[1] ≈ 4.0 # 24.0 / 3! + @test jprob.discrete_jump_aggregation.maj_rates[1] ≈ 4.0 # 24.0 / 3! end # Test 5: rescale_rates_on_update propagated through JumpSet merge and JumpProblem varargs @@ -161,3 +172,31 @@ end # Mismatched rescale_rates_on_update via JumpProblem varargs — should error @test_throws ErrorException JumpProblem(dprob, Direct(), maj_true, maj_false) end + +# Test 6: Mapper-backed MAJ merge raises error +@testset "Parameterized MAJ merge error" begin + maj_p1 = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) + maj_p2 = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) + @test_throws ErrorException JumpSet(; massaction_jumps = [maj_p1, maj_p2]) + + dprob = DiscreteProblem([100, 0], (0.0, 1.0), [1.0]) + @test_throws ErrorException JumpProblem(dprob, Direct(), maj_p1, maj_p2) +end + +# Test 7: Immutability and aliasing +@testset "Immutability and aliasing after remake" begin + p = [6.0] + maj = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) + dprob = DiscreteProblem([100, 0], (0.0, 100.0), p) + jprob = JumpProblem(dprob, Direct(), maj) + init(jprob, SSAStepper()) + rates_before = copy(jprob.discrete_jump_aggregation.maj_rates) + + # remake with new p, then init — original MAJ is not mutated + jprob2 = remake(jprob; p = [12.0]) + @test jprob.massaction_jump === jprob2.massaction_jump # shared MAJ + @test jprob.massaction_jump.scaled_rates === nothing # still nothing + init(jprob2, SSAStepper()) + # aggregation is shared, so maj_rates reflects latest init + @test jprob2.discrete_jump_aggregation.maj_rates[1] ≈ 2.0 # 12.0 / 3! +end diff --git a/test/ssa_callback_test.jl b/test/ssa_callback_test.jl index 1cbf678b3..10115a5ac 100644 --- a/test/ssa_callback_test.jl +++ b/test/ssa_callback_test.jl @@ -105,14 +105,19 @@ sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCall @test all(p .== [0.0, 1.0]) @test sol[1, end] == 100 -# test scale_rates kwarg +# test scale_rates / rescale_rates_on_update behavior p .= [1.0] dprob = DiscreteProblem(u₀, tspan, p) maj5 = MassActionJump([[1 => 2]], [[1 => -1, 2 => 1]]; param_idxs = [1]) jprob = JumpProblem(dprob, Direct(), maj5, save_positions = (false, false)) -@test all(jprob.massaction_jump.scaled_rates .== [0.5]) -jprob = JumpProblem(dprob, Direct(), maj5, save_positions = (false, false), scale_rates = false) -@test all(jprob.massaction_jump.scaled_rates .== [1.0]) +@test jprob.massaction_jump.scaled_rates === nothing # parameterized MAJ +integ = init(jprob, SSAStepper()) +@test all(jprob.discrete_jump_aggregation.maj_rates .== [0.5]) +# with scale_rates = false on the MAJ itself +maj5_noscale = MassActionJump([[1 => 2]], [[1 => -1, 2 => 1]]; param_idxs = [1], scale_rates = false) +jprob = JumpProblem(dprob, Direct(), maj5_noscale, save_positions = (false, false)) +integ = init(jprob, SSAStepper()) +@test all(jprob.discrete_jump_aggregation.maj_rates .== [1.0]) # test for https://github.com/SciML/JumpProcesses.jl/issues/239 maj6 = MassActionJump([[1 => 1], [2 => 1]], [[1 => -1, 2 => 1], [1 => 1, 2 => -1]]; From a7dfc6bfaddeac1fcbb8f8a064bf7c95ab448c52 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 3 Mar 2026 16:36:06 -0500 Subject: [PATCH 22/52] Fix test failures: allow built-in mapper merges, fix spatial evalrxrate - Narrow MAJ merge guards to only block custom (non-MassActionJumpParamMapper) mapper types. Built-in parameterized MAJ merges via varargs and JumpSet vectors work correctly and are tested in ssa_callback_test.jl. - Fix 3-arg evalrxrate call in test/spatial/reaction_rates.jl to use eval_massaction_rate which handles both MassActionJump and SpatialMassActionJump correctly. - Add fill_scaled_rates! call before spatial rate tests to populate maj_rates. - Update scale_rates_field_test.jl merge test to verify built-in merges succeed and custom mapper merges error. Co-Authored-By: Claude Opus 4.5 --- src/jumps.jl | 12 +++++++----- test/scale_rates_field_test.jl | 22 ++++++++++++++++------ test/spatial/reaction_rates.jl | 5 ++++- 3 files changed, 27 insertions(+), 12 deletions(-) diff --git a/src/jumps.jl b/src/jumps.jl index c27038dd8..708748876 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -568,8 +568,8 @@ function JumpSet(vjs, cjs, rj, majv::Vector{T}) where {T <: MassActionJump} error("JumpSets do not accept empty mass action jump collections; use \"nothing\" instead.") end - if any(using_params, majv) - error("Cannot merge MassActionJumps backed by parameter mappers. " * + if any(m -> using_params(m) && !(m.param_mapper isa MassActionJumpParamMapper), majv) + error("Cannot merge MassActionJumps backed by custom parameter mappers. " * "Construct a single MassActionJump with all reactions instead.") end @@ -731,9 +731,11 @@ massaction_jump_combine(maj1::MassActionJump, maj2::Nothing) = maj1 massaction_jump_combine(maj1::Nothing, maj2::MassActionJump) = maj2 massaction_jump_combine(maj1::Nothing, maj2::Nothing) = maj1 function massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) - (using_params(maj1) || using_params(maj2)) && - error("Cannot merge MassActionJumps backed by parameter mappers. " * - "Construct a single MassActionJump with all reactions instead.") + for m in (maj1, maj2) + (using_params(m) && !(m.param_mapper isa MassActionJumpParamMapper)) && + error("Cannot merge MassActionJumps backed by custom parameter mappers. " * + "Construct a single MassActionJump with all reactions instead.") + end (maj1.rescale_rates_on_update == maj2.rescale_rates_on_update) || error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") majump_merge!(maj1, maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, diff --git a/test/scale_rates_field_test.jl b/test/scale_rates_field_test.jl index a5c961070..ee5be07bf 100644 --- a/test/scale_rates_field_test.jl +++ b/test/scale_rates_field_test.jl @@ -128,8 +128,7 @@ end end # Test 5: rescale_rates_on_update propagated through JumpSet merge and JumpProblem varargs -# Use explicit-rate MAJs since the JumpSet vector merge path doesn't support -# parameterized (Nothing-rated) MAJs. +# Uses explicit-rate MAJs to test rescale_rates_on_update propagation and mismatch errors. @testset "rescale_rates_on_update propagated through merge paths" begin reactant_stoch2 = [[2 => 3]] net_stoch2 = [[2 => -3, 1 => 1]] @@ -173,14 +172,25 @@ end @test_throws ErrorException JumpProblem(dprob, Direct(), maj_true, maj_false) end -# Test 6: Mapper-backed MAJ merge raises error -@testset "Parameterized MAJ merge error" begin +# Test 6: Custom mapper-backed MAJ merge raises error; built-in mapper merge works +@testset "MAJ merge behavior" begin + # Built-in MassActionJumpParamMapper merges should succeed maj_p1 = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) maj_p2 = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) - @test_throws ErrorException JumpSet(; massaction_jumps = [maj_p1, maj_p2]) + merged = JumpSet(; massaction_jumps = [maj_p1, maj_p2]) + @test JumpProcesses.get_num_majumps(merged.massaction_jump) == 2 dprob = DiscreteProblem([100, 0], (0.0, 1.0), [1.0]) - @test_throws ErrorException JumpProblem(dprob, Direct(), maj_p1, maj_p2) + jprob = JumpProblem(dprob, Direct(), maj_p1, maj_p2) + @test JumpProcesses.get_num_majumps(jprob.massaction_jump) == 2 + + # Custom mapper merge should error + mapper1 = PreScaledMapper([1], reactant_stoch) + mapper2 = PreScaledMapper([1], reactant_stoch) + maj_c1 = MassActionJump(reactant_stoch, net_stoch; param_mapper = mapper1, scale_rates = false) + maj_c2 = MassActionJump(reactant_stoch, net_stoch; param_mapper = mapper2, scale_rates = false) + @test_throws ErrorException JumpSet(; massaction_jumps = [maj_c1, maj_c2]) + @test_throws ErrorException JumpProblem(dprob, Direct(), maj_c1, maj_c2) end # Test 7: Immutability and aliasing diff --git a/test/spatial/reaction_rates.jl b/test/spatial/reaction_rates.jl index 8199dd39e..2e6b5ac39 100644 --- a/test/spatial/reaction_rates.jl +++ b/test/spatial/reaction_rates.jl @@ -36,10 +36,13 @@ rx_rates_list = [JP.RxRates(num_nodes, ma_jumps), JP.RxRates(num_nodes, spatial_ for rx_rates in rx_rates_list @test JP.num_rxs(rx_rates) == length(rates) show(io, "text/plain", rx_rates) + if rx_rates.ma_jumps isa MassActionJump + JP.fill_scaled_rates!(rx_rates.maj_rates, rx_rates.ma_jumps, nothing) + end for site in 1:num_nodes JP.update_rx_rates!(rx_rates, 1:num_rxs, integrator, site) @test JP.total_site_rx_rate(rx_rates, site) == 1.1 - rx_props = [JP.evalrxrate(u[:, site], rx, ma_jumps) for rx in 1:num_rxs] + rx_props = [JP.eval_massaction_rate(u, rx, rx_rates.ma_jumps, site, rx_rates.maj_rates) for rx in 1:num_rxs] rx_probs = rx_props / sum(rx_props) d = Dict{Int, Int}() for i in 1:num_samples From 0e0221467de9b7d79c2f419480611ddccaccec80 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 3 Mar 2026 17:17:23 -0500 Subject: [PATCH 23/52] Fix merge aliasing bug and remove deprecated JumpProblem kwargs - setup_majump_to_merge now copies arrays before creating the merge accumulator, preventing in-place mutation from corrupting the original MAJ net_stoch/reactant_stoch arrays - massaction_jump_combine(::Nothing, ::MassActionJump) now creates a safe copy via setup_majump_to_merge instead of returning the original MAJ - Remove scale_rates and useiszero as JumpProblem kwargs; they are now properties of the MassActionJump itself. Passing them raises an informative ArgumentError Co-Authored-By: Claude Opus 4.5 --- src/jumps.jl | 23 ++++++++++++++++++++--- src/problem.jl | 7 ++++++- test/scale_rates_field_test.jl | 8 +++++--- 3 files changed, 31 insertions(+), 7 deletions(-) diff --git a/src/jumps.jl b/src/jumps.jl index 708748876..5d4f5f317 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -460,6 +460,7 @@ end function to_collection(ratemap::MassActionJumpParamMapper{Int}) MassActionJumpParamMapper([ratemap.param_idxs]) end +to_collection(ratemap::MassActionJumpParamMapper{<:AbstractVector}) = ratemap function Base.merge!(pmap1::MassActionJumpParamMapper{U}, pmap2::MassActionJumpParamMapper{U}) where {U <: AbstractVector} @@ -647,10 +648,12 @@ function check_majump_type(maj::MassActionJump{Nothing, T, U, V}) where {T, U, V end # if given containers of rates and stoichiometry directly create a jump +# copy arrays since majump_merge! will mutate them in-place function setup_majump_to_merge(sr::T, rs::AbstractVector{S}, ns::AbstractVector{U}, pmapper, rescale_rates_on_update::Bool) where {T <: AbstractVector, S <: AbstractArray, U <: AbstractArray} - MassActionJump(sr, rs, ns, pmapper; scale_rates = false, rescale_rates_on_update) + MassActionJump(copy(sr), copy(rs), copy(ns), pmapper; + scale_rates = false, nocopy = true, rescale_rates_on_update) end # if just given the data for one jump (and not in a container) wrap in a vector @@ -662,7 +665,17 @@ function setup_majump_to_merge(sr::S, rs::T, ns::U, pmapper, scale_rates = false, rescale_rates_on_update) end -# if no rate field setup yet +# if no rate field setup yet — collection case (rs is already Vector{<:AbstractArray}) +# copy arrays and mapper since majump_merge! will mutate them in-place +function setup_majump_to_merge(::Nothing, rs::AbstractVector{S}, ns::AbstractVector{U}, + pmapper, + rescale_rates_on_update::Bool) where {S <: AbstractArray, U <: AbstractArray} + pm = (pmapper === nothing) ? nothing : deepcopy(pmapper) + MassActionJump(nothing, copy(rs), copy(ns), pm; + scale_rates = false, nocopy = true, rescale_rates_on_update) +end + +# if no rate field setup yet — single reaction case (wrap in a vector) function setup_majump_to_merge(::Nothing, rs::T, ns::U, pmapper, rescale_rates_on_update::Bool) where {T <: AbstractArray, U <: AbstractArray} MassActionJump(nothing, [rs], [ns], @@ -728,8 +741,12 @@ function majump_merge!(maj::MassActionJump{T, S, U, V}, sr::T, rs::S, ns::U, end massaction_jump_combine(maj1::MassActionJump, maj2::Nothing) = maj1 -massaction_jump_combine(maj1::Nothing, maj2::MassActionJump) = maj2 massaction_jump_combine(maj1::Nothing, maj2::Nothing) = maj1 +# copy the MAJ so it is safe to mutate during subsequent merges +function massaction_jump_combine(maj1::Nothing, maj2::MassActionJump) + setup_majump_to_merge(maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, + maj2.param_mapper, maj2.rescale_rates_on_update) +end function massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) for m in (maj1, maj2) (using_params(m) && !(m.param_mapper isa MassActionJumpParamMapper)) && diff --git a/src/problem.jl b/src/problem.jl index 924b218fe..05838c463 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -230,13 +230,18 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS vr_aggregator::VariableRateAggregator = VR_FRM(), save_positions = prob isa DiffEqBase.AbstractDiscreteProblem ? (false, true) : (true, true), - scale_rates = true, useiszero = true, spatial_system = nothing, hopping_constants = nothing, callback = nothing, tstops = nothing, use_vrj_bounds = true, kwargs...) if haskey(kwargs, :rng) throw(ArgumentError("`rng` is no longer a keyword argument for `JumpProblem`. Pass `rng` to `solve` or `init` instead, e.g. `solve(jprob, SSAStepper(); rng = my_rng)`.")) end + if haskey(kwargs, :scale_rates) + throw(ArgumentError("`scale_rates` is no longer a keyword argument for `JumpProblem`. Set `scale_rates` on the `MassActionJump` directly instead.")) + end + if haskey(kwargs, :useiszero) + throw(ArgumentError("`useiszero` is no longer a keyword argument for `JumpProblem`. Set `useiszero` on the `MassActionJump` directly instead.")) + end # keep the original MAJ (including {Nothing} parameterized MAJs); # fill_scaled_rates! in each aggregator's initialize! handles rate setup diff --git a/test/scale_rates_field_test.jl b/test/scale_rates_field_test.jl index ee5be07bf..7044659ed 100644 --- a/test/scale_rates_field_test.jl +++ b/test/scale_rates_field_test.jl @@ -80,7 +80,7 @@ end maj = MassActionJump(reactant_stoch, net_stoch; param_mapper = mapper, scale_rates = false) dprob = DiscreteProblem([100, 0], (0.0, 100.0), p) - jprob = JumpProblem(dprob, Direct(), maj; scale_rates = false) + jprob = JumpProblem(dprob, Direct(), maj) expected_scaled = k / factorial(3) # 1.0 @@ -159,7 +159,7 @@ end # Two MAJs with matching rescale_rates_on_update = false via JumpProblem varargs maj_f1 = MassActionJump([1.0], reactant_stoch, net_stoch; scale_rates = false) maj_f2 = MassActionJump([2.0], reactant_stoch2, net_stoch2; scale_rates = false) - jprob_f = JumpProblem(dprob, Direct(), maj_f1, maj_f2; scale_rates = false) + jprob_f = JumpProblem(dprob, Direct(), maj_f1, maj_f2) @test jprob_f.massaction_jump.rescale_rates_on_update == false # Two MAJs with matching rescale_rates_on_update = true via JumpProblem varargs @@ -181,7 +181,9 @@ end @test JumpProcesses.get_num_majumps(merged.massaction_jump) == 2 dprob = DiscreteProblem([100, 0], (0.0, 1.0), [1.0]) - jprob = JumpProblem(dprob, Direct(), maj_p1, maj_p2) + maj_p3 = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) + maj_p4 = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1]) + jprob = JumpProblem(dprob, Direct(), maj_p3, maj_p4) @test JumpProcesses.get_num_majumps(jprob.massaction_jump) == 2 # Custom mapper merge should error From 95f6f0824a93d01974b2c4d4f5751ed7afebb638 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Wed, 4 Mar 2026 12:46:46 -0500 Subject: [PATCH 24/52] Add v10_working branch to CI triggers and pin OrdinaryDiffEqCore Co-Authored-By: Claude Opus 4.5 --- .github/workflows/Downgrade.yml | 1 + .github/workflows/Tests.yml | 1 + .github/workflows/ThreadSafety.yml | 1 + Project.toml | 2 +- 4 files changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml index 54db9bada..3fe5b7a7b 100644 --- a/.github/workflows/Downgrade.yml +++ b/.github/workflows/Downgrade.yml @@ -3,6 +3,7 @@ on: pull_request: branches: - master + - v10_working paths-ignore: - 'docs/**' push: diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index dd8ce04f5..3d0b7ef34 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -8,6 +8,7 @@ on: pull_request: branches: - master + - v10_working paths-ignore: - 'docs/**' diff --git a/.github/workflows/ThreadSafety.yml b/.github/workflows/ThreadSafety.yml index b82eda2db..2dfbb4643 100644 --- a/.github/workflows/ThreadSafety.yml +++ b/.github/workflows/ThreadSafety.yml @@ -9,6 +9,7 @@ on: pull_request: branches: - master + - v10_working paths-ignore: - 'docs/**' diff --git a/Project.toml b/Project.toml index 242da5aee..2c6fc33a2 100644 --- a/Project.toml +++ b/Project.toml @@ -45,7 +45,7 @@ KernelAbstractions = "0.9" LinearAlgebra = "1" LinearSolve = "3" OrdinaryDiffEq = "6" -OrdinaryDiffEqCore = "3.11" +OrdinaryDiffEqCore = "=3.15" Pkg = "1" PoissonRandom = "0.4" Random = "1" From 90e7fdce41c0cb4bab2580653d217cff4aa3a36d Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Fri, 20 Mar 2026 15:00:27 -0400 Subject: [PATCH 25/52] Fix CI failures from master merge: update tests for v10 API changes - Move rng kwarg from JumpProblem to solve/init in merged test files (extended_jump_array.jl, callbacks.jl) - Remove stale seed! import from JumpProcesses.jl (QA failure) - Add fill_scaled_rates! in spatial bracketing test to populate maj_rates - Remove old prob_func(prob, i, repeat) signature in thread_safety.jl (use default EnsembleProblem which already deepcopies) Co-Authored-By: Claude Opus 4.5 --- src/JumpProcesses.jl | 2 +- test/callbacks.jl | 8 ++++---- test/extended_jump_array.jl | 4 ++-- test/spatial/bracketing.jl | 2 ++ test/thread_safety.jl | 6 ++---- 5 files changed, 11 insertions(+), 11 deletions(-) diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl index e32c4fb8e..2c1ea6e9e 100644 --- a/src/JumpProcesses.jl +++ b/src/JumpProcesses.jl @@ -5,7 +5,7 @@ using Reexport: Reexport, @reexport # Explicit imports from standard libraries using LinearAlgebra: LinearAlgebra, mul! -using Random: Random, randexp, seed! +using Random: Random, randexp # Explicit imports from external packages using DocStringExtensions: DocStringExtensions, FIELDS, TYPEDEF diff --git a/test/callbacks.jl b/test/callbacks.jl index b0b796aff..4fe727325 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -490,9 +490,9 @@ end crj = ConstantRateJump(rate, affect!) sde_prob = SDEProblem(f, g, 1.0, (0.0, 1.0)) - jprob = JumpProblem(sde_prob, Direct(), crj; rng) + jprob = JumpProblem(sde_prob, Direct(), crj) - integrator = init(jprob, EM(); dt = 0.01) + integrator = init(jprob, EM(); dt = 0.01, rng) # Count discrete callbacks — the jump should appear exactly once. # If both JumpProcesses and StochasticDiffEq add it, we'd see duplicates. @@ -508,8 +508,8 @@ end vaffect!(integrator) = (integrator.u[1] += 0.5) vrj = VariableRateJump(vrate, vaffect!) - jprob_vr = JumpProblem(sde_prob_vr, Direct(), vrj; vr_aggregator = VR_FRM(), rng) - integrator_vr = init(jprob_vr, EM(); dt = 0.01) + jprob_vr = JumpProblem(sde_prob_vr, Direct(), vrj; vr_aggregator = VR_FRM()) + integrator_vr = init(jprob_vr, EM(); dt = 0.01, rng) # VariableRateJumps produce continuous callbacks n_continuous = length(integrator_vr.opts.callback.continuous_callbacks) diff --git a/test/extended_jump_array.jl b/test/extended_jump_array.jl index 9f6d59a5e..b0789da17 100644 --- a/test/extended_jump_array.jl +++ b/test/extended_jump_array.jl @@ -154,8 +154,8 @@ let affect!(integrator) = (integrator.u[1] += 1; nothing) vrj = VariableRateJump(rate, affect!) oprob = ODEProblem(f!, [0.0], (0.0, 1.0)) - jprob = JumpProblem(oprob, Direct(), vrj; vr_aggregator = VR_FRM(), + jprob = JumpProblem(oprob, Direct(), vrj; vr_aggregator = VR_FRM()) + sol = solve(jprob, Rodas5P(linsolve = QRFactorization()); rng = StableRNG(789)) - sol = solve(jprob, Rodas5P(linsolve = QRFactorization())) @test sol.retcode == ReturnCode.Success end diff --git a/test/spatial/bracketing.jl b/test/spatial/bracketing.jl index 31c1e23bf..7afcaba87 100644 --- a/test/spatial/bracketing.jl +++ b/test/spatial/bracketing.jl @@ -19,6 +19,8 @@ netstoch = [[1 => -1]] majump = MassActionJump(majump_rates, reactstoch, netstoch) rx_rates = JP.LowHigh(JP.RxRates(n, majump)) +JP.fill_scaled_rates!(rx_rates.low.maj_rates, majump, nothing) +JP.fill_scaled_rates!(rx_rates.high.maj_rates, majump, nothing) # set up hop rates hop_constants = [1.0] diff --git a/test/thread_safety.jl b/test/thread_safety.jl index 9f7c44918..460dec735 100644 --- a/test/thread_safety.jl +++ b/test/thread_safety.jl @@ -39,8 +39,7 @@ let for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) jump_prob = JumpProblem(ode_prob, Direct(), vrj; vr_aggregator = agg) - prob_func(prob, i, repeat) = deepcopy(prob) - prob = EnsembleProblem(jump_prob, prob_func = prob_func) + prob = EnsembleProblem(jump_prob) sol = solve(prob, Tsit5(), EnsembleThreads(), trajectories = 400, save_everystep = false) firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol)] @@ -58,8 +57,7 @@ let for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) jump_prob = JumpProblem(sde_prob, Direct(), vrj; vr_aggregator = agg) - prob_func(prob, i, repeat) = deepcopy(prob) - prob = EnsembleProblem(jump_prob; prob_func) + prob = EnsembleProblem(jump_prob) sol = solve(prob, SRIW1(), EnsembleThreads(); trajectories = 400, save_everystep = false) firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol)] From 363587f3b006623d4c80ba69f0014e6b0b1e6b38 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Fri, 20 Mar 2026 17:24:22 -0400 Subject: [PATCH 26/52] Fix remaining CI: add SpatialMassActionJump fill_scaled_rates! no-op, update bracketing test - Add fill_scaled_rates! no-op for SpatialMassActionJump (rates stored in struct, no working copy needed). Placed in spatial_massaction_jump.jl since the type is defined after jumps.jl in the include order. - Update test/bracketing.jl: pass maj_rates to get_majump_brackets, add maj_rates field to DummyAggregator, call fill_scaled_rates! to init rates. Co-Authored-By: Claude Opus 4.5 --- src/spatial/spatial_massaction_jump.jl | 4 ++++ test/bracketing.jl | 11 +++++++---- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/src/spatial/spatial_massaction_jump.jl b/src/spatial/spatial_massaction_jump.jl index 52bcc7929..9706a1828 100644 --- a/src/spatial/spatial_massaction_jump.jl +++ b/src/spatial/spatial_massaction_jump.jl @@ -96,6 +96,10 @@ function SpatialMassActionJump(ma_jumps::MassActionJump{T, S, U, V}; scale_rates scale_rates = scale_rates, useiszero = useiszero, nocopy = nocopy) end +# SpatialMassActionJump stores rates in the struct itself (uniform_rates/spatial_rates), +# so fill_scaled_rates! is a no-op. +fill_scaled_rates!(dest, maj::SpatialMassActionJump, params) = nothing + ############################################## function get_num_majumps(smaj::SpatialMassActionJump{ diff --git a/test/bracketing.jl b/test/bracketing.jl index 182689432..6138e971a 100644 --- a/test/bracketing.jl +++ b/test/bracketing.jl @@ -36,9 +36,11 @@ reactstoch = [[1 => 1]] netstoch = [[1 => -1]] majump = MassActionJump(majump_rates, reactstoch, netstoch) +maj_rates_work = Vector{Float64}(undef, JP.get_num_majumps(majump)) +JP.fill_scaled_rates!(maj_rates_work, majump, nothing) reaction_index = 1 -@test JP.get_majump_brackets(ulow, uhigh, reaction_index, majump)[1] == majump_rates[1] * ulow[1] # low -@test JP.get_majump_brackets(ulow, uhigh, reaction_index, majump)[2] == majump_rates[1] * uhigh[1] # high +@test JP.get_majump_brackets(ulow, uhigh, reaction_index, majump, maj_rates_work)[1] == majump_rates[1] * ulow[1] # low +@test JP.get_majump_brackets(ulow, uhigh, reaction_index, majump, maj_rates_work)[2] == majump_rates[1] * uhigh[1] # high # constant rate rate(u, params, t) = 1 / u[1] @@ -56,6 +58,7 @@ mutable struct DummyAggregator{T, M, R, BD} <: cur_rate_high::Vector{T} sum_rate::T ma_jumps::M + maj_rates::Vector{T} rates::R bracket_data::BD end @@ -63,7 +66,7 @@ end cur_rate_low = [0.0, 0.0] cur_rate_high = [0.0, 0.0] sum_rate = 0.0 -p = DummyAggregator([0], [0], cur_rate_low, cur_rate_high, sum_rate, majump, [rate], bd) +p = DummyAggregator([0], [0], cur_rate_low, cur_rate_high, sum_rate, majump, maj_rates_work, [rate], bd) u = [100] JP.update_u_brackets!(p, u) @@ -77,7 +80,7 @@ reaction_index = 2 @test JP.get_jump_brackets(reaction_index, p, params, t)[1] == rate(p.uhigh, params, t) @test JP.get_jump_brackets(reaction_index, p, params, t)[2] == rate(p.ulow, params, t) -p = DummyAggregator([0], [0], cur_rate_low, cur_rate_high, sum_rate, majump, [rate], bd) +p = DummyAggregator([0], [0], cur_rate_low, cur_rate_high, sum_rate, majump, maj_rates_work, [rate], bd) JP.set_bracketing!(p, u, params, t) @test p.ulow[1]≈u[1] * (1 - fluctuation_rate) atol=1 @test p.uhigh[1]≈u[1] * (1 + fluctuation_rate) atol=1 From f31ae2f85c427b36ed6931f9aaaf8d557c67f0e6 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 21 Mar 2026 22:27:13 -0400 Subject: [PATCH 27/52] Fix scalar param_idxs bug and merge aliasing for parameterized MAJs Bug 1: A single parameterized MassActionJump with scalar param_idxs (e.g. param_idxs = 1) would segfault when passed directly to JumpProblem without merging, because the scalar MassActionJumpParamMapper{Int} had no callable method. Fix: convert scalar param_idxs to a one-element vector at construction time, and add a JumpSet constructor for single-reaction parameterized MAJs to normalize them to collection format. Bug 2: Merging two single-reaction parameterized MAJs with vector-backed param_idxs would mutate the first MAJ's mapper via aliasing in to_collection. Fix: to_collection now copies the param_idxs vector. Also: add docstring for the mapper callable API, remove dead code for Int-typed mappers, and add tests for both bugs. Co-Authored-By: Claude Opus 4.5 --- src/jumps.jl | 44 +++++++++++++++++++-------------------- test/ssa_callback_test.jl | 25 ++++++++++++++++++++++ 2 files changed, 47 insertions(+), 22 deletions(-) diff --git a/src/jumps.jl b/src/jumps.jl index 5d4f5f317..e1e60d057 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -421,7 +421,7 @@ function MassActionJump(rs, ns; param_idxs = nothing, param_mapper = nothing, if param_mapper === nothing (param_idxs === nothing) && error("If no parameter indices are given via param_idxs, an explicit parameter mapping must be passed in via param_mapper.") - pmapper = MassActionJumpParamMapper(param_idxs) + pmapper = MassActionJumpParamMapper(param_idxs isa Integer ? [param_idxs] : param_idxs) else (param_idxs !== nothing) && error("Only one of param_idxs and param_mapper should be passed.") @@ -442,12 +442,21 @@ struct MassActionJumpParamMapper{U} param_idxs::U end -# In-place mapper callable API: `(mapper)(dest, maj, params)` -# -# Fill `dest` with the scaled rates for the mass action jump given `params`. -# Custom mappers (e.g. MTK's JumpSysMajParamMapper) should implement this -# 3-arg callable. Use `maj.rescale_rates_on_update` to decide whether to -# apply stoichiometric scaling via `scalerates!`. +""" + (mapper)(dest::AbstractVector, maj::MassActionJump, params) + +In-place mapper callable API for parameterized `MassActionJump`s (those with +`scaled_rates = nothing`). Called by [`fill_scaled_rates!`](@ref) during +aggregator initialization and reinitialization to populate the working rate +vector from parameters. + +`dest` should be filled with the current rates for each reaction. If +`maj.rescale_rates_on_update` is `true`, the mapper should also apply +stoichiometric scaling via [`scalerates!`](@ref). + +Custom mappers (e.g. ModelingToolkitBase's `JumpSysMajParamMapper`) should +implement this 3-arg callable to support immutable `MassActionJump`s. +""" function (mapper::MassActionJumpParamMapper{U})(dest::AbstractVector, maj::MassActionJump, params) where {U <: AbstractArray} @inbounds for i in eachindex(dest) @@ -457,27 +466,13 @@ function (mapper::MassActionJumpParamMapper{U})(dest::AbstractVector, nothing end -function to_collection(ratemap::MassActionJumpParamMapper{Int}) - MassActionJumpParamMapper([ratemap.param_idxs]) -end -to_collection(ratemap::MassActionJumpParamMapper{<:AbstractVector}) = ratemap +to_collection(ratemap::MassActionJumpParamMapper) = MassActionJumpParamMapper(copy(ratemap.param_idxs)) function Base.merge!(pmap1::MassActionJumpParamMapper{U}, pmap2::MassActionJumpParamMapper{U}) where {U <: AbstractVector} append!(pmap1.param_idxs, pmap2.param_idxs) end -function Base.merge!(pmap1::MassActionJumpParamMapper{U}, - pmap2::MassActionJumpParamMapper{V}) where {U <: AbstractVector, - V <: Int} - push!(pmap1.param_idxs, pmap2.param_idxs) -end - -function Base.merge(pmap1::MassActionJumpParamMapper{Int}, - pmap2::MassActionJumpParamMapper{Int}) - MassActionJumpParamMapper([pmap1.param_idxs, pmap2.param_idxs]) -end - """ fill_scaled_rates!(dest::AbstractVector, maj::MassActionJump, params) @@ -547,6 +542,11 @@ end function JumpSet(vj, cj, rj, maj::MassActionJump{S, T, U, V}) where {S <: Number, T, U, V} JumpSet(vj, cj, rj, check_majump_type(maj)) end +function JumpSet(vj, cj, rj, + maj::MassActionJump{Nothing, T, U, V}) where { + T <: AbstractArray{<:Pair}, U <: AbstractArray{<:Pair}, V} + JumpSet(vj, cj, rj, check_majump_type(maj)) +end JumpSet(jump::ConstantRateJump) = JumpSet((), (jump,), nothing, nothing) JumpSet(jump::VariableRateJump) = JumpSet((jump,), (), nothing, nothing) diff --git a/test/ssa_callback_test.jl b/test/ssa_callback_test.jl index 10115a5ac..212ce9f42 100644 --- a/test/ssa_callback_test.jl +++ b/test/ssa_callback_test.jl @@ -231,3 +231,28 @@ let @test 3.0 ∈ sol5.t @test 6.0 ∈ sol5.t end + +# test that a single parameterized MAJ with scalar param_idxs works without merging +@testset "Single parameterized MAJ with scalar param_idxs" begin + maj = MassActionJump([1 => 1, 2 => 1], [1 => -1, 2 => -1, 3 => 1]; + param_idxs = 1) + p = [0.1] + u0 = [100, 100, 0] + dprob = DiscreteProblem(u0, (0.0, 10.0), p) + jprob = JumpProblem(dprob, Direct(), maj) + sol = solve(jprob, SSAStepper(); rng = StableRNG(12345)) + @test sol.retcode == ReturnCode.Success + @test sol[3, end] > 0 +end + +# test that merging MAJs does not mutate the input MAJs' param_mapper +@testset "Merging MAJs does not mutate inputs" begin + maj1 = MassActionJump([1 => 1], [1 => -1]; param_idxs = [1]) + maj2 = MassActionJump([2 => 1], [2 => -1]; param_idxs = [2]) + orig1 = copy(maj1.param_mapper.param_idxs) + orig2 = copy(maj2.param_mapper.param_idxs) + dprob = DiscreteProblem([100, 100], (0.0, 1.0), [0.1, 0.2]) + jprob = JumpProblem(dprob, Direct(), maj1, maj2) + @test maj1.param_mapper.param_idxs == orig1 + @test maj2.param_mapper.param_idxs == orig2 +end From f5fa5fceeed2f95760a518611630f6ddb531ed79 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 21 Mar 2026 22:32:26 -0400 Subject: [PATCH 28/52] Move mapper docstring to struct, document v10 breaking changes in HISTORY.md - Move mapper callable API docstring from the method to MassActionJumpParamMapper struct where it belongs - Add HISTORY.md entries for v10 breaking changes: scale_rates/useiszero removal from JumpProblem, immutable parameterized MAJ design, update_parameters! removal with reset_aggregated_jumps! guidance, custom mapper 3-arg callable requirement, scalar param_idxs normalization Co-Authored-By: Claude Opus 4.5 --- HISTORY.md | 25 +++++++++++++++++++++++++ src/jumps.jl | 24 +++++++++++++++--------- 2 files changed, 40 insertions(+), 9 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index d010778af..b12999b04 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -21,6 +21,31 @@ solver pathways (SSAStepper, ODE, SDE, tau-leaping). - `SSAIntegrator` now supports the `SciMLBase` RNG interface (`has_rng`, `get_rng`, `set_rng!`). + - **Breaking**: The `scale_rates` and `useiszero` keyword arguments have been + removed from `JumpProblem`. Set them on the `MassActionJump` directly: + ```julia + # Before (no longer works): + jprob = JumpProblem(dprob, Direct(), maj; scale_rates = false) + + # After: + maj = MassActionJump(rates, reactant_stoch, net_stoch; scale_rates = false) + jprob = JumpProblem(dprob, Direct(), maj) + ``` + - **Breaking**: Parameterized `MassActionJump`s (those constructed with + `param_idxs` or a custom `param_mapper`) are now immutable — rates are + computed from parameters at aggregator initialization rather than being + materialized into the jump at `JumpProblem` construction time. This means: + - `update_parameters!` has been removed. Mass action rates are now + automatically recomputed from the current parameter values whenever the + aggregator reinitializes. After modifying parameters (e.g. in a + callback), call `reset_aggregated_jumps!(integrator)` to trigger + reinitialization with the updated parameter values. + - Custom parameter mappers (e.g. ModelingToolkitBase's + `JumpSysMajParamMapper`) must implement the 3-arg callable API: + `(mapper)(dest::AbstractVector, maj::MassActionJump, params)`. + See [`MassActionJumpParamMapper`](@ref) for details. + - Scalar `param_idxs` (e.g. `param_idxs = 1`) is now internally converted to + a one-element vector. The scalar form continues to work as before. ## 9.14 diff --git a/src/jumps.jl b/src/jumps.jl index e1e60d057..f1e1bf5af 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -438,25 +438,31 @@ using_params(maj::Nothing) = false @inline get_num_majumps(maj::MassActionJump) = length(maj.net_stoch) @inline get_num_majumps(maj::Nothing) = 0 -struct MassActionJumpParamMapper{U} - param_idxs::U -end - """ + MassActionJumpParamMapper{U} + +Default parameter mapper for parameterized `MassActionJump`s (those constructed +with `param_idxs` instead of explicit `scaled_rates`). Stores the indices into +the parameter vector `p` that correspond to each reaction's rate constant. + +Implements the in-place mapper callable API: + (mapper)(dest::AbstractVector, maj::MassActionJump, params) -In-place mapper callable API for parameterized `MassActionJump`s (those with -`scaled_rates = nothing`). Called by [`fill_scaled_rates!`](@ref) during -aggregator initialization and reinitialization to populate the working rate -vector from parameters. +which is called by [`fill_scaled_rates!`](@ref) during aggregator initialization +and reinitialization to populate the working rate vector from parameters. `dest` should be filled with the current rates for each reaction. If `maj.rescale_rates_on_update` is `true`, the mapper should also apply stoichiometric scaling via [`scalerates!`](@ref). Custom mappers (e.g. ModelingToolkitBase's `JumpSysMajParamMapper`) should -implement this 3-arg callable to support immutable `MassActionJump`s. +implement this 3-arg callable to support parameterized `MassActionJump`s. """ +struct MassActionJumpParamMapper{U} + param_idxs::U +end + function (mapper::MassActionJumpParamMapper{U})(dest::AbstractVector, maj::MassActionJump, params) where {U <: AbstractArray} @inbounds for i in eachindex(dest) From ba5fe777a009b5d355952e0521980fd28dd92e5c Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Thu, 9 Apr 2026 08:29:11 -0400 Subject: [PATCH 29/52] Update SciMLBase compat to v3 and remove fork source SciMLBase 3 now includes the needed EnsembleProblem changes, so the fork dependency is no longer needed. Co-Authored-By: Claude Opus 4.5 --- Project.toml | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/Project.toml b/Project.toml index bd0840737..1f95b10f7 100644 --- a/Project.toml +++ b/Project.toml @@ -53,7 +53,7 @@ Random = "1" RecursiveArrayTools = "3.35, 4" Reexport = "1.2" SafeTestsets = "0.1" -SciMLBase = "2.147" +SciMLBase = "3" StableRNGs = "1" StaticArrays = "1.9.8" Statistics = "1" @@ -62,9 +62,6 @@ SymbolicIndexingInterface = "0.3.36" Test = "1" julia = "1.10" -[sources] -SciMLBase = {url = "https://github.com/isaacsas/SciMLBase.jl.git", rev = "ensemble_rng_redesign"} - [extras] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" From bb96d1c9df94781f533462a66e8cb9aa3200cc2f Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 22 Sep 2026 15:22:08 -0400 Subject: [PATCH 30/52] Align stochastic jump initialization with backend RNG and dispatch Route SDE/RODE init through StochasticDiffEqCore's JumpProblem initializer so solve and init share seed precedence, TaskLocalRNG conversion, seed metadata, and callback setup. Avoid merging stored callbacks twice. Cover the full OrdinaryDiffEqCore algorithm intersection to remove the generic initializer ambiguity while retaining narrower stochastic routes. Preserve explicit legacy alias_jump overrides through the existing backend alias_jumps field without changing other alias settings. Keep ODE/DAE and SSAStepper behavior, transitional helpers, and development version unchanged. Document the backend-specific RNG policy and actual dispatch pathways. Add registered regressions for RNG/seed trajectories and metadata, callback merging, native alias settings, and legacy overrides. Validated on Julia 1.12.7 with current released dependencies and an isolated minimal graph at SciMLBase 3.1, DiffEqBase 7.0, and OrdinaryDiffEqCore 4.0. Focused checks, DAE/backend tau smoke, both extension load orders, ambiguity checks, and fresh package precompilation pass. Existing compat floors remain unchanged. DelayDiffEq integrator RNG support and the independently reproduced backend alias_p=false issue remain recorded upstream follow-ups. No step-3 algorithm migration or broader alias API change is included. --- docs/src/api.md | 46 ++++++--- ext/JumpProcessesOrdinaryDiffEqCoreExt.jl | 32 ++++++- test/runtests.jl | 1 + test/stochastic_dispatch.jl | 112 ++++++++++++++++++++++ 4 files changed, 174 insertions(+), 17 deletions(-) create mode 100644 test/stochastic_dispatch.jl diff --git a/docs/src/api.md b/docs/src/api.md index a080c6335..6a886821f 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -153,28 +153,35 @@ sol = solve(jprob, SSAStepper(); seed = 1234) ### Resolution priority -When both `rng` and `seed` are passed to the same `solve`/`init` call, `rng` -takes priority: +For SSAStepper, OrdinaryDiffEq ODE/DAE solvers, SimpleTauLeaping, and +SimpleExplicitTauLeaping, an explicit `rng` takes priority over `seed`: | User provides | Result | |---|---| | `rng` via `solve`/`init` | Uses that `rng` | | `seed` via `solve`/`init` | Creates `Xoshiro(seed)` | -| Nothing | Uses `Random.default_rng()` (SSAStepper, ODE, tau-leaping) or a randomly-seeded `Xoshiro` (SDE) | +| Nothing | Uses `Random.default_rng()` | + +StochasticDiffEq handles SDE/RODE RNGs for both `solve` and `init`. An explicit +RNG other than `TaskLocalRNG` takes priority over seeds. With no RNG, or with +`TaskLocalRNG`, a nonzero solve/init `seed` takes priority, followed by the +underlying problem's nonzero `seed`; otherwise StochasticDiffEq creates a +randomly seeded `Xoshiro`. It converts `TaskLocalRNG` to `Xoshiro` rather than +storing it on the stochastic integrator. In this pathway, `seed=0` means no +seed override, whereas SSAStepper and the ODE/DAE pathway use `Xoshiro(0)`. ### Behavior by solver pathway | Solver | Default RNG (nothing passed) | `rng` / `seed` support | |---|---|---| | `SSAStepper` | `Random.default_rng()` | Full support via `solve`/`init` kwargs | -| ODE solvers (e.g., `Tsit5`) | `Random.default_rng()` | Full support via `solve`/`init` kwargs | -| SDE solvers (e.g., `SRIW1`) | Randomly-seeded `Xoshiro` | Full support; `TaskLocalRNG` is auto-converted to `Xoshiro` | +| OrdinaryDiffEq ODE/DAE solvers (e.g., `Tsit5`, `DFBDF`) | `Random.default_rng()` | Full support via `solve`/`init` kwargs | +| StochasticDiffEq SDE/RODE solvers (e.g., `SRIW1`, `RandomEM`) | `Xoshiro` from the stored problem seed, or a random seed | Full support; `TaskLocalRNG` follows the seed policy above and is converted to `Xoshiro` | | `SimpleTauLeaping` | `Random.default_rng()` | Full support via `solve` kwargs | !!! note For reproducible simulations, always pass an explicit `rng` or `seed`. - The default RNG is shared global state and may produce different results - depending on prior usage. + Julia's default RNG is task-local; results may depend on prior draws in the task. # Private / Developer API @@ -191,21 +198,32 @@ backends. | Solver type | `__solve` handled by | `__init` handled by | Uses `__jump_init`? | |---|---|---|---| | `SSAStepper` | JumpProcesses (`solve.jl`) | JumpProcesses (`SSA_stepper.jl`) | No | -| ODE (e.g., `Tsit5`) | JumpProcesses (`solve.jl`) | JumpProcesses (`solve.jl`) → OrdinaryDiffEq | Yes | -| SDE (e.g., `SRIW1`) | StochasticDiffEq | StochasticDiffEq | No | +| OrdinaryDiffEq ODE/DAE (e.g., `Tsit5`, `DFBDF`) | JumpProcesses (`solve.jl`) | JumpProcesses' OrdinaryDiffEqCore extension → OrdinaryDiffEq | Yes | +| StochasticDiffEq SDE/RODE (e.g., `SRIW1`, `RandomEM`) | StochasticDiffEqCore | JumpProcesses' OrdinaryDiffEqCore extension → StochasticDiffEqCore's JumpProblem initializer | No | +| StochasticDiffEq jump algorithms (e.g., `TauLeaping`) | StochasticDiffEqCore | StochasticDiffEqCore's specialized jump-algorithm initializer | No | | `SimpleTauLeaping` | JumpProcesses (`simple_regular_solve.jl`, custom `DiffEqBase.solve`) | N/A | No | For **SSAStepper**, `rng` is resolved via `resolve_rng` in `SSA_stepper.jl`'s `__init` and stored on the [`SSAIntegrator`](@ref). -For **ODE solvers**, `rng` is resolved via `resolve_rng` in `__jump_init` +For **OrdinaryDiffEq ODE/DAE solvers**, `rng` is resolved via `resolve_rng` in `__jump_init` (`solve.jl`) and forwarded to OrdinaryDiffEq's `init`, which stores it on the `ODEIntegrator`. -For **SDE solvers**, StochasticDiffEq handles the full solve/init pathway -directly (JumpProcesses' ambiguity-fix `__solve` method is never dispatched to). -StochasticDiffEq has its own `_resolve_rng` that additionally handles -`TaskLocalRNG` conversion and the problem's stored seed. +For **StochasticDiffEq SDE/RODE solvers**, the backend's more specific `__solve` +handles solves. JumpProcesses' extension resolves the public `init` dispatch +ambiguity by invoking the backend's existing initializer with the original +JumpProblem and unchanged RNG/seed inputs. Both routes therefore use +the backend's JumpProblem initializer, including its RNG policy, jump-state +copying, and callback setup. Older supported StochasticDiffEqCore releases route +`solve` through `__init` and therefore also use the extension. DiffEqBase merges +stored problem keywords and user callbacks before dispatching `init`; the +extension does not merge them a second time. + +An explicit legacy `alias_jump` passed to stochastic `init` retains its existing +jump-copy override through the backend's `alias_jumps` field, preserving other +alias settings. With that keyword omitted, the backend controls jump aliasing. +This compatibility path does not change the alias API of SSAStepper or ODE solvers. For **tau-leaping**, JumpProcesses defines a custom `DiffEqBase.solve` that bypasses the standard `__solve`/`__init` pathway. It calls `resolve_rng` diff --git a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl index 77f2631e4..3419bf529 100644 --- a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl +++ b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl @@ -13,16 +13,42 @@ end function SciMLBase.__init( _jump_prob::JumpProcesses.JumpProblem{IIP, P}, - alg::Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm}; + alg::Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm, + StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; kwargs...) where {IIP, P} + # Cover the full intersection with OrdinaryDiffEqCore's initializer. The + # narrower stochastic method below preserves its backend-owned setup. _jump_init(_jump_prob, alg; kwargs...) end function SciMLBase.__init( _jump_prob::JumpProcesses.JumpProblem{IIP, P}, alg::Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; - kwargs...) where {IIP, P} - _jump_init(_jump_prob, alg; kwargs...) + alias_jump = nothing, alias = nothing, kwargs...) where {IIP, P} + if alias_jump !== nothing + # Preserve the legacy override previously consumed by __jump_init. The + # backend already supports alias_jumps; retain its other alias settings. + aliases = if alias isa Union{SciMLBase.SDEAliasSpecifier, SciMLBase.RODEAliasSpecifier} + alias + else + constructor = alg isa StochasticDiffEqRODEAlgorithm ? + SciMLBase.RODEAliasSpecifier : SciMLBase.SDEAliasSpecifier + alias isa Bool ? constructor(; alias) : constructor() + end + fields = NamedTuple{fieldnames(typeof(aliases))}( + ntuple(i -> getfield(aliases, i), fieldcount(typeof(aliases)))) + alias = typeof(aliases)(; fields..., alias_jumps = alias_jump) + end + # Use the stochastic backend's JumpProblem initialization, as its solve path does. + # Keep rng and seed unresolved so it handles TaskLocalRNG conversion, problem + # seeds, and seed metadata consistently. Keeping the wrapper also preserves its + # jump aliasing and callback setup. The broader problem signature selects the + # backend method without recursively dispatching to this ambiguity fix. + # DiffEqBase.init_call already merged stored kwargs and user callbacks. + invoke(SciMLBase.__init, + Tuple{Union{SciMLBase.AbstractRODEProblem, JumpProcesses.JumpProblem}, + Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}}, + _jump_prob, alg; alias, kwargs...) end end diff --git a/test/runtests.jl b/test/runtests.jl index ef76d99f1..f2af5b418 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -43,6 +43,7 @@ end @time @safetestset "Saveat Regression test" begin include("saveat_regression.jl") end @time @safetestset "Save_positions test" begin include("save_positions.jl") end @time @safetestset "RNG kwarg tests" begin include("rng_kwarg_tests.jl") end + @time @safetestset "Stochastic dispatch tests" begin include("stochastic_dispatch.jl") end @time @safetestset "Ensemble Uniqueness test" begin include("ensemble_uniqueness.jl") end @time @safetestset "Ensemble Problem Tests" begin include("ensemble_problems.jl") end @time @safetestset "A + B <--> C" begin include("reversible_binding.jl") end diff --git a/test/stochastic_dispatch.jl b/test/stochastic_dispatch.jl new file mode 100644 index 000000000..fee265140 --- /dev/null +++ b/test/stochastic_dispatch.jl @@ -0,0 +1,112 @@ +using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq, SciMLBase +using Random, StableRNGs, Test + +function stochastic_jump_problem(kind; seed = 0, variable = false, callback = nothing) + f!(du, u, p, t) = (du .= -0.1 .* u) + g!(du, u, p, t) = (du .= 0.2) + f_rode!(du, u, p, t, W) = (du .= -0.1 .* u .+ 0.2 .* W) + prob = kind === :SDE ? + SDEProblem(f!, g!, [10.0], (0.0, 0.5), zeros(Int, 2); seed) : + RODEProblem(f_rode!, [10.0], (0.0, 0.5), zeros(Int, 2); seed) + rate(u, p, t) = 10.0 + affect!(integrator) = (integrator.u[1] += 1) + jump = variable ? VariableRateJump(rate, affect!) : ConstantRateJump(rate, affect!) + JumpProblem(prob, Direct(), jump; callback) +end + +stochastic_rng_options(kind) = + kind === :default ? (;) : + kind === :seed ? (; seed = 456) : + kind === :zero_seed ? (; seed = 0) : + kind === :task_rng_seed ? (; rng = Random.default_rng(), seed = 456) : + kind === :task_rng ? (; rng = Random.default_rng()) : + (; rng = StableRNG(789), seed = 456) + +@testset "Stochastic solve/init use the same RNG policy" begin + @testset "$kind stored seed=$stored_seed options=$options_kind variable=$variable" for + (kind, alg, variable) in ((:SDE, EM(), false), (:SDE, EM(), true), + (:RODE, RandomEM(), false)), + stored_seed in (0, 123), + options_kind in (:default, :seed, :zero_seed, :task_rng_seed, :task_rng, + :explicit_rng_seed) + Random.seed!(999) + sol = solve(stochastic_jump_problem(kind; seed = stored_seed, variable), alg; + dt = 0.01, stochastic_rng_options(options_kind)...) + + Random.seed!(999) + options = stochastic_rng_options(options_kind) + integrator = init(stochastic_jump_problem(kind; seed = stored_seed, variable), alg; + dt = 0.01, options...) + if options_kind === :explicit_rng_seed + @test SciMLBase.get_rng(integrator) === options.rng + else + @test SciMLBase.get_rng(integrator) isa Xoshiro + end + solve!(integrator) + @test successful_retcode(sol) + @test successful_retcode(integrator.sol) + @test sol.t == integrator.sol.t + @test sol.u == integrator.sol.u + @test sol.seed == integrator.sol.seed + expected_seed = options_kind === :explicit_rng_seed ? 0 : + options_kind in (:seed, :task_rng_seed) ? 456 : stored_seed + if expected_seed != 0 || options_kind === :explicit_rng_seed + @test sol.seed == expected_seed + end + end +end + +@testset "Stochastic initialization preserves callback merging" begin + @testset "merge_callbacks=$merge_callbacks entry=$entry" for + merge_callbacks in (true, false), entry in (:solve, :init) + cb1 = DiscreteCallback((u, t, integrator) -> t == 0.25, + integrator -> (integrator.p[1] += 1); save_positions = (false, false)) + cb2 = DiscreteCallback((u, t, integrator) -> t == 0.25, + integrator -> (integrator.p[2] += 1); save_positions = (false, false)) + jprob = stochastic_jump_problem(:SDE; callback = cb1) + options = (; seed = 123, dt = 0.01, tstops = [0.25], callback = cb2, merge_callbacks) + sol = if entry === :solve + solve(jprob, EM(); options...) + else + integrator = init(jprob, EM(); options...) + solve!(integrator) + integrator.sol + end + @test successful_retcode(sol) + @test sol.prob.p == [Int(merge_callbacks), 1] + end +end + +@testset "Stochastic initialization respects the backend jump-copy option" begin + for alias_jumps in (true, false) + jprob = stochastic_jump_problem(:SDE) + integrator = init(jprob, EM(); dt = 0.01, seed = 123, + alias = SciMLBase.SDEAliasSpecifier(; alias_jumps)) + aliases_aggregation = any(integrator.opts.callback.discrete_callbacks) do callback + callback.condition === jprob.discrete_jump_aggregation + end + @test aliases_aggregation == alias_jumps + solve!(integrator) + @test successful_retcode(integrator.sol) + end +end + +@testset "Stochastic initialization preserves explicit legacy alias_jump" begin + for (kind, alg) in ((:SDE, EM()), (:RODE, RandomEM())), + alias_jump in (true, false), native_alias in (nothing, true, false, :specifier) + jprob = stochastic_jump_problem(kind) + constructor = kind === :SDE ? SciMLBase.SDEAliasSpecifier : SciMLBase.RODEAliasSpecifier + alias = native_alias === :specifier ? + constructor(; alias_jumps = !alias_jump, alias_u0 = false) : native_alias + integrator = init(jprob, alg; dt = 0.01, seed = 123, alias_jump, alias) + aliases_aggregation = any(integrator.opts.callback.discrete_callbacks) do callback + callback.condition === jprob.discrete_jump_aggregation + end + @test aliases_aggregation == alias_jump + if native_alias in (false, :specifier) + @test integrator.u !== jprob.prob.u0 + end + solve!(integrator) + @test successful_retcode(integrator.sol) + end +end From 5468eac1e6cafaacaa14713370e259ef285d397b Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Thu, 24 Sep 2026 09:52:00 -0400 Subject: [PATCH 31/52] Use upstream stochastic JumpProblem initialization Remove the invoke workaround and temporary legacy alias adapter. Retain the broad OrdinaryDiffEqCore intersection method, and update stochastic dispatch tests and API documentation for backend-owned initialization. Validated 377 assertions using StochasticDiffEqCore from OrdinaryDiffEq PR #4614 at 77c07bffe2. The temporary source override remains outside the package; the released Core runtime compatibility floor is still pending. --- docs/src/api.md | 27 +++++++++--------- ext/JumpProcessesOrdinaryDiffEqCoreExt.jl | 34 ++--------------------- test/stochastic_dispatch.jl | 25 ++++++++++++----- 3 files changed, 33 insertions(+), 53 deletions(-) diff --git a/docs/src/api.md b/docs/src/api.md index 6a886821f..fe8e53a4d 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -199,7 +199,7 @@ backends. |---|---|---|---| | `SSAStepper` | JumpProcesses (`solve.jl`) | JumpProcesses (`SSA_stepper.jl`) | No | | OrdinaryDiffEq ODE/DAE (e.g., `Tsit5`, `DFBDF`) | JumpProcesses (`solve.jl`) | JumpProcesses' OrdinaryDiffEqCore extension → OrdinaryDiffEq | Yes | -| StochasticDiffEq SDE/RODE (e.g., `SRIW1`, `RandomEM`) | StochasticDiffEqCore | JumpProcesses' OrdinaryDiffEqCore extension → StochasticDiffEqCore's JumpProblem initializer | No | +| StochasticDiffEq SDE/RODE (e.g., `SRIW1`, `RandomEM`) | StochasticDiffEqCore | StochasticDiffEqCore's JumpProblem initializer | No | | StochasticDiffEq jump algorithms (e.g., `TauLeaping`) | StochasticDiffEqCore | StochasticDiffEqCore's specialized jump-algorithm initializer | No | | `SimpleTauLeaping` | JumpProcesses (`simple_regular_solve.jl`, custom `DiffEqBase.solve`) | N/A | No | @@ -211,19 +211,18 @@ For **OrdinaryDiffEq ODE/DAE solvers**, `rng` is resolved via `resolve_rng` in ` `ODEIntegrator`. For **StochasticDiffEq SDE/RODE solvers**, the backend's more specific `__solve` -handles solves. JumpProcesses' extension resolves the public `init` dispatch -ambiguity by invoking the backend's existing initializer with the original -JumpProblem and unchanged RNG/seed inputs. Both routes therefore use -the backend's JumpProblem initializer, including its RNG policy, jump-state -copying, and callback setup. Older supported StochasticDiffEqCore releases route -`solve` through `__init` and therefore also use the extension. DiffEqBase merges -stored problem keywords and user callbacks before dispatching `init`; the -extension does not merge them a second time. - -An explicit legacy `alias_jump` passed to stochastic `init` retains its existing -jump-copy override through the backend's `alias_jumps` field, preserving other -alias settings. With that keyword omitted, the backend controls jump aliasing. -This compatibility path does not change the alias API of SSAStepper or ODE solvers. +and `__init` methods handle the original JumpProblem and unchanged RNG/seed +inputs. Both routes use the backend's JumpProblem initializer, including its RNG +policy, jump-state copying, and callback setup. DiffEqBase merges stored problem +keywords and user callbacks before dispatching `init`. + +Control stochastic jump-state copying through the backend's alias specifier: +`alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)` for SDE problems, +or `SciMLBase.RODEAliasSpecifier` for RODE problems. Set the `alias_jumps` field +to `true` to reuse the original jump state. When that field is unspecified, the +backend aliases on thread 1 and copies on other threads. Both `solve` and `init` +follow this policy; `alias_jumps` is a field of the specifier, not a standalone +keyword argument. For **tau-leaping**, JumpProcesses defines a custom `DiffEqBase.solve` that bypasses the standard `__solve`/`__init` pathway. It calls `resolve_rng` diff --git a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl index 3419bf529..cf9697b45 100644 --- a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl +++ b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl @@ -16,39 +16,9 @@ function SciMLBase.__init( alg::Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm, StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; kwargs...) where {IIP, P} - # Cover the full intersection with OrdinaryDiffEqCore's initializer. The - # narrower stochastic method below preserves its backend-owned setup. + # Cover the full intersection with OrdinaryDiffEqCore's initializer. + # StochasticDiffEqCore's narrower methods own SDE/RODE initialization. _jump_init(_jump_prob, alg; kwargs...) end -function SciMLBase.__init( - _jump_prob::JumpProcesses.JumpProblem{IIP, P}, - alg::Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; - alias_jump = nothing, alias = nothing, kwargs...) where {IIP, P} - if alias_jump !== nothing - # Preserve the legacy override previously consumed by __jump_init. The - # backend already supports alias_jumps; retain its other alias settings. - aliases = if alias isa Union{SciMLBase.SDEAliasSpecifier, SciMLBase.RODEAliasSpecifier} - alias - else - constructor = alg isa StochasticDiffEqRODEAlgorithm ? - SciMLBase.RODEAliasSpecifier : SciMLBase.SDEAliasSpecifier - alias isa Bool ? constructor(; alias) : constructor() - end - fields = NamedTuple{fieldnames(typeof(aliases))}( - ntuple(i -> getfield(aliases, i), fieldcount(typeof(aliases)))) - alias = typeof(aliases)(; fields..., alias_jumps = alias_jump) - end - # Use the stochastic backend's JumpProblem initialization, as its solve path does. - # Keep rng and seed unresolved so it handles TaskLocalRNG conversion, problem - # seeds, and seed metadata consistently. Keeping the wrapper also preserves its - # jump aliasing and callback setup. The broader problem signature selects the - # backend method without recursively dispatching to this ambiguity fix. - # DiffEqBase.init_call already merged stored kwargs and user callbacks. - invoke(SciMLBase.__init, - Tuple{Union{SciMLBase.AbstractRODEProblem, JumpProcesses.JumpProblem}, - Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}}, - _jump_prob, alg; alias, kwargs...) -end - end diff --git a/test/stochastic_dispatch.jl b/test/stochastic_dispatch.jl index fee265140..89419036c 100644 --- a/test/stochastic_dispatch.jl +++ b/test/stochastic_dispatch.jl @@ -22,6 +22,14 @@ stochastic_rng_options(kind) = kind === :task_rng ? (; rng = Random.default_rng()) : (; rng = StableRNG(789), seed = 456) +@testset "Stochastic initialization dispatches to the backend" begin + for (kind, alg) in ((:SDE, EM()), (:RODE, RandomEM())) + jprob = stochastic_jump_problem(kind) + method = which(SciMLBase.__init, (typeof(jprob), typeof(alg))) + @test method.module === StochasticDiffEq.StochasticDiffEqCore + end +end + @testset "Stochastic solve/init use the same RNG policy" begin @testset "$kind stored seed=$stored_seed options=$options_kind variable=$variable" for (kind, alg, variable) in ((:SDE, EM(), false), (:SDE, EM(), true), @@ -91,19 +99,22 @@ end end end -@testset "Stochastic initialization preserves explicit legacy alias_jump" begin +@testset "Stochastic initialization preserves native alias settings" begin for (kind, alg) in ((:SDE, EM()), (:RODE, RandomEM())), - alias_jump in (true, false), native_alias in (nothing, true, false, :specifier) + alias_mode in (:default, :alias, :copy, :specifier_alias, :specifier_copy) jprob = stochastic_jump_problem(kind) constructor = kind === :SDE ? SciMLBase.SDEAliasSpecifier : SciMLBase.RODEAliasSpecifier - alias = native_alias === :specifier ? - constructor(; alias_jumps = !alias_jump, alias_u0 = false) : native_alias - integrator = init(jprob, alg; dt = 0.01, seed = 123, alias_jump, alias) + alias_jumps = alias_mode === :default ? Threads.threadid() == 1 : + alias_mode in (:alias, :specifier_alias) + alias = alias_mode === :default ? nothing : + alias_mode in (:alias, :copy) ? alias_jumps : + constructor(; alias_jumps, alias_u0 = false) + integrator = init(jprob, alg; dt = 0.01, seed = 123, alias) aliases_aggregation = any(integrator.opts.callback.discrete_callbacks) do callback callback.condition === jprob.discrete_jump_aggregation end - @test aliases_aggregation == alias_jump - if native_alias in (false, :specifier) + @test aliases_aggregation == alias_jumps + if alias_mode in (:copy, :specifier_alias, :specifier_copy) @test integrator.u !== jprob.prob.u0 end solve!(integrator) From b811f6405412f75d91ef933729796d0ab9d38873 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 29 Sep 2026 08:17:17 -0400 Subject: [PATCH 32/52] Reconcile tau-leaping and kernel RNG and mass-action state Use solve-level RNG resolution and fresh mass-action rates in CPU tau algorithms. Materialize current host rates for kernels, reject unsupported RNG controls, and preserve RNG forwarding in serial fallbacks. Fix scalar mass-action dispatch and removed-keyword diagnostics. Preserve SortingDirect learned ordering and migrate stale test, example, and benchmark RNG calls with coverage for parameter refresh and state reuse. Validation: seven focused lanes passed 2,692 assertions on Julia 1.12.7 with the pinned PR #4614 dependency graph, including four-thread CPU kernel checks. Actual CUDA validation remains outstanding. --- benchmark/benchmarks.jl | 24 +-- docs/src/api.md | 7 +- docs/src/tutorials/simple_poisson_process.md | 19 +-- ext/JumpProcessesKernelAbstractionsExt.jl | 36 +++- ext/explicit_tau_leaping.jl | 25 +-- ext/ssa_stepper.jl | 27 +-- src/SSA_stepper.jl | 7 +- src/aggregators/rssa.jl | 5 - src/jumps.jl | 5 + src/problem.jl | 6 + src/simple_regular_solve.jl | 42 +++-- test/bracketing.jl | 4 +- test/explicit_tau_kernel_tests.jl | 4 + test/fp_unknowns.jl | 3 +- test/kernel_contract_tests.jl | 127 ++++++++++++++ test/massaction_state.jl | 147 ++++++++++++++++ test/regular_jumps.jl | 166 +++++++++++++++++-- test/runtests.jl | 2 + test/sortingdirect_reuse.jl | 58 +++++++ test/ssa_aggregator_smoke.jl | 4 +- test/ssa_kernel_tests.jl | 5 + test/variable_rate.jl | 4 +- 22 files changed, 627 insertions(+), 100 deletions(-) create mode 100644 test/kernel_contract_tests.jl create mode 100644 test/massaction_state.jl create mode 100644 test/sortingdirect_reuse.jl diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 6a59991e5..8d7167676 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -42,16 +42,16 @@ SUITE["construct"]["jumpproblem_constantrate"] = @benchmarkable JumpProblem( ) # ============================================================================= -# Solves (SSAStepper, fixed rng for reproducibility) +# Solves (SSAStepper, fresh seeded RNG for each sample) # ============================================================================= SUITE["solve"] = BenchmarkGroup() -jprob_maj = JumpProblem(dprob, Direct(), maj; rng = StableRNG(12345)) -jprob_cr = JumpProblem(dprob, Direct(), jump1, jump2; rng = StableRNG(12345)) +jprob_maj = JumpProblem(dprob, Direct(), maj) +jprob_cr = JumpProblem(dprob, Direct(), jump1, jump2) -SUITE["solve"]["massaction"] = @benchmarkable solve($jprob_maj, SSAStepper()) -SUITE["solve"]["constantrate"] = @benchmarkable solve($jprob_cr, SSAStepper()) +SUITE["solve"]["massaction"] = @benchmarkable solve($jprob_maj, SSAStepper(); rng) setup=(rng = StableRNG(12345)) evals=1 +SUITE["solve"]["constantrate"] = @benchmarkable solve($jprob_cr, SSAStepper(); rng) setup=(rng = StableRNG(12345)) evals=1 # ============================================================================= # Aggregators @@ -59,12 +59,12 @@ SUITE["solve"]["constantrate"] = @benchmarkable solve($jprob_cr, SSAStepper()) SUITE["aggregators"] = BenchmarkGroup() -jprob_rdirect = JumpProblem(dprob, RDirect(), maj; rng = StableRNG(12345)) -jprob_sorting = JumpProblem(dprob, SortingDirect(), maj; rng = StableRNG(12345)) -jprob_nrm = JumpProblem(dprob, NRM(), maj; rng = StableRNG(12345)) +jprob_rdirect = JumpProblem(dprob, RDirect(), maj) +jprob_sorting = JumpProblem(dprob, SortingDirect(), maj) +jprob_nrm = JumpProblem(dprob, NRM(), maj) -SUITE["aggregators"]["RDirect"] = @benchmarkable solve($jprob_rdirect, SSAStepper()) +SUITE["aggregators"]["RDirect"] = @benchmarkable solve($jprob_rdirect, SSAStepper(); rng) setup=(rng = StableRNG(12345)) evals=1 SUITE["aggregators"]["SortingDirect"] = @benchmarkable solve( - $jprob_sorting, SSAStepper() -) -SUITE["aggregators"]["NRM"] = @benchmarkable solve($jprob_nrm, SSAStepper()) + $jprob_sorting, SSAStepper(); rng +) setup=(rng = StableRNG(12345)) evals=1 +SUITE["aggregators"]["NRM"] = @benchmarkable solve($jprob_nrm, SSAStepper(); rng) setup=(rng = StableRNG(12345)) evals=1 diff --git a/docs/src/api.md b/docs/src/api.md index fe8e53a4d..5ca6d1348 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -153,8 +153,8 @@ sol = solve(jprob, SSAStepper(); seed = 1234) ### Resolution priority -For SSAStepper, OrdinaryDiffEq ODE/DAE solvers, SimpleTauLeaping, and -SimpleExplicitTauLeaping, an explicit `rng` takes priority over `seed`: +For SSAStepper, OrdinaryDiffEq ODE/DAE solvers, and all five CPU tau-leaping +algorithms, an explicit `rng` takes priority over `seed`: | User provides | Result | |---|---| @@ -178,6 +178,7 @@ seed override, whereas SSAStepper and the ODE/DAE pathway use `Xoshiro(0)`. | OrdinaryDiffEq ODE/DAE solvers (e.g., `Tsit5`, `DFBDF`) | `Random.default_rng()` | Full support via `solve`/`init` kwargs | | StochasticDiffEq SDE/RODE solvers (e.g., `SRIW1`, `RandomEM`) | `Xoshiro` from the stored problem seed, or a random seed | Full support; `TaskLocalRNG` follows the seed policy above and is converted to `Xoshiro` | | `SimpleTauLeaping` | `Random.default_rng()` | Full support via `solve` kwargs | +| `SimpleExplicitTauLeaping`, `SimpleImplicitTauLeaping`, `SimpleTrapezoidalLeaping`, `SimpleAdaptiveTauLeaping` | `Random.default_rng()` | Full support via `solve` kwargs | !!! note For reproducible simulations, always pass an explicit `rng` or `seed`. @@ -201,7 +202,7 @@ backends. | OrdinaryDiffEq ODE/DAE (e.g., `Tsit5`, `DFBDF`) | JumpProcesses (`solve.jl`) | JumpProcesses' OrdinaryDiffEqCore extension → OrdinaryDiffEq | Yes | | StochasticDiffEq SDE/RODE (e.g., `SRIW1`, `RandomEM`) | StochasticDiffEqCore | StochasticDiffEqCore's JumpProblem initializer | No | | StochasticDiffEq jump algorithms (e.g., `TauLeaping`) | StochasticDiffEqCore | StochasticDiffEqCore's specialized jump-algorithm initializer | No | -| `SimpleTauLeaping` | JumpProcesses (`simple_regular_solve.jl`, custom `DiffEqBase.solve`) | N/A | No | +| All five CPU tau-leaping algorithms | JumpProcesses (`simple_regular_solve.jl`, custom `DiffEqBase.solve`) | N/A | No | For **SSAStepper**, `rng` is resolved via `resolve_rng` in `SSA_stepper.jl`'s `__init` and stored on the [`SSAIntegrator`](@ref). diff --git a/docs/src/tutorials/simple_poisson_process.md b/docs/src/tutorials/simple_poisson_process.md index 11f5c5f0e..20cd2e97e 100644 --- a/docs/src/tutorials/simple_poisson_process.md +++ b/docs/src/tutorials/simple_poisson_process.md @@ -27,6 +27,7 @@ installed via using Pkg Pkg.add("JumpProcesses") Pkg.add("Plots") +Pkg.add("SciMLBase") ``` Let's also load our packages and set some defaults for our plot formatting @@ -371,28 +372,24 @@ with ``N(t)`` a Poisson counting process with constant transition rate ``\lambda``, and the ``C_i`` independent and identical samples from a uniform distribution over ``\{-1,1\}``. We can simulate such a process as follows. -We first ensure that we use the same random number generator as JumpProcesses. -Custom RNGs can be passed to `solve` or `init` via the `rng` keyword argument. -If no RNG is specified, JumpProcesses uses `Random.default_rng()`, which is also -available as `JumpProcesses.DEFAULT_RNG`. Let's use the default - -```@example tut1 -rng = JumpProcesses.DEFAULT_RNG -``` +The affect function can use `SciMLBase.get_rng(integrator)` to draw from the +same random number generator as the jump solver. Custom RNGs can be passed to +`solve` or `init` via the `rng` keyword argument. Let's assume `u[1]` is ``N(t)`` and `u[2]` is ``G(t)``. We now proceed as in the previous examples ```@example tut1 +using SciMLBase + rate3(u, p, t) = p.λ -# define the affect function via a closure -affect3! = integrator -> let rng = rng +function affect3!(integrator) # N(t) <-- N(t) + 1 integrator.u[1] += 1 # G(t) <-- G(t) + C_{N(t)} - integrator.u[2] += rand(rng, (-1, 1)) + integrator.u[2] += rand(SciMLBase.get_rng(integrator), (-1, 1)) nothing end crj = ConstantRateJump(rate3, affect3!) diff --git a/ext/JumpProcessesKernelAbstractionsExt.jl b/ext/JumpProcessesKernelAbstractionsExt.jl index 5df7b3876..c5275252b 100644 --- a/ext/JumpProcessesKernelAbstractionsExt.jl +++ b/ext/JumpProcessesKernelAbstractionsExt.jl @@ -5,6 +5,26 @@ using KernelAbstractions, Adapt using StaticArrays using PoissonRandom, Random +# Kernels draw from the backend's own RNG; a host RNG object cannot be used +# inside device code. CPU kernels use Julia's task-local generator, which can be +# seeded here. Device generators must be seeded with their backend's API before +# calling solve (for example, CUDA.seed! for CUDA). +function seed_kernel_backend!(backend, seed, rng, kwargs) + rng === nothing || throw(ArgumentError( + "EnsembleGPUKernel does not support a host `rng` with multiple trajectories. " * + "Use `seed` with the CPU backend or seed the device RNG through its backend API.")) + haskey(kwargs, :rng_func) && throw(ArgumentError( + "EnsembleGPUKernel does not support `rng_func` with multiple trajectories. " * + "Kernels draw from the backend's own RNG.")) + if seed !== nothing + backend isa CPU || throw(ArgumentError( + "EnsembleGPUKernel only supports the `seed` keyword with the CPU backend. " * + "Seed the device RNG through its backend API before calling solve.")) + Random.seed!(seed) + end + return nothing +end + include("ssa_stepper.jl") include("explicit_tau_leaping.jl") @@ -13,15 +33,17 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, ensemblealg::EnsembleGPUKernel; trajectories, seed = nothing, + rng = nothing, dt = error("dt is required for SimpleTauLeaping."), kwargs...) if trajectories == 1 return SciMLBase.__solve(ensembleprob, alg, EnsembleSerial(); trajectories = 1, - seed, dt, kwargs...) + seed, rng, dt, kwargs...) end ensemblealg.backend === nothing ? backend = CPU() : backend = ensemblealg.backend + seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob @@ -35,7 +57,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, # Run vectorized solve ts, us = vectorized_solve( - probs, jump_prob, SimpleTauLeaping(); backend, trajectories, seed, dt) + probs, jump_prob, SimpleTauLeaping(); backend, trajectories, dt) # Convert to CPU for inspection _ts = Array(ts) @@ -84,8 +106,7 @@ end # SimpleTauLeaping kernel @kernel function simple_tau_leaping_kernel( @Const(probs_data), _us, _ts, dt, @Const(rj_data), - current_u_buf, rate_cache_buf, counts_buf, local_dc_buf, - seed::UInt64) + current_u_buf, rate_cache_buf, counts_buf, local_dc_buf) i = @index(Global, Linear) # Get thread-local buffers @@ -152,7 +173,7 @@ end # Vectorized solve function function vectorized_solve(probs, prob::JumpProblem, alg::SimpleTauLeaping; - backend, trajectories, seed, dt, kwargs...) + backend, trajectories, dt, kwargs...) # Extract common jump data rj = prob.regular_jump rj_data = JumpData(rj.rate, rj.c, rj.numjumps) @@ -199,13 +220,10 @@ function vectorized_solve(probs, prob::JumpProblem, alg::SimpleTauLeaping; init_event = init_kernel(probs_data_gpu, current_u_buf; ndrange = n_trajectories) KernelAbstractions.synchronize(backend) - # Seed for Poisson sampling - seed = seed === nothing ? UInt64(12345) : UInt64(seed); - # Launch main kernel kernel = simple_tau_leaping_kernel(backend) main_event = kernel(probs_data_gpu, us, ts, dt, rj_data_gpu, - current_u_buf, rate_cache_buf, counts_buf, local_dc_buf, seed; + current_u_buf, rate_cache_buf, counts_buf, local_dc_buf; ndrange = n_trajectories) KernelAbstractions.synchronize(backend) diff --git a/ext/explicit_tau_leaping.jl b/ext/explicit_tau_leaping.jl index 38bc6a524..7af58e362 100644 --- a/ext/explicit_tau_leaping.jl +++ b/ext/explicit_tau_leaping.jl @@ -213,10 +213,12 @@ trajectory takes is not known in advance and the solution is sampled onto a fixe time grid instead. Each returned solution therefore contains exactly the grid points, as if the problem had been built with `save_positions = (false, false)`. -Randomness comes from the backend's own device RNG rather than the `rng` stored -in the `JumpProblem`. `seed` is applied to the ambient generator, so it makes a -run reproducible on backends that draw from it, such as `CPU()`; seeding a GPU -backend is done through that backend's own `seed!`. +For multiple trajectories, randomness comes from the backend's own RNG. `seed` +is supported with `CPU()` and seeds Julia's task-local generator; device backends +must be seeded through their own API before calling `solve` (for example, +`CUDA.seed!`). Passing a host `rng` or `rng_func`, or `seed` with a device backend, +is rejected. +A single trajectory uses `EnsembleSerial()` and its usual solve-level RNG inputs. The reaction data is uploaded to the device once and shared by every thread, so all trajectories solve the same problem and a `prob_func` is not supported. @@ -226,24 +228,23 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, ensemblealg::EnsembleGPUKernel; trajectories, seed = nothing, + rng = nothing, dtmin = nothing, saveat = nothing, save_start = true, save_end = true, callback = nothing, kwargs...) + callback === nothing || + error("EnsembleGPUKernel with SimpleExplicitTauLeaping does not support callbacks.") + if trajectories == 1 return SciMLBase.__solve(ensembleprob, alg, EnsembleSerial(); trajectories = 1, - seed, dtmin, saveat, save_start, save_end, callback, kwargs...) + seed, rng, dtmin, saveat, save_start, save_end, kwargs...) end - callback === nothing || - error("EnsembleGPUKernel with SimpleExplicitTauLeaping does not support callbacks, \ - since they would have to run inside the GPU kernel.") - - seed !== nothing && Random.seed!(seed) - ensemblealg.backend === nothing ? backend = CPU() : backend = ensemblealg.backend + seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob jump_prob isa JumpProblem || @@ -289,7 +290,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, # lives in registers rather than being uploaded once per thread. u0 = SVector{state_dim, ET}(prob.u0) saveat_gpu = adapt(backend, save_times) - maj_gpu = GPUMassActionJump(maj, backend, TT) + maj_gpu = GPUMassActionJump(maj, prob.p, backend, TT) # g_i depends only on the reactant stoichiometry, so the per-species highest # order and largest stoichiometry are precomputed once on the host. diff --git a/ext/ssa_stepper.jl b/ext/ssa_stepper.jl index 7e6d14823..f666da319 100644 --- a/ext/ssa_stepper.jl +++ b/ext/ssa_stepper.jl @@ -54,12 +54,17 @@ function flatten_stoich(stoch, ::Type{IT}) where {IT <: Integer} species, coeffs, offsets end -function GPUMassActionJump(maj::JumpProcesses.MassActionJump, backend, +function GPUMassActionJump(maj::JumpProcesses.MassActionJump, params, backend, ::Type{TT}) where {TT <: AbstractFloat} rs_species, rs_coeffs, rs_offsets = flatten_stoich(maj.reactant_stoch, Int32) ns_species, ns_coeffs, ns_offsets = flatten_stoich(maj.net_stoch, Int32) - GPUMassActionJump(adapt(backend, convert(Vector{TT}, maj.scaled_rates)), + # The host definition can be parameter-mapped, with scaled_rates === nothing. + # Materialize fresh numeric working rates before adapting them to the device. + rates = Vector{TT}(undef, JumpProcesses.get_num_majumps(maj)) + JumpProcesses.fill_scaled_rates!(rates, maj, params) + + GPUMassActionJump(adapt(backend, rates), adapt(backend, rs_species), adapt(backend, rs_coeffs), adapt(backend, rs_offsets), @@ -263,10 +268,12 @@ grid points, as if the problem had been built with `save_positions = (false, false)`; the aggregator chosen in the `JumpProblem` is also ignored, since the kernel always runs the Direct method. -Randomness comes from the backend's own device RNG rather than the `rng` stored -in the `JumpProblem`. `seed` is applied to the ambient generator, so it makes a -run reproducible on backends that draw from it, such as `CPU()`; seeding a GPU -backend is done through that backend's own `seed!`. +For multiple trajectories, randomness comes from the backend's own RNG. `seed` +is supported with `CPU()` and seeds Julia's task-local generator; device backends +must be seeded through their own API before calling `solve` (for example, +`CUDA.seed!`). Passing a host `rng` or `rng_func`, or `seed` with a device backend, +is rejected. +A single trajectory uses `EnsembleSerial()` and its usual solve-level RNG inputs. The reaction data is uploaded to the device once and shared by every thread, so all trajectories solve the same problem and a `prob_func` is not supported. @@ -276,6 +283,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, ensemblealg::EnsembleGPUKernel; trajectories, seed = nothing, + rng = nothing, saveat = nothing, save_start = true, save_end = true, @@ -283,16 +291,15 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, kwargs...) if trajectories == 1 return SciMLBase.__solve(ensembleprob, alg, EnsembleSerial(); trajectories = 1, - seed, saveat, save_start, save_end, callback, kwargs...) + seed, rng, saveat, save_start, save_end, callback, kwargs...) end callback === nothing || error("EnsembleGPUKernel with SSAStepper does not support callbacks, since they \ would have to run inside the GPU kernel.") - seed !== nothing && Random.seed!(seed) - ensemblealg.backend === nothing ? backend = CPU() : backend = ensemblealg.backend + seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob jump_prob isa JumpProblem || @@ -340,7 +347,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, # lives in registers rather than being uploaded once per thread. u0 = SVector{state_dim, ET}(prob.u0) saveat_gpu = adapt(backend, save_times) - maj_gpu = GPUMassActionJump(maj, backend, TT) + maj_gpu = GPUMassActionJump(maj, prob.p, backend, TT) us = allocate(backend, ET, (trajectories, state_dim, nsave)) diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index fd1b73b3a..919ace32d 100644 --- a/src/SSA_stepper.jl +++ b/src/SSA_stepper.jl @@ -462,7 +462,12 @@ function should_continue_solve(integrator::SSAIntegrator) integrator.keep_stepping && (has_jump || has_tstop) end -function reset_aggregated_jumps!(integrator::SSAIntegrator, uprev = nothing) +function reset_aggregated_jumps!(integrator::SSAIntegrator, uprev = nothing; kwargs...) + if haskey(kwargs, :update_jump_params) + throw(ArgumentError("`update_jump_params` keyword argument has been removed. " * + "Rate updates are now handled automatically by `initialize!` " * + "via `fill_scaled_rates!`.")) + end reset_aggregated_jumps!(integrator, uprev, integrator.cb) nothing end diff --git a/src/aggregators/rssa.jl b/src/aggregators/rssa.jl index 22c50c922..23b273bd6 100644 --- a/src/aggregators/rssa.jl +++ b/src/aggregators/rssa.jl @@ -137,11 +137,6 @@ function generate_jumps!(p::RSSAJumpAggregation, integrator, u, params, t) nothing end -# alt erlang sampling above -#rerl = one(sum_rate) -#rerl *= rand(p.rng) -#p.next_jump_time = t + (-one(sum_rate) / sum_rate) * log(rerl) - ######################## SSA specific helper routines ######################### """ diff --git a/src/jumps.jl b/src/jumps.jl index dc5d316cc..b627ae033 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -762,6 +762,11 @@ function JumpSet(vj, cj, rj, T <: AbstractArray{<:Pair}, U <: AbstractArray{<:Pair}, V} JumpSet(vj, cj, rj, check_majump_type(maj)) end +function JumpSet(vj::AbstractJump, cj::AbstractJump, rj::AbstractJump, + maj::MassActionJump{Nothing, T, U, V}) where { + T <: AbstractArray{<:Pair}, U <: AbstractArray{<:Pair}, V} + JumpSet(vj, cj, rj, check_majump_type(maj)) +end JumpSet(jump::ConstantRateJump) = JumpSet((), (jump,), nothing, nothing) JumpSet(jump::VariableRateJump) = JumpSet((jump,), (), nothing, nothing) diff --git a/src/problem.jl b/src/problem.jl index 84bd0ad3d..29e23b6d2 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -319,6 +319,12 @@ function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; if haskey(kwargs, :rng) throw(ArgumentError("`rng` is no longer a keyword argument for `JumpProblem`. Pass `rng` to `solve` or `init` instead, e.g. `solve(jprob, SSAStepper(); rng = my_rng)`.")) end + if haskey(kwargs, :scale_rates) + throw(ArgumentError("`scale_rates` is no longer a keyword argument for `JumpProblem`. Set `scale_rates` on the `MassActionJump` directly instead.")) + end + if haskey(kwargs, :useiszero) + throw(ArgumentError("`useiszero` is no longer a keyword argument for `JumpProblem`. Set `useiszero` on the `MassActionJump` directly instead.")) + end # Validate no spatial systems (not currently supported) (spatial_system !== nothing || hopping_constants !== nothing) && diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index 27fe53ffd..a235f6d32 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -11,7 +11,8 @@ the timestep through the `dt` keyword to `solve`. The algorithm constructor has no fields or keyword arguments. The `solve` method accepts: - `dt`: Required fixed timestep. - - `seed`: Optional random seed for the jump problem RNG. + - `rng`: Optional RNG used by this solve; takes precedence over `seed`. + - `seed`: Optional seed used to construct a solve-local RNG. - `saveat`: Optional scalar interval or collection of save times. - `save_start`: Whether to save the initial time. Defaults follow SciML save conventions. - `save_end`: Whether to save the final time. Defaults follow SciML save conventions. @@ -127,6 +128,8 @@ end SimpleImplicitTauLeaping(; epsilon = 0.05) = SimpleImplicitTauLeaping(epsilon) +SciMLBase.supports_solve_rng(::JumpProblem, ::SimpleImplicitTauLeaping) = true + """ SimpleTrapezoidalLeaping(; epsilon = 0.05) @@ -169,6 +172,8 @@ end SimpleTrapezoidalLeaping(; epsilon = 0.05) = SimpleTrapezoidalLeaping(epsilon) +SciMLBase.supports_solve_rng(::JumpProblem, ::SimpleTrapezoidalLeaping) = true + """ $(TYPEDEF) @@ -234,6 +239,8 @@ function SimpleAdaptiveTauLeaping(; stiffness_ratio_threshold, implicit_epsilon_factor) end +SciMLBase.supports_solve_rng(::JumpProblem, ::SimpleAdaptiveTauLeaping) = true + 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 \ @@ -773,7 +780,7 @@ end function DiffEqBase.solve( jump_prob::JumpProblem, alg::Union{SimpleImplicitTauLeaping, SimpleTrapezoidalLeaping}; - seed = nothing, + seed = nothing, rng = nothing, dtmin = nothing, saveat = nothing, save_start = nothing, save_end = nothing ) @@ -781,18 +788,18 @@ function DiffEqBase.solve( error("$(nameof(typeof(alg))) can only be used with PureLeaping JumpProblem with a MassActionJump.") prob = jump_prob.prob - rng = jump_prob.rng + _rng = resolve_rng(rng, seed) tspan = prob.tspan if dtmin === nothing dtmin = 1.0e-10 * one(typeof(tspan[2])) end - (seed !== nothing) && seed!(rng, seed) - maj = jump_prob.massaction_jump numjumps = get_num_majumps(maj) - rate = massaction_rate(maj, numjumps) + maj_rates = Vector{typeof(tspan[2])}(undef, numjumps) + fill_scaled_rates!(maj_rates, maj, prob.p) + rate = massaction_rate(maj, maj_rates, numjumps) u0 = copy(prob.u0) p = prob.p @@ -828,7 +835,7 @@ function DiffEqBase.solve( ) simple_implicit_tau_leaping_loop!( - prob, alg, u_current, u_new, t_current, t_end, p, rng, + prob, alg, u_current, u_new, t_current, t_end, p, _rng, rate, nu, hor, max_hor, max_stoich, numjumps, epsilon, dtmin, saveat_times, usave, tsave, du, counts, rate_cache, rate_current, maj, save_end @@ -1024,7 +1031,7 @@ end function DiffEqBase.solve( jump_prob::JumpProblem, alg::SimpleAdaptiveTauLeaping; - seed = nothing, + seed = nothing, rng = nothing, dtmin = nothing, saveat = nothing, save_start = nothing, save_end = nothing ) @@ -1032,18 +1039,18 @@ function DiffEqBase.solve( error("SimpleAdaptiveTauLeaping can only be used with PureLeaping JumpProblem with a MassActionJump.") prob = jump_prob.prob - rng = jump_prob.rng + _rng = resolve_rng(rng, seed) tspan = prob.tspan if dtmin === nothing dtmin = 1.0e-10 * one(typeof(tspan[2])) end - (seed !== nothing) && seed!(rng, seed) - maj = jump_prob.massaction_jump numjumps = get_num_majumps(maj) - rate = massaction_rate(maj, numjumps) + maj_rates = Vector{typeof(tspan[2])}(undef, numjumps) + fill_scaled_rates!(maj_rates, maj, prob.p) + rate = massaction_rate(maj, maj_rates, numjumps) u0 = copy(prob.u0) p = prob.p @@ -1083,7 +1090,7 @@ function DiffEqBase.solve( ) simple_adaptive_tau_leaping_loop!( - prob, alg, u_current, u_new, t_current, t_end, p, rng, + prob, alg, u_current, u_new, t_current, t_end, p, _rng, rate, nu, hor, max_hor, max_stoich, numjumps, numspecies, alg.epsilon, dtmin, saveat_times, usave, tsave, du, counts, rate_cache, rate_effective, rate_current, maj, alg.implicit_alg, alg.eigenvalue_check, @@ -1114,6 +1121,15 @@ Ensemble algorithm marker for GPU execution of tau-leaping ensemble simulations. - `backend`: Backend object used by the GPU extension. - `cpu_offload`: Fraction of trajectories to offload to CPU execution. +## Random number generation + +Kernel ensembles use backend-local RNGs and reject explicit host `rng` and +`rng_func` keywords. +The `seed` solve keyword is supported with the KernelAbstractions CPU backend; +other backends require their own seeding API and reject `seed`. With +`trajectories = 1`, execution uses the serial CPU solver, which accepts its +usual `rng`, `seed`, and ensemble `rng_func` keywords. + ## Returns - A `SciMLBase.EnsembleAlgorithm` value for use as the ensemble algorithm argument to diff --git a/test/bracketing.jl b/test/bracketing.jl index 5d1841cbe..cec65800f 100644 --- a/test/bracketing.jl +++ b/test/bracketing.jl @@ -161,14 +161,14 @@ function nonmonotonic_jump_prob(alg) jump = ConstantRateJump(nonmonotonic_rate, nonmonotonic_affect!; bounds = nonmonotonic_bounds) prob = DiscreteProblem(u0, (0.0, tf)) - JumpProblem(prob, alg, jump; save_positions = (false, false), rng, vartojumps_map, + JumpProblem(prob, alg, jump; save_positions = (false, false), vartojumps_map, jumptovars_map) end function runSSAs(jump_prob) Asamp = zeros(Int, Nsims) for i in 1:Nsims - Asamp[i] = solve(jump_prob, SSAStepper())[1, end] + Asamp[i] = solve(jump_prob, SSAStepper(); rng)[1, end] end mean(Asamp) end diff --git a/test/explicit_tau_kernel_tests.jl b/test/explicit_tau_kernel_tests.jl index 8e0904416..1bb9e6df4 100644 --- a/test/explicit_tau_kernel_tests.jl +++ b/test/explicit_tau_kernel_tests.jl @@ -1,6 +1,7 @@ using JumpProcesses using Test, Statistics using KernelAbstractions, Adapt +include("kernel_contract_tests.jl") # Body of the SimpleExplicitTauLeaping kernel tests, run against both a GPU # backend and the KernelAbstractions CPU backend so the same behaviour is covered @@ -159,5 +160,8 @@ function run_explicit_tau_kernel_tests(backend, nsims) EnsembleGPUKernel(backend); trajectories = 5, saveat = 1.0) end + run_kernel_massaction_rate_tests(backend, SimpleExplicitTauLeaping(), PureLeaping()) + rng_prob = leaping_prob([5.0], [Pair{Int, Int}[]], [[1 => 1]], [0.0], (0.0, 4.0)) + run_kernel_rng_tests(backend, SimpleExplicitTauLeaping(), rng_prob; saveat = 1.0) return nothing end diff --git a/test/fp_unknowns.jl b/test/fp_unknowns.jl index 564643461..041a958ea 100644 --- a/test/fp_unknowns.jl +++ b/test/fp_unknowns.jl @@ -35,8 +35,7 @@ function test(rng) Ymeans = zeros(length(SSAalgs)) for (j, agg) in enumerate(SSAalgs) jprob = JumpProblem(dprob, agg, maj; save_positions = (false, false), - vartojumps_map = vtoj, jumptovars_map = jtov, dep_graph = dg, - scale_rates = false) + vartojumps_map = vtoj, jumptovars_map = jtov, dep_graph = dg) for i in 1:Nsims sol = solve(jprob, SSAStepper(); rng) Xmeans[j] += sol[1, end] diff --git a/test/kernel_contract_tests.jl b/test/kernel_contract_tests.jl new file mode 100644 index 000000000..ada4d5f15 --- /dev/null +++ b/test/kernel_contract_tests.jl @@ -0,0 +1,127 @@ +using Random +import SciMLBase + +# No allocation or kernel-launch methods: seed rejection must happen before +# either is needed, even when no actual GPU is available to the test process. +struct UnseedableKernelBackend <: KernelAbstractions.GPU end + +function run_kernel_rng_tests(backend, alg, jump_prob; kwargs...) + ensemble_prob = EnsembleProblem(jump_prob) + # Use enough trajectories for the CPU backend to launch multiple tasks in a + # threaded test process, so replay covers task-local RNG inheritance too. + kernel_solve(; options...) = solve(ensemble_prob, alg, EnsembleGPUKernel(backend); + trajectories = 4096, kwargs..., options...) + paths(sol) = [(traj.t, traj.u) for traj in sol.u] + + @testset "Kernel RNG inputs: $(typeof(alg))" begin + rng = Xoshiro(17) + untouched_rng = copy(rng) + @test_throws ArgumentError kernel_solve(; rng) + @test rand(rng, UInt) == rand(untouched_rng, UInt) + @test_throws ArgumentError kernel_solve(; rng = Xoshiro(17), seed = 123) + @test_throws ArgumentError kernel_solve(; rng_func = ctx -> Xoshiro(17)) + + if backend isa CPU + sol1 = kernel_solve(; seed = 123) + sol2 = kernel_solve(; seed = 123) + sol3 = kernel_solve(; seed = 456) + @test paths(sol1) == paths(sol2) + @test paths(sol1) != paths(sol3) + else + @test_throws ArgumentError kernel_solve(; seed = 123) + end + + # An unsupported device seed must not reseed the host generator. + Random.seed!(321) + expected = rand(UInt) + Random.seed!(321) + @test_throws ArgumentError solve(ensemble_prob, alg, + EnsembleGPUKernel(UnseedableKernelBackend()); + trajectories = 16, seed = 123, kwargs...) + @test rand(UInt) == expected + + # A single trajectory still delegates all RNG controls to EnsembleSerial, + # including the ensemble layer's precedence when both rng and seed occur. + for rng_options in (() -> (; seed = 123), + () -> (; rng = Xoshiro(17)), + () -> (; rng = Xoshiro(17), seed = 123), + () -> (; seed = 123, rng_func = ctx -> Xoshiro(17))) + kernel = solve(ensemble_prob, alg, EnsembleGPUKernel(backend); + trajectories = 1, kwargs..., rng_options()...) + serial = solve(ensemble_prob, alg, EnsembleSerial(); + trajectories = 1, kwargs..., rng_options()...) + @test paths(kernel) == paths(serial) + @test all(SciMLBase.successful_retcode, kernel.u) + end + end +end + +function run_kernel_massaction_rate_tests(backend, alg, aggregator) + ext = Base.get_extension(JumpProcesses, :JumpProcessesKernelAbstractionsExt) + + @testset "Kernel current-parameter rates: $(typeof(alg))" begin + for RT in (Float32, Float64) + rs = [Pair{Int, Int}[], [1 => 1], [1 => 2], [1 => 3]] + ns = [[2 => 1], [1 => -1, 2 => 1], [1 => -2, 2 => 1], [1 => -3, 2 => 1]] + params = RT[1, 2, 6, 24] + fixed = MassActionJump(params, rs, ns) + mapped = MassActionJump(rs, ns; param_idxs = [1, 2, 3, 4]) + fixed_gpu = ext.GPUMassActionJump(fixed, params, backend, RT) + mapped_gpu = ext.GPUMassActionJump(mapped, params, backend, RT) + @test Array(fixed_gpu.scaled_rates) == RT[1, 2, 3, 4] + @test Array(mapped_gpu.scaled_rates) == RT[1, 2, 3, 4] + @test eltype(mapped_gpu.scaled_rates) === RT + @test mapped.scaled_rates === nothing + @test params == RT[1, 2, 6, 24] + + # Refreshing a conversion leaves earlier device data and the fixed + # definition alone. Fixed constants do not acquire dependence on p. + params .*= 2 + refreshed = ext.GPUMassActionJump(mapped, params, backend, RT) + fixed_again = ext.GPUMassActionJump(fixed, params, backend, RT) + @test Array(refreshed.scaled_rates) == RT[2, 4, 6, 8] + @test Array(mapped_gpu.scaled_rates) == RT[1, 2, 3, 4] + @test Array(fixed_again.scaled_rates) == RT[1, 2, 3, 4] + @test fixed.scaled_rates == RT[1, 2, 3, 4] + + # A custom mapper owns its scaling policy, including pre-scaled rates. + prescaled = MassActionJump(rs, ns; + param_mapper = (dest, maj, p) -> (dest .= p), + rescale_rates_on_update = false) + custom_gpu = ext.GPUMassActionJump(prescaled, params, backend, RT) + @test Array(custom_gpu.scaled_rates) == params + + # Exercise each solver's conversion call, including a changed-parameter + # remake of a third-order reaction with a scalar parameter index. + maj = MassActionJump([[1 => 3]], [[1 => -3, 2 => 1]]; param_idxs = 1) + u0 = RT[30, 0] + jp = JumpProblem(DiscreteProblem(u0, (zero(RT), one(RT)), RT[0]), + aggregator, maj; save_positions = (false, false)) + active = remake(jp; p = RT[0.06]) + inactive_sol = solve(EnsembleProblem(jp), alg, EnsembleGPUKernel(backend); + trajectories = 16, saveat = one(RT)) + active_sol = solve(EnsembleProblem(active), alg, EnsembleGPUKernel(backend); + trajectories = 16, saveat = one(RT)) + @test all(traj -> all(u -> u == u0, traj.u), inactive_sol.u) + @test any(traj -> traj.u[end][2] > 0, active_sol.u) + @test all(traj -> all(u -> u[1] + 3u[2] == 30, traj.u), active_sol.u) + @test all(SciMLBase.successful_retcode, active_sol.u) + @test jp.prob.p == RT[0] + @test jp.massaction_jump.scaled_rates === nothing + @test active.massaction_jump.scaled_rates === nothing + + active.prob.p[1] = RT(0.12) + updated_gpu = ext.GPUMassActionJump( + active.massaction_jump, active.prob.p, backend, RT) + @test Array(updated_gpu.scaled_rates) ≈ RT[0.02] + end + end +end + +function run_regular_kernel_rng_tests(backend) + birth_rate!(out, u, p, t) = (out[1] = p[1]) + birth_c!(du, u, p, t, counts, mark) = (du[1] = counts[1]) + rj = RegularJump(birth_rate!, birth_c!, 1) + jp = JumpProblem(DiscreteProblem([0.0], (0.0, 2.0), (20.0,)), PureLeaping(), rj) + run_kernel_rng_tests(backend, SimpleTauLeaping(), jp; dt = 0.25) +end diff --git a/test/massaction_state.jl b/test/massaction_state.jl new file mode 100644 index 000000000..83ad47d13 --- /dev/null +++ b/test/massaction_state.jl @@ -0,0 +1,147 @@ +using JumpProcesses, Test + +@testset "Mass-action working rates follow current parameters" begin + # 3A -> B has a 3! scaling factor, which must be applied exactly once. + reactant_stoch = [[1 => 3]] + net_stoch = [[1 => -3, 2 => 1]] + @testset "$(typeof(algorithm))" for algorithm in JumpProcesses.JUMP_AGGREGATORS + maj = MassActionJump(reactant_stoch, net_stoch; param_idxs = 1) + prob = DiscreteProblem([9, 0], (0.0, 0.1), [6.0]) + jprob = JumpProblem(prob, algorithm, maj) + # A parameter change after construction must be seen by initialization. + prob.p[1] = 12.0 + integrator = init(jprob, SSAStepper(); seed = 12345, alias_jump = true) + @test jprob.discrete_jump_aggregation.maj_rates == [2.0] + integrator.p[1] = 18.0 + reset_aggregated_jumps!(integrator) + @test jprob.discrete_jump_aggregation.maj_rates == [3.0] + @test jprob.massaction_jump.scaled_rates === nothing + + remade = remake(jprob; p = [24.0]) + init(remade, SSAStepper(); seed = 12345, alias_jump = true) + @test remade.discrete_jump_aggregation.maj_rates == [4.0] + @test jprob.prob.p == [18.0] + @test remade.massaction_jump === jprob.massaction_jump + end +end + +@testset "Spatial mass-action working rates follow current parameters" begin + @testset "$(typeof(algorithm))" for algorithm in (NSM(), DirectCRDirect()) + maj = MassActionJump([[1 => 3]], [[1 => -3, 2 => 1]]; param_idxs = 1) + prob = DiscreteProblem([6 9; 0 0], (0.0, 0.1), [6.0]) + jprob = JumpProblem(prob, algorithm, maj; + spatial_system = CartesianGrid((2,)), hopping_constants = zeros(2)) + prob.p[1] = 12.0 + integrator = init(jprob, SSAStepper(); seed = 12345, alias_jump = true) + rx_rates = jprob.discrete_jump_aggregation.rx_rates + @test rx_rates.maj_rates == [2.0] + @test rx_rates.rates == [240.0 1008.0] + integrator.p[1] = 18.0 + reset_aggregated_jumps!(integrator) + @test rx_rates.maj_rates == [3.0] + @test rx_rates.rates == [360.0 1512.0] + @test jprob.massaction_jump.scaled_rates === nothing + + remade = remake(jprob; p = [24.0]) + init(remade, SSAStepper(); seed = 12345, alias_jump = true) + @test remade.discrete_jump_aggregation.rx_rates.maj_rates == [4.0] + @test remade.discrete_jump_aggregation.rx_rates.rates == [480.0 2016.0] + @test jprob.prob.p == [18.0] + end +end + +@testset "Flattening materializes current scaled rates once" begin + @testset "$parameterized" for parameterized in (false, true) + reactant_stoch = [[1 => 3]] + net_stoch = [[1 => -3, 2 => 1]] + maj = parameterized ? MassActionJump(reactant_stoch, net_stoch; param_idxs = 1) : + MassActionJump([12.0], reactant_stoch, net_stoch) + prob = DiscreteProblem([6 9; 0 0], (0.0, 0.1), [6.0]) + prob.p[1] = 12.0 + jprob = JumpProblem(prob, Direct(), maj; + spatial_system = CartesianGrid((2,)), hopping_constants = zeros(2, 2)) + @test jprob.massaction_jump.scaled_rates == [0.0, 0.0, 0.0, 0.0, 2.0, 2.0] + @test jprob.prob.u0 == [6, 0, 9, 0] + init(jprob, SSAStepper(); seed = 12345, alias_jump = true) + @test jprob.discrete_jump_aggregation.maj_rates == jprob.massaction_jump.scaled_rates + if parameterized + @test maj.scaled_rates === nothing + else + @test maj.scaled_rates == [2.0] + end + end +end + +@testset "Mass-action merges preserve source arrays and parameter indices" begin + @testset "$parameterized, $collection" for parameterized in (false, true), + collection in (false, true) + rs1, rs2 = [[1 => 3]], [[2 => 1]] + ns1, ns2 = [[1 => -3, 2 => 1]], [[2 => -1, 1 => 3]] + idxs1, idxs2 = [1], [2] + maj1 = parameterized ? MassActionJump(rs1, ns1; param_idxs = idxs1) : + MassActionJump([6.0], rs1, ns1) + maj2 = parameterized ? MassActionJump(rs2, ns2; param_idxs = idxs2) : + MassActionJump([2.0], rs2, ns2) + jset = collection ? JumpSet(; massaction_jumps = [maj1, maj2]) : JumpSet(maj1, maj2) + @test maj1.reactant_stoch == [[1 => 3]] + @test maj2.reactant_stoch == [[2 => 1]] + @test ns1 == [[1 => -3, 2 => 1]] + @test ns2 == [[2 => -1, 1 => 3]] + @test idxs1 == [1] + @test idxs2 == [2] + @test jset.massaction_jump.net_stoch == [only(ns1), only(ns2)] + dest = zeros(2) + JumpProcesses.fill_scaled_rates!(dest, jset.massaction_jump, [6.0, 2.0]) + @test dest == [1.0, 2.0] + if parameterized + @test maj1.param_mapper.param_idxs == [1] + @test maj2.param_mapper.param_idxs == [2] + @test jset.massaction_jump.param_mapper.param_idxs == [1, 2] + else + @test maj1.scaled_rates == [1.0] + @test maj2.scaled_rates == [2.0] + end + end +end + +@testset "Removed reset rate-refresh keyword is diagnosed" begin + maj = MassActionJump([[1 => 1]], [[1 => -1]]; param_idxs = 1) + jprob = JumpProblem(DiscreteProblem([5], (0.0, 1.0), [1.0]), Direct(), maj) + integrator = init(jprob, SSAStepper(); seed = 12345) + for update_jump_params in (false, true) + @test_throws ArgumentError reset_aggregated_jumps!(integrator; update_jump_params) + @test_throws r"update_jump_params.*removed" reset_aggregated_jumps!(integrator; + update_jump_params) + end +end + +@testset "Four-jump construction with a scalar mass-action jump" begin + rate(u, p, t) = 1.0 + affect!(integrator) = nothing + crj = ConstantRateJump(rate, affect!) + for parameterized in (false, true) + maj = parameterized ? MassActionJump([1 => 3], [1 => -3]; param_idxs = 1) : + MassActionJump(6.0, [1 => 3], [1 => -3]) + jset = JumpSet(crj, crj, crj, maj) + @test jset.constant_jumps == (crj, crj, crj) + @test jset.variable_jumps == () + @test jset.regular_jump === nothing + @test jset.massaction_jump.reactant_stoch == [[1 => 3]] + dest = zeros(1) + JumpProcesses.fill_scaled_rates!(dest, jset.massaction_jump, [6.0]) + @test dest == [1.0] + end +end + +@testset "Removed JumpProblem rate keywords are diagnosed" begin + maj = MassActionJump([[1 => 3]], [[1 => -3]]; param_idxs = 1) + prob = DiscreteProblem([6], (0.0, 1.0), [6.0]) + for algorithm in (Direct(), PureLeaping()), value in (false, true) + @test_throws ArgumentError JumpProblem(prob, algorithm, maj; scale_rates = value) + @test_throws r"scale_rates.*no longer.*MassActionJump" JumpProblem(prob, algorithm, + maj; scale_rates = value) + @test_throws ArgumentError JumpProblem(prob, algorithm, maj; useiszero = value) + @test_throws r"useiszero.*no longer.*MassActionJump" JumpProblem(prob, algorithm, + maj; useiszero = value) + end +end diff --git a/test/regular_jumps.jl b/test/regular_jumps.jl index 0142a5a86..eaeef4dc9 100644 --- a/test/regular_jumps.jl +++ b/test/regular_jumps.jl @@ -1,6 +1,6 @@ using JumpProcesses, DiffEqBase, SciMLBase using Test, LinearAlgebra, Statistics -using StableRNGs +using StableRNGs, Random rng = StableRNG(12345) @testset "Successful leaping return codes" begin @@ -15,9 +15,9 @@ rng = StableRNG(12345) ) @testset "$(nameof(typeof(alg)))" begin jump = alg isa SimpleTauLeaping ? regular_jump : massaction_jump - jp = JumpProblem(prob, PureLeaping(), jump; rng = StableRNG(12345)) + jp = JumpProblem(prob, PureLeaping(), jump) kwargs = alg isa SimpleTauLeaping ? (; dt = 0.01) : (;) - sol = solve(jp, alg; kwargs...) + sol = solve(jp, alg; rng = StableRNG(12345), kwargs...) @test sol.t[end] == last(prob.tspan) @test sol.retcode == ReturnCode.Success @test successful_retcode(sol) @@ -25,6 +25,140 @@ rng = StableRNG(12345) end end +@testset "CPU tau-leaping solve RNG ownership" begin + prob = DiscreteProblem([0.0], (0.0, 2.0)) + regular_jump = RegularJump( + (out, u, p, t) -> (out[1] = 20.0), + (du, u, p, t, counts, mark) -> (du[1] = counts[1]), 1 + ) + massaction_jump = MassActionJump([20.0], [Pair{Int, Int}[]], [[1 => 1]]) + for alg in ( + SimpleTauLeaping(), SimpleExplicitTauLeaping(), + SimpleImplicitTauLeaping(), SimpleTrapezoidalLeaping(), + SimpleAdaptiveTauLeaping(), + SimpleAdaptiveTauLeaping(implicit_alg = SimpleTrapezoidalLeaping()), + ) + @testset "$alg" begin + jump = alg isa SimpleTauLeaping ? regular_jump : massaction_jump + jp = JumpProblem(prob, PureLeaping(), jump) + kwargs = alg isa SimpleTauLeaping ? (; dt = 0.05) : (;) + @test SciMLBase.supports_solve_rng(jp, alg) + + # Explicit RNGs are consumed, and override seed without being reseeded. + rng = StableRNG(1234) + before = copy(rng) + sol_rng = solve(jp, alg; rng, seed = 999, saveat = 0.5, kwargs...) + sol_rng_replay = solve(jp, alg; rng = before, saveat = 0.5, kwargs...) + @test sol_rng.t == sol_rng_replay.t + @test sol_rng.u == sol_rng_replay.u + next_draw = rand(rng) + @test next_draw == rand(before) + @test next_draw != rand(StableRNG(1234)) + + # A seed creates a fresh Xoshiro on every solve, including problem reuse. + sol_seed = solve(jp, alg; seed = 1234, saveat = 0.5, kwargs...) + sol_seed_replay = solve(jp, alg; seed = 1234, saveat = 0.5, kwargs...) + sol_xoshiro = solve(jp, alg; rng = Xoshiro(1234), saveat = 0.5, kwargs...) + sol_changed = solve(jp, alg; seed = 5678, saveat = 0.5, kwargs...) + @test sol_seed.t == sol_seed_replay.t == sol_xoshiro.t + @test sol_seed.u == sol_seed_replay.u == sol_xoshiro.u + @test sol_seed.u != sol_changed.u + @test successful_retcode(solve(jp, alg; kwargs...)) + + # The support trait lets ensembles forward their trajectory RNGs. + ens1 = solve(EnsembleProblem(jp), alg, EnsembleSerial(); + trajectories = 3, rng = StableRNG(1234), saveat = 0.5, kwargs...) + ens2 = solve(EnsembleProblem(jp), alg, EnsembleSerial(); + trajectories = 3, rng = StableRNG(1234), saveat = 0.5, kwargs...) + @test [sol.u for sol in ens1.u] == [sol.u for sol in ens2.u] + @test length(unique([sol.u for sol in ens1.u])) > 1 + end + end +end + +@testset "CPU tau-leaping current mass-action parameters" begin + reactants = [Pair{Int, Int}[], [1 => 2], [1 => 3]] + net = [[1 => 1], [1 => -2], [1 => -3]] + original_rates = [5.0, 0.01, 0.001] + updated_rates = [10.0, 0.02, 0.002] + maj = MassActionJump(reactants, net; param_idxs = [1, 2, 3]) + # A custom mapper producing already-scaled rates must not be scaled again. + prescaled_mapper = (dest, maj, params) -> (dest .= params ./ [1, 2, 6]; nothing) + custom_maj = MassActionJump(reactants, net; + param_mapper = prescaled_mapper, scale_rates = false) + for alg in ( + SimpleExplicitTauLeaping(), SimpleImplicitTauLeaping(), + SimpleTrapezoidalLeaping(), SimpleAdaptiveTauLeaping(), + SimpleAdaptiveTauLeaping(implicit_alg = SimpleTrapezoidalLeaping()), + ) + @testset "$alg" begin + prob = DiscreteProblem([30.0], (0.0, 2.0), copy(original_rates)) + jp = JumpProblem(prob, PureLeaping(), maj) + fixed = JumpProblem(prob, PureLeaping(), + MassActionJump(original_rates, reactants, net)) + baseline = solve(jp, alg; seed = 1234, saveat = 0.5) + fixed_sol = solve(fixed, alg; seed = 1234, saveat = 0.5) + @test baseline.t == fixed_sol.t + @test baseline.u == fixed_sol.u + + updated = remake(jp; p = copy(updated_rates)) + updated_fixed = JumpProblem(updated.prob, PureLeaping(), + MassActionJump(updated_rates, reactants, net)) + updated_custom = JumpProblem(updated.prob, PureLeaping(), custom_maj) + expected = solve(updated_fixed, alg; seed = 1234, saveat = 0.5) + remade_sol = solve(updated, alg; seed = 1234, saveat = 0.5) + custom_sol = solve(updated_custom, alg; seed = 1234, saveat = 0.5) + @test remade_sol.t == expected.t == custom_sol.t + @test remade_sol.u == expected.u == custom_sol.u + @test remade_sol.u != baseline.u + @test solve(jp, alg; seed = 1234, saveat = 0.5).u == baseline.u + + # Parameter mutation before a new solve also refreshes the working rates. + jp.prob.p .= updated_rates + @test solve(jp, alg; seed = 1234, saveat = 0.5).u == expected.u + # Fixed coefficients continue to ignore parameters, even after mutation. + @test solve(fixed, alg; seed = 1234, saveat = 0.5).u == baseline.u + @test maj.scaled_rates === nothing + @test custom_maj.scaled_rates === nothing + @test original_rates == [5.0, 0.01, 0.001] + @test updated_rates == [10.0, 0.02, 0.002] + @test successful_retcode(remade_sol) + end + end +end + +@testset "Parameterized tau-leaping implicit drift and Jacobian" begin + # The nonlinear solver uses finite differences with floating trial states, + # including when the population input is integer-valued. + maj = MassActionJump([[1 => 1]], [[1 => -1]]; param_idxs = 1) + p = [0.4] + rates = zeros(1) + JumpProcesses.fill_scaled_rates!(rates, maj, p) + rate = JumpProcesses.massaction_rate(maj, rates, 1) + nu = reshape([-1.0], 1, 1) + for (alg, expected) in ( + (SimpleImplicitTauLeaping(), 100 / 1.2), + (SimpleTrapezoidalLeaping(), 100 * 0.9 / 1.1), + ) + predicted, converged = JumpProcesses.solve_implicit( + [100], zeros(1), zeros(1), nu, p, 0.0, 0.5, rate, 1, alg) + @test converged + @test predicted ≈ [expected] + end + + # Three reactions and two species produce a rectangular propensity Jacobian. + # The drift Jacobian is square and includes the scaled third-order derivative. + maj = MassActionJump([[1 => 1], [2 => 1], [1 => 3]], + [[1 => -1], [2 => -1], [1 => -3]]; param_idxs = [1, 2, 3]) + p = [0.4, 0.1, 0.006] + rates = zeros(3) + JumpProcesses.fill_scaled_rates!(rates, maj, p) + rate = JumpProcesses.massaction_rate(maj, rates, 3) + nu = [-1.0 0.0 -3.0; 0.0 -1.0 0.0] + jac = JumpProcesses.compute_drift_jacobian([10, 20], rate, nu, 3, 2, p, 0.0) + @test jac ≈ [-0.4 - 3 * 0.001 * (3 * 10^2 - 6 * 10 + 2) 0.0; 0.0 -0.1] +end + @test SimpleImplicitTauLeaping().epsilon == 0.05 @test SimpleTrapezoidalLeaping().epsilon == 0.05 @@ -101,24 +235,24 @@ end sol_implicit = solve( EnsembleProblem(jump_prob_maj), SimpleImplicitTauLeaping(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare + trajectories = Nsims, saveat = t_compare, rng ) sol_trapezoidal = solve( EnsembleProblem(jump_prob_maj), SimpleTrapezoidalLeaping(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare + trajectories = Nsims, saveat = t_compare, rng ) sol_adaptive_implicit = solve( EnsembleProblem(jump_prob_maj), SimpleAdaptiveTauLeaping(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare + trajectories = Nsims, saveat = t_compare, rng ) sol_adaptive_trapezoidal = solve( EnsembleProblem(jump_prob_maj), SimpleAdaptiveTauLeaping(implicit_alg = SimpleTrapezoidalLeaping()), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare + trajectories = Nsims, saveat = t_compare, rng ) mean_I_direct = compute_mean_at_saves(sol_direct, Nsims, npts, 2) @@ -198,24 +332,24 @@ end sol_implicit = solve( EnsembleProblem(jump_prob_maj), SimpleImplicitTauLeaping(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare + trajectories = Nsims, saveat = t_compare, rng ) sol_trapezoidal = solve( EnsembleProblem(jump_prob_maj), SimpleTrapezoidalLeaping(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare + trajectories = Nsims, saveat = t_compare, rng ) sol_adaptive_implicit = solve( EnsembleProblem(jump_prob_maj), SimpleAdaptiveTauLeaping(), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare + trajectories = Nsims, saveat = t_compare, rng ) sol_adaptive_trapezoidal = solve( EnsembleProblem(jump_prob_maj), SimpleAdaptiveTauLeaping(implicit_alg = SimpleTrapezoidalLeaping()), EnsembleSerial(); - trajectories = Nsims, saveat = t_compare + trajectories = Nsims, saveat = t_compare, rng ) mean_I_direct = compute_mean_at_saves(sol_direct, Nsims, npts, 3) @@ -263,13 +397,13 @@ end net_stoich = [[1 => -1, 2 => 1], [2 => -1, 3 => 1], [3 => -1, 4 => 1]] maj = MassActionJump([0.3 / 1000, 0.2, 0.1], reactant_stoich, net_stoich) prob = DiscreteProblem([999.0, 0.0, 10.0, 0.0], (0.0, 20.0)) - jprob = JumpProblem(prob, PureLeaping(), maj; rng) + jprob = JumpProblem(prob, PureLeaping(), maj) total0 = sum(prob.u0) for implicit_alg in (SimpleImplicitTauLeaping(), SimpleTrapezoidalLeaping()) for eigenvalue_check in (false, true) alg = SimpleAdaptiveTauLeaping(; implicit_alg, eigenvalue_check) - sol = solve(jprob, alg; seed = 1234) + sol = solve(jprob, alg; rng = StableRNG(1234)) @test sol.t[end] == 20.0 @test all(u -> all(>=(0), u), sol.u) # every reaction moves one individual between compartments @@ -284,12 +418,12 @@ end # which threw a BoundsError for such a model. maj = MassActionJump([5.0], [Pair{Int, Int}[]], [[1 => 1]]) prob = DiscreteProblem([0.0], (0.0, 4.0)) - jprob = JumpProblem(prob, PureLeaping(), maj; rng) + jprob = JumpProblem(prob, PureLeaping(), maj) @test JumpProcesses.compute_hor([Pair{Int, Int}[], [1 => 2]], 2) == [0, 2] for alg in (SimpleExplicitTauLeaping(), SimpleImplicitTauLeaping()) - sol = solve(jprob, alg; seed = 1234, saveat = 2.0) + sol = solve(jprob, alg; rng = StableRNG(1234), saveat = 2.0) @test sol.t == [0.0, 2.0, 4.0] @test issorted([u[1] for u in sol.u]) # A can only accumulate end @@ -297,7 +431,7 @@ end # mean of A(t) is the Poisson mean 5t Nsims = 4000 sol = solve(EnsembleProblem(jprob), SimpleExplicitTauLeaping(), EnsembleSerial(); - trajectories = Nsims, saveat = 4.0) + trajectories = Nsims, saveat = 4.0, rng) @test isapprox(mean(sol.u[i].u[end][1] for i in 1:Nsims), 20.0, rtol = 0.05) end diff --git a/test/runtests.jl b/test/runtests.jl index f2af5b418..bd909d33a 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -24,6 +24,7 @@ end @time @safetestset "Split Coupled Tests" begin include("splitcoupled.jl") end @time @safetestset "SSA Tests" begin include("ssa_tests.jl") end @time @safetestset "SSA Aggregator Smoke Tests" begin include("ssa_aggregator_smoke.jl") end + @time @safetestset "SortingDirect retained search order" begin include("sortingdirect_reuse.jl") end @time @safetestset "Tau Leaping Tests" begin include("regular_jumps.jl") end @time @safetestset "Simple SSA Callback Test" begin include("ssa_callback_test.jl") end @time @safetestset "SIR Discrete Callback Test" begin include("sir_model.jl") end @@ -34,6 +35,7 @@ end @time @safetestset "Mass Action Jump Tests; Special Cases" begin include("degenerate_rx_cases.jl") end @time @safetestset "Mass Action Jump Tests; Floating Point Inputs" begin include("fp_unknowns.jl") end @time @safetestset "scale_rates Field Tests" begin include("scale_rates_field_test.jl") end + @time @safetestset "Mass-action working state" begin include("massaction_state.jl") end @time @safetestset "Direct allocations test" begin include("allocations.jl") end @time @safetestset "Bracketing Tests" begin include("bracketing.jl") end @time @safetestset "Composition-Rejection Table Tests" begin include("table_test.jl") end diff --git a/test/sortingdirect_reuse.jl b/test/sortingdirect_reuse.jl new file mode 100644 index 000000000..134461563 --- /dev/null +++ b/test/sortingdirect_reuse.jl @@ -0,0 +1,58 @@ +using JumpProcesses, SciMLBase, Test + +function sortingdirect_reuse_problem(representation) + prob = DiscreteProblem(zeros(Int, 3), (0.0, 2.0), [1.0, 2.0, 9.0]) + reactant_stoch = [Pair{Int, Int}[] for _ in 1:3] + net_stoch = [[i => 1] for i in 1:3] + if representation === :parameterized_massaction + maj = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1, 2, 3]) + return JumpProblem(prob, SortingDirect(), maj) + elseif representation === :fixed_massaction + maj = MassActionJump(prob.p, reactant_stoch, net_stoch) + return JumpProblem(prob, SortingDirect(), maj) + end + + jumps = ntuple(3) do i + rate(u, p, t) = p[i] + function affect!(integrator) + integrator.u[i] += 1 + nothing + end + ConstantRateJump(rate, affect!) + end + JumpProblem(prob, SortingDirect(), jumps...; dep_graph = [[i] for i in 1:3]) +end + +@testset "SortingDirect preserves learned search order" begin + @testset "$representation" for representation in + (:parameterized_massaction, :fixed_massaction, :constant) + jprob = sortingdirect_reuse_problem(representation) + sol = solve(jprob, SSAStepper(); seed = 12345, alias_jump = true) + @test SciMLBase.successful_retcode(sol) + @test sum(last(sol.u)) > 0 + + # Seeded replay requires identical initial search order, as well as RNG + # state. Each copy begins with the order learned by the completed solve. + replay = solve(deepcopy(jprob), SSAStepper(); seed = 12345, alias_jump = true) + repeated = solve(deepcopy(jprob), SSAStepper(); seed = 12345, alias_jump = true) + @test replay.t == repeated.t + @test replay.u == repeated.u + + # Initialization and rate resets retain the ordering optimization. + aggregation = jprob.discrete_jump_aggregation + aggregation.jump_search_order .= [3, 1, 2] + integrator = init(jprob, SSAStepper(); seed = 12345, alias_jump = true) + @test aggregation.jump_search_order == [3, 1, 2] + if representation !== :fixed_massaction + integrator.p .= [3.0, 4.0, 5.0] + end + reset_aggregated_jumps!(integrator) + @test aggregation.jump_search_order == [3, 1, 2] + @test aggregation.cur_rates == + (representation === :fixed_massaction ? [1.0, 2.0, 9.0] : [3.0, 4.0, 5.0]) + if representation === :parameterized_massaction + @test aggregation.maj_rates == integrator.p + @test jprob.massaction_jump.scaled_rates === nothing + end + end +end diff --git a/test/ssa_aggregator_smoke.jl b/test/ssa_aggregator_smoke.jl index 2c207e9c6..8674772f8 100644 --- a/test/ssa_aggregator_smoke.jl +++ b/test/ssa_aggregator_smoke.jl @@ -12,8 +12,8 @@ function check_ssa_smoke_solution(sol, prob) end function check_ssa_smoke_problem(jprob) - # Independent trials start with identical mutable aggregator state. In - # particular, SortingDirect retains its learned search order between solves. + # Independent trials start with identical mutable aggregator state. SortingDirect + # intentionally retains its learned search order; see sortingdirect_reuse.jl. sol_default = solve(deepcopy(jprob), SSAStepper()) check_ssa_smoke_solution(sol_default, jprob.prob) diff --git a/test/ssa_kernel_tests.jl b/test/ssa_kernel_tests.jl index 81328c0e1..58beb3bcc 100644 --- a/test/ssa_kernel_tests.jl +++ b/test/ssa_kernel_tests.jl @@ -1,6 +1,7 @@ using JumpProcesses using Test, Statistics using KernelAbstractions, Adapt +include("kernel_contract_tests.jl") # Body of the SSA kernel tests, run against both a GPU backend and the # KernelAbstractions CPU backend so the same behaviour is covered without a GPU. @@ -197,5 +198,9 @@ function run_ssa_kernel_tests(backend, nsims) @test_throws ErrorException solve(varying, SSAStepper(), EnsembleGPUKernel(backend); trajectories = 5, saveat = 1.0) end + run_kernel_massaction_rate_tests(backend, SSAStepper(), Direct()) + rng_prob = majump_prob([5.0], [Pair{Int, Int}[]], [[1 => 1]], [0], (0.0, 4.0)) + run_kernel_rng_tests(backend, SSAStepper(), rng_prob; saveat = 1.0) + run_regular_kernel_rng_tests(backend) return nothing end diff --git a/test/variable_rate.jl b/test/variable_rate.jl index 70e469924..6dfd996fc 100644 --- a/test/variable_rate.jl +++ b/test/variable_rate.jl @@ -634,9 +634,9 @@ let u0 = [150.0] tspan = (0.0, 10.0) sde_prob = SDEProblem(f!, g!, u0, tspan, p; noise_rate_prototype = zeros(1, 2)) - jprob = JumpProblem(sde_prob, Direct(), birth, death; vr_aggregator = VR_FRM(), rng) + jprob = JumpProblem(sde_prob, Direct(), birth, death; vr_aggregator = VR_FRM()) - sol = solve(jprob, LambaEM(); callback = switch_cb, adaptive = true) + sol = solve(jprob, LambaEM(); rng, callback = switch_cb, adaptive = true) @test SciMLBase.successful_retcode(sol) @test sol.t[end] == tspan[2] # The bug drove `jump_u` past ±1e15 within a few steps; with the fix it From 323896a8484a172a09d4c172d90ce9e6c0887017 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 29 Sep 2026 09:45:02 -0400 Subject: [PATCH 33/52] Document v10 migration and maintain CI coverage Explain solve-level and kernel RNG controls, mass-action working rates, custom mapper scaling, removed APIs, and mutable jump-state isolation. Add a migration guide and keep the 9.x integration version metadata. Preserve grouped, 32-bit, GPU, downgrade, and dedicated four-thread CI coverage while removing staging triggers and aligning action versions. Strengthen ensemble replay and uniqueness checks, and fix trajectory counts in the thread-safety tests. Validation: 102 focused Julia assertions and 11 static checks passed. The documentation review follow-up passed all 21 affected checks; runtime, test, and CI contents remained unchanged. The broader validation matrix and full documentation build remain pending. --- .github/workflows/Downgrade.yml | 1 - .github/workflows/Tests.yml | 1 - .github/workflows/ThreadSafety.yml | 7 +- HISTORY.md | 69 ++++++-- docs/pages.jl | 1 + docs/src/api.md | 32 +++- docs/src/faq.md | 11 +- docs/src/jump_types.md | 20 ++- docs/src/migration.md | 164 ++++++++++++++++++ .../tutorials/discrete_stochastic_example.md | 5 + ext/explicit_tau_leaping.jl | 17 +- ext/ssa_stepper.jl | 17 +- src/jumps.jl | 33 ++-- src/problem.jl | 12 +- src/simple_regular_solve.jl | 20 ++- test/ensemble_uniqueness.jl | 12 +- test/thread_safety.jl | 49 ++++-- 17 files changed, 397 insertions(+), 74 deletions(-) create mode 100644 docs/src/migration.md diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml index 37ccde5fe..26c77b953 100644 --- a/.github/workflows/Downgrade.yml +++ b/.github/workflows/Downgrade.yml @@ -3,7 +3,6 @@ on: pull_request: branches: - master - - v10_working paths-ignore: - 'docs/**' push: diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index 50c77dd50..45dc20491 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -8,7 +8,6 @@ on: pull_request: branches: - master - - v10_working paths-ignore: - 'docs/**' diff --git a/.github/workflows/ThreadSafety.yml b/.github/workflows/ThreadSafety.yml index 2dfbb4643..36c7220e2 100644 --- a/.github/workflows/ThreadSafety.yml +++ b/.github/workflows/ThreadSafety.yml @@ -9,7 +9,6 @@ on: pull_request: branches: - master - - v10_working paths-ignore: - 'docs/**' @@ -29,10 +28,10 @@ jobs: - "pre" runs-on: ubuntu-latest steps: - - uses: actions/checkout@v4 + - uses: actions/checkout@v7 - name: "Setup Julia ${{ matrix.version }}" - uses: julia-actions/setup-julia@v2 + uses: julia-actions/setup-julia@v3 with: version: "${{ matrix.version }}" @@ -58,7 +57,7 @@ jobs: uses: julia-actions/julia-processcoverage@v1 - name: "Report Coverage" - uses: codecov/codecov-action@v5 + uses: codecov/codecov-action@v7 if: ${{ github.event.pull_request.head.repo.full_name == github.repository }} with: files: lcov.info diff --git a/HISTORY.md b/HISTORY.md index a33d25516..a4ec2a767 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -2,6 +2,10 @@ ## 10.0 (Breaking) +This section describes the forthcoming 10.0 release. During integration, +`Project.toml` retains version 9.33.1 for development and testing; these breaking +changes are not intended for a 9.x release. + - **Breaking**: The `rng` keyword argument has been removed from `JumpProblem`. Pass `rng` to `solve` or `init` instead: ```julia @@ -13,14 +17,46 @@ jprob = JumpProblem(dprob, Direct(), jump) sol = solve(jprob, SSAStepper(); rng = Xoshiro(1234)) ``` - - RNG state is now owned by the integrator, not the aggregator. This - eliminates data races when sharing a `JumpProblem` across threads and - ensures a single, consistent RNG priority across all solver pathways: - `rng` > `seed` > `Random.default_rng()`. - - `rng` and `seed` kwargs are fully supported on `solve`/`init` for all - solver pathways (SSAStepper, ODE, SDE, tau-leaping). + - **Breaking**: `JumpProcesses.DEFAULT_RNG` has been removed. Custom jump + affects should draw from `SciMLBase.get_rng(integrator)`. Add SciMLBase as + a direct project dependency and load it with `using SciMLBase` to use this + interface. Pass `rng` or `seed` to `solve`/`init` to control the solver's RNG. + - Jump callbacks now sample from the integrator's RNG. `JumpProblem`s, + aggregators, and callbacks no longer store their own RNGs. Aggregators and + callbacks still contain mutable working state, so concurrent solves must + isolate that state; sharing an arbitrary `JumpProblem` across threads or + tasks is not made safe by this change. + - `SSAStepper` and the OrdinaryDiffEq ODE/DAE pathways accept `rng` and `seed` + through `solve`/`init`; all five CPU tau-leaping algorithms accept them + through `solve`. These pathways use `rng` first, then `Xoshiro(seed)` if a + seed is supplied, then `Random.default_rng()`. + - StochasticDiffEq SDE/RODE `solve` and `init` follow the backend's RNG policy. + An explicit RNG other than `TaskLocalRNG` takes priority over seeds. With + no RNG or with `TaskLocalRNG`, a nonzero solve/init seed takes priority, + followed by the underlying problem's nonzero seed, then a randomly seeded + `Xoshiro`. `TaskLocalRNG` is converted to `Xoshiro`; `seed = 0` means no seed + override on this pathway, whereas the JumpProcesses RNG resolver uses + `Xoshiro(0)`. + - **Breaking**: The singular `alias_jump` keyword is no longer consumed on + the StochasticDiffEq `init` pathway. For SDE jump-state copying, pass + `alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)` to `solve` or + `init`; use `alias_jumps = true` in the specifier to reuse jump state. RODE + problems use `SciMLBase.RODEAliasSpecifier`. Other JumpProcesses-owned + pathways retain their existing `alias_jump` behavior. - `SSAIntegrator` now supports the `SciMLBase` RNG interface (`has_rng`, `get_rng`, `set_rng!`). + - **Breaking**: Passing `seed` to a multi-trajectory `EnsembleGPUKernel` + solve with a GPU device backend now raises `ArgumentError`, instead of + accepting a seed that does not control the device RNG. This applies to + `SSAStepper`, `SimpleTauLeaping`, and `SimpleExplicitTauLeaping`. Seed the + device through the backend's API before solving (for example, + `CUDA.seed!`). Multi-trajectory kernel solves also reject host `rng` + objects and `rng_func`. The `CPU()` backend still supports `seed` using + Julia's task-local RNG; repeatability requires the same thread count, + workload/trajectory count, CPU backend settings, and Julia/package versions. + Identical trajectories across thread counts are not guaranteed. With `trajectories = 1`, + `EnsembleGPUKernel` retains its serial CPU fallback and that pathway's + usual RNG/seed rules. - **Breaking**: The `scale_rates` and `useiszero` keyword arguments have been removed from `JumpProblem`. Set them on the `MassActionJump` directly: ```julia @@ -32,18 +68,27 @@ jprob = JumpProblem(dprob, Direct(), maj) ``` - **Breaking**: Parameterized `MassActionJump`s (those constructed with - `param_idxs` or a custom `param_mapper`) are now immutable — rates are - computed from parameters at aggregator initialization rather than being - materialized into the jump at `JumpProblem` construction time. This means: + `param_idxs` or a custom `param_mapper`) retain `scaled_rates === nothing`. + Solvers and aggregators own the working rate buffers and fill them from + current parameters at initialization and reset. `remake` no longer mutates + shared rate coefficients in the jump definition. Fixed-rate and symbolic + definitions retain their supported `scaled_rates` representation. This means: - `update_parameters!` has been removed. Mass action rates are now automatically recomputed from the current parameter values whenever the aggregator reinitializes. After modifying parameters (e.g. in a callback), call `reset_aggregated_jumps!(integrator)` to trigger reinitialization with the updated parameter values. + - The `update_jump_params` keyword has been removed from + `reset_aggregated_jumps!`; supplying either `true` or `false` raises an + explanatory error. Rate refresh is automatic during reset. - Custom parameter mappers (e.g. ModelingToolkitBase's - `JumpSysMajParamMapper`) must implement the 3-arg callable API: - `(mapper)(dest::AbstractVector, maj::MassActionJump, params)`. - See [`MassActionJumpParamMapper`](@ref) for details. + `JumpSysMajParamMapper`) must replace the old extraction/mutation + callables with `(mapper)(dest::AbstractVector, maj::MassActionJump, params)` + and fill `dest` with the scaled rates. The built-in mapper honors + `maj.rescale_rates_on_update`; custom mappers own their scaling policy + and must avoid scaling already-scaled symbolic expressions twice. + See the [custom mapper migration guide](docs/src/migration.md#custom-parameter-mappers) + for an example. - Scalar `param_idxs` (e.g. `param_idxs = 1`) is now internally converted to a one-element vector. The scalar form continues to work as before. diff --git a/docs/pages.jl b/docs/pages.jl index 1f973f241..761708c58 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -10,5 +10,6 @@ pages = ["index.md", "Type Documentation" => Any["Jumps, JumpProblem, and Aggregators" => "jump_types.md", "Jump solvers" => "jump_solve.md"], "FAQ" => "faq.md", + "Migrating to 10.0" => "migration.md", "API" => "api.md" ] diff --git a/docs/src/api.md b/docs/src/api.md index 5ca6d1348..1ca9eeed7 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -126,7 +126,9 @@ from SciMLBase directly. ## Random Number Generator Control JumpProcesses supports controlling the random number generator (RNG) used for -jump sampling via the `rng` and `seed` keyword arguments to `solve` or `init`. +jump sampling via the `rng` and `seed` keyword arguments to `solve` or `init`, +as supported by the solver pathway below. `JumpProblem` no longer accepts `rng`; +see [Migrating to JumpProcesses 10](@ref migration_v10). ### `rng` keyword argument @@ -168,7 +170,8 @@ RNG other than `TaskLocalRNG` takes priority over seeds. With no RNG, or with underlying problem's nonzero `seed`; otherwise StochasticDiffEq creates a randomly seeded `Xoshiro`. It converts `TaskLocalRNG` to `Xoshiro` rather than storing it on the stochastic integrator. In this pathway, `seed=0` means no -seed override, whereas SSAStepper and the ODE/DAE pathway use `Xoshiro(0)`. +seed override, whereas SSAStepper, the ODE/DAE pathway, and CPU tau-leaping +solvers use `Xoshiro(0)`. ### Behavior by solver pathway @@ -181,8 +184,24 @@ seed override, whereas SSAStepper and the ODE/DAE pathway use `Xoshiro(0)`. | `SimpleExplicitTauLeaping`, `SimpleImplicitTauLeaping`, `SimpleTrapezoidalLeaping`, `SimpleAdaptiveTauLeaping` | `Random.default_rng()` | Full support via `solve` kwargs | !!! note - For reproducible simulations, always pass an explicit `rng` or `seed`. + Use an explicit `rng` or `seed` where the solver supports it, and start from + identical model and mutable aggregator state. An explicit RNG is advanced + by the solve; construct or copy it from the same starting state for replay. Julia's default RNG is task-local; results may depend on prior draws in the task. + `SortingDirect` retains its learned search order, which can also affect replay. + +### Kernel ensembles + +Multi-trajectory [`EnsembleGPUKernel`](@ref) solves use backend-local RNGs and +reject host `rng` and `rng_func` inputs. With `KernelAbstractions.CPU()`, `seed` +seeds Julia's task-local RNG. Replay assumes the same thread count, workload +and trajectory count, CPU backend settings, and Julia/package versions; +identical trajectories across thread counts are not guaranteed. This differs +from CPU ensemble methods that assign RNGs explicitly to each trajectory. + +GPU device backends reject `seed` with `ArgumentError`; seed the backend +before solving, for example with `CUDA.seed!`. With `trajectories = 1`, the +serial CPU fallback accepts its usual RNG/seed inputs. # Private / Developer API @@ -225,6 +244,13 @@ backend aliases on thread 1 and copies on other threads. Both `solve` and `init` follow this policy; `alias_jumps` is a field of the specifier, not a standalone keyword argument. +The singular `alias_jump` keyword is no longer consumed on the SDE `init` +path; use the backend alias specifier above. JumpProcesses-owned SSAStepper +and ODE/DAE initialization retain `alias_jump`, defaulting to `true` on +thread 1 and `false` on other threads. Independent mutable jump state is +still required for concurrent solves, including when using +`EnsembleProblem(...; safetycopy = false)`. + For **tau-leaping**, JumpProcesses defines a custom `DiffEqBase.solve` that bypasses the standard `__solve`/`__init` pathway. It calls `resolve_rng` directly with the `rng` and `seed` kwargs from the `solve` call. diff --git a/docs/src/faq.md b/docs/src/faq.md index 47d77c3cb..07ae35307 100644 --- a/docs/src/faq.md +++ b/docs/src/faq.md @@ -75,7 +75,11 @@ sol = solve(jprob, SSAStepper(); rng = StableRNG(1234)) A `seed` keyword argument is also supported as a shorthand for creating a `Xoshiro` generator: `solve(jprob, SSAStepper(); seed = 1234)`. -By default, JumpProcesses uses Julia's built-in `Random.default_rng()`. +By default, `SSAStepper`, the OrdinaryDiffEq ODE/DAE pathway, and all five CPU +tau-leaping solvers use Julia's built-in `Random.default_rng()`. +SDE/RODE solvers follow StochasticDiffEq's RNG/seed policy, and GPU kernel +ensembles have backend-specific controls; see +[Random Number Generator Control](@ref). ## What are these aggregators and aggregations in JumpProcesses? @@ -133,6 +137,11 @@ When using an ODE or SDE time stepper, any callback should work. underlying jump simulation algorithms know to reinitialize their internal data structures. Omitting this call will lead to incorrect behavior! +The reset refreshes parameter-mapped mass-action working rates from the current +parameters. The former `update_jump_params` keyword is no longer accepted; no +flag is needed. Fixed coefficients remain fixed. See the +[10.0 migration guide](@ref migration_v10) for mapper and constructor changes. + A simple example that uses a `MassActionJump` and changes the parameters at a specified time in the simulation using a `DiscreteCallback` is diff --git a/docs/src/jump_types.md b/docs/src/jump_types.md index 181964d91..cf6f02941 100644 --- a/docs/src/jump_types.md +++ b/docs/src/jump_types.md @@ -112,7 +112,17 @@ MassActionJump(reactant_stoich, net_stoich; scale_rates = true, param_idxs = not Gillespie, and hence need to be rescaled. *The default, `scale_rates = true`, corresponds to rescaling the passed in rate constants.* See below. - `param_idxs` is a vector of the indices within the parameter vector, `p`, that - correspond to the rate constant for each jump. + correspond to the rate constant for each jump, or a scalar index for a + single reaction. Alternatively, supply a custom `param_mapper` callable + `(dest, maj, params)` that fills the final scaled coefficients in `dest`. + +Parameterized definitions have `scaled_rates === nothing`. Working rates are +owned by the solver or aggregator and refreshed from the current parameters +at initialization and jump-state resets. After changing parameters in a +callback, call [`reset_aggregated_jumps!`](@ref). Definitions constructed with +explicit fixed coefficients retain their stored rates and do not read `p`. +See [Migrating to JumpProcesses 10](@ref migration_v10) for custom mapper +scaling and the removed `JumpProblem` constructor keywords. **Notes for Mass Action Jumps** @@ -418,8 +428,12 @@ condition or simulation parameters without having to create and initialize a new `JumpProblem`. In such situations `remake` can be used to change the initial condition, time span, and the parameter vector. **Note,** the new `JumpProblem` will alias internal data structures from the old problem, including core -components of the SSA aggregators. As such, only the new problem generated by -`remake` should be used for subsequent simulations. +components of the SSA aggregators. Sequential solves refresh parameter-mapped +mass-action working rates from each problem's current parameters; `remake` +does not overwrite rates in a shared parameter-mapped definition. For concurrent +solves, create independent problem copies. `SortingDirect` also retains its +learned search order, so a reused problem can have different initial sampling +state even with the same seed. As an example, consider the following SIR model: diff --git a/docs/src/migration.md b/docs/src/migration.md new file mode 100644 index 000000000..c3824a0a0 --- /dev/null +++ b/docs/src/migration.md @@ -0,0 +1,164 @@ +# [Migrating to JumpProcesses 10](@id migration_v10) + +This guide describes the forthcoming 10.0 API. During integration, the package +still carries 9.x development version metadata; that does not make the changes +below compatible with released 9.x code. + +## Move RNG inputs to the solver + +`JumpProblem` no longer accepts `rng`. Pass it to `solve` or `init` instead: + +```@example migration_rng +using JumpProcesses, Random + +maj = MassActionJump([Pair{Int, Int}[]], [[1 => 1]]; param_idxs = [1]) +dprob = DiscreteProblem([0], (0.0, 1.0), [2.0]) +jprob = JumpProblem(dprob, Direct(), maj) +sol = solve(jprob, SSAStepper(); rng = Xoshiro(1234)) +nothing # hide +``` + +For `SSAStepper`, ODE/DAE solvers, and all five CPU tau-leaping algorithms, +an explicit `rng` is used and advanced without reseeding it. Otherwise, `seed` +creates a fresh `Xoshiro(seed)`; with neither input, the solver uses +`Random.default_rng()`. The CPU tau-leaping methods accept these inputs through +`solve`, not `init`. + +StochasticDiffEq uses its own SDE/RODE policy for both `solve` and `init`. An +explicit RNG other than `TaskLocalRNG` takes precedence. With no RNG or a +`TaskLocalRNG`, it uses a nonzero solve/init seed, then the underlying problem's +nonzero seed, then a random seed. It stores a `Xoshiro`, not the `TaskLocalRNG`, +in that case. A stochastic seed of zero means no seed override. See +[Random Number Generator Control](@ref) for the pathway table. + +Jump `affect!` functions that draw random numbers should use +`SciMLBase.get_rng(integrator)` to draw from the solver's RNG. Add SciMLBase +as a direct dependency of your project and load it with `using SciMLBase` +before using this interface; JumpProcesses does not re-export `get_rng`. +The old `JumpProcesses.DEFAULT_RNG` constant is removed. + +## Keep mass-action definitions separate from working rates + +In 9.x, parameter updates could overwrite a `MassActionJump`'s stored scaled +rates. In 10.0, a parameter-mapped definition has `scaled_rates === nothing`. +The solver or aggregator owns a numeric working-rate buffer, filled from the +current parameters at initialization and when the jump state is reset. CPU +tau-leaping solvers fill their buffer at each solve; kernels materialize host +rates before transferring them to the backend. + +Use `param_idxs` or a custom `param_mapper` for rates that depend on parameters. +For example, the definition above can be reused with different parameters: + +```@example migration_rng +updated = remake(jprob; p = [4.0]) +updated_sol = solve(updated, SSAStepper(); seed = 1234) +nothing # hide +``` + +Definitions constructed with explicit fixed coefficients retain their stored +rates, which are copied into working buffers. Changing `p` does not turn fixed +coefficients into parameter-dependent rates. Definition arrays can still be +mutable; do not mutate shared rate or stoichiometry arrays during solves. + +### Constructor and callback keyword changes + +Set `scale_rates` and `useiszero` on `MassActionJump`, not `JumpProblem`. +For example, pre-scaled coefficients can be supplied with +`MassActionJump(rates, reactant_stoich, net_stoich; scale_rates = false)`. +Passing either keyword to `JumpProblem` now raises `ArgumentError`. + +The old `update_parameters!` mutator is removed, along with the +`update_jump_params` keyword of `reset_aggregated_jumps!`. +After changing parameters or the state in a callback, call +`reset_aggregated_jumps!(integrator)` without that keyword. It refreshes +parameter-mapped working rates as part of reinitializing jump state; parameter +changes during a solve are not detected automatically without a reset. + +```@example migration_rng +function change_rate!(integrator) + integrator.p[1] = 4.0 + reset_aggregated_jumps!(integrator) + nothing +end +condition(u, t, integrator) = t == 0.5 +callback = DiscreteCallback(condition, change_rate!) +callback_sol = solve(remake(jprob; p = [2.0]), SSAStepper(); + seed = 1234, callback, tstops = [0.5]) +nothing # hide +``` + +### Custom parameter mappers + +Replace mappers that mutate `maj.scaled_rates` or return a new rate vector with +the in-place callable `(dest, maj, params)`. Fill every entry of `dest`, leave +the definition and input parameters unchanged, and return `nothing`. +`fill_scaled_rates!` delegates to this callable; it does not apply scaling +afterward. The mapper must produce the final scaled coefficients. + +For example, this mapper uses a product of two parameters as the rate constant +for a reaction consuming two molecules. It includes the stoichiometric factor +`1 / 2!` directly, so the definition uses `scale_rates = false`: + +```@example migration_mapper +using JumpProcesses + +function map_rates!(dest, maj, params) + dest[1] = params[1] * params[2] / 2 + nothing +end + +maj = MassActionJump([[1 => 2]], [[1 => -2]]; + param_mapper = map_rates!, scale_rates = false) +dprob = DiscreteProblem([20], (0.0, 1.0), [0.2, 0.5]) +jprob = JumpProblem(dprob, Direct(), maj) +sol = solve(jprob, SSAStepper(); seed = 1234) +nothing # hide +``` + +The built-in `param_idxs` mapper applies stoichiometric scaling according to +`rescale_rates_on_update`, which defaults to `scale_rates`. A custom mapper +that already produces scaled coefficients should use `scale_rates = false` +and avoid scaling them again. Built-in parameter-index mappers can be merged; +merging arbitrary custom mappers is not supported. Combine those reactions in +one definition with one mapper instead. + +## Kernel ensemble RNG controls + +For multiple trajectories, `EnsembleGPUKernel` uses its backend's RNG and +rejects host `rng` and `rng_func` inputs. `seed` is supported with +`KernelAbstractions.CPU()`. Replay assumes the same thread count, workload and +trajectory count, CPU backend settings, and Julia/package versions. Identical +trajectories across thread counts are not guaranteed. + +With a GPU device backend, a non-`nothing` `seed` now raises `ArgumentError`; +previously it was accepted without controlling the device RNG. Seed through +the backend API before `solve`, for example `CUDA.seed!`. These rules apply to +the `SSAStepper`, `SimpleTauLeaping`, and `SimpleExplicitTauLeaping` kernels. +With `trajectories = 1`, execution uses the serial CPU fallback and its usual +RNG/seed controls. + +## Copying and concurrent solves + +Integrator RNGs and separate working-rate buffers do not make a shared +`JumpProblem` safe for concurrent solves: callbacks and aggregators remain +mutable. Use independent problem copies for manually launched tasks. With +`EnsembleProblem`, `safetycopy = true` requests a separate problem copy per +trajectory. If using `safetycopy = false`, ensure your `prob_func` and solver +settings preserve the required isolation; do not assume arbitrary shared +mutable objects are safe. +Do not return one shared mutable RNG from a custom ensemble `rng_func`. + +The existing `alias_jump` keyword remains available for JumpProcesses-owned +SSAStepper and ODE/DAE initialization. Its default is `true` on thread 1 and +`false` on other threads. The SDE/RODE backend instead uses +`alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)` or +`SciMLBase.RODEAliasSpecifier(; alias_jumps = false)` to request copied jump +state. The singular `alias_jump` keyword is no longer consumed on the SDE +`init` path; use the backend specifier. `alias_jumps` is a field of that +specifier, not a standalone keyword. A common JumpProcesses alias API is +not introduced by this migration. + +`SortingDirect` retains its learned search order across solves and resets. +For identical seeded trajectories, start from identical learned state as well +as identical model inputs; reusing a problem with a different learned order +can change the trajectory. diff --git a/docs/src/tutorials/discrete_stochastic_example.md b/docs/src/tutorials/discrete_stochastic_example.md index 99db5078d..fb2389ad2 100644 --- a/docs/src/tutorials/discrete_stochastic_example.md +++ b/docs/src/tutorials/discrete_stochastic_example.md @@ -477,6 +477,11 @@ function that gets evaluated is with ``\hat{k} = k / \prod_{i=1}^{N} R_i!`` the renormalized rate constant. Passing the keyword argument `scale_rates = false` will disable `MassActionJump`s internally rescaling the rate constant by ``(\prod_{i=1}^{N} R_i!)^{-1}``. +Set this keyword on `MassActionJump`, not `JumpProblem`. For the `param_idxs` +form used here, scaling is applied when the solver fills its working rates +from the current parameters at initialization or a jump-state reset. If a +callback changes those parameters, call `reset_aggregated_jumps!(integrator)` +after the change; no `update_jump_params` keyword is needed. For chemical reaction systems Catalyst.jl automatically groups reactions into their optimal jump representation. diff --git a/ext/explicit_tau_leaping.jl b/ext/explicit_tau_leaping.jl index 7af58e362..72dffdc63 100644 --- a/ext/explicit_tau_leaping.jl +++ b/ext/explicit_tau_leaping.jl @@ -214,11 +214,18 @@ time grid instead. Each returned solution therefore contains exactly the grid points, as if the problem had been built with `save_positions = (false, false)`. For multiple trajectories, randomness comes from the backend's own RNG. `seed` -is supported with `CPU()` and seeds Julia's task-local generator; device backends -must be seeded through their own API before calling `solve` (for example, -`CUDA.seed!`). Passing a host `rng` or `rng_func`, or `seed` with a device backend, -is rejected. -A single trajectory uses `EnsembleSerial()` and its usual solve-level RNG inputs. +is supported with `CPU()` and seeds Julia's task-local generator. CPU-kernel replay +with the same seed requires the same thread count, workload (including trajectory +count), CPU backend settings, and Julia and package versions. Identical trajectories +across thread counts are not guaranteed. These conditions apply to CPU kernel +execution; CPU ensemble algorithms that explicitly assign RNGs to trajectories +have their own RNG behavior. + +Device backends must be seeded through their own API before calling `solve` +(for example, `CUDA.seed!`). Passing a host `rng` or `rng_func`, or `seed` with a +device backend, raises `ArgumentError`. With `trajectories = 1`, execution delegates +to `EnsembleSerial()` on the CPU and accepts its usual `rng`, `seed`, and ensemble +`rng_func` keywords. The reaction data is uploaded to the device once and shared by every thread, so all trajectories solve the same problem and a `prob_func` is not supported. diff --git a/ext/ssa_stepper.jl b/ext/ssa_stepper.jl index f666da319..e32a0127d 100644 --- a/ext/ssa_stepper.jl +++ b/ext/ssa_stepper.jl @@ -269,11 +269,18 @@ grid points, as if the problem had been built with also ignored, since the kernel always runs the Direct method. For multiple trajectories, randomness comes from the backend's own RNG. `seed` -is supported with `CPU()` and seeds Julia's task-local generator; device backends -must be seeded through their own API before calling `solve` (for example, -`CUDA.seed!`). Passing a host `rng` or `rng_func`, or `seed` with a device backend, -is rejected. -A single trajectory uses `EnsembleSerial()` and its usual solve-level RNG inputs. +is supported with `CPU()` and seeds Julia's task-local generator. CPU-kernel replay +with the same seed requires the same thread count, workload (including trajectory +count), CPU backend settings, and Julia and package versions. Identical trajectories +across thread counts are not guaranteed. These conditions apply to CPU kernel +execution; CPU ensemble algorithms that explicitly assign RNGs to trajectories +have their own RNG behavior. + +Device backends must be seeded through their own API before calling `solve` +(for example, `CUDA.seed!`). Passing a host `rng` or `rng_func`, or `seed` with a +device backend, raises `ArgumentError`. With `trajectories = 1`, execution delegates +to `EnsembleSerial()` on the CPU and accepts its usual `rng`, `seed`, and ensemble +`rng_func` keywords. The reaction data is uploaded to the device once and shared by every thread, so all trajectories solve the same problem and a `prob_func` is not supported. diff --git a/src/jumps.jl b/src/jumps.jl index b627ae033..2f847e59e 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -469,7 +469,14 @@ $(FIELDS) `reactant_stoch` from the input. Note, if `scale_rates=true` this will potentially modify both of these. - `param_idxs = nothing`, indexes in the parameter vector, `JumpProblem.prob.p`, - that correspond to each reaction's rate. + that correspond to each reaction's rate; a scalar index is accepted for one reaction. +- `param_mapper = nothing`, an alternative to `param_idxs`: a callable + `(dest, maj, params)` that fills the final scaled working rates in `dest` and + returns `nothing`. It must not mutate the definition or input parameters. +- `rescale_rates_on_update = scale_rates`, the scaling policy used by the built-in + parameter mapper. Custom mappers implement their own scaling and should honor + this policy when mapping unscaled coefficients; pre-scaled output must not be + scaled again. See the tutorial and main docs for details. @@ -495,8 +502,9 @@ jprob = JumpProblem(prob, Direct(), maj) ``` ## Notes -- By default, reaction rates are rescaled when constructing the `MassActionJump` - as explained in the [main +- By default, fixed reaction rates are rescaled when constructing the `MassActionJump`; + the built-in parameter mapper rescales rates when filling each solver's working + buffer at initialization/reset, as explained in the [main docs](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Defining-a-Mass-Action-Jump). Disable this with the kwarg `scale_rates=false`. - Also see the [main @@ -507,11 +515,11 @@ jprob = JumpProblem(prob, Direct(), maj) """ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump """ - The (scaled) reaction rate constants. When stored within a `JumpProblem`, this vector is - shared across remade problems. Users are responsible for maintaining consistent - stoichiometric scaling and thread safety if mutating directly. In general, prefer using - the parameter index or mapping form for rates that need to change, rather than mutating - `scaled_rates` directly. + The stored (scaled) fixed reaction rate constants, or `nothing` for a + parameter-mapped definition. Numeric working rates live in solver/aggregator + buffers and are refreshed at initialization/reset. Stored definition arrays + can be shared across remade problems; do not mutate them during concurrent + solves. Prefer parameter indices or a custom mapper for rates that change. """ scaled_rates::T """The reactant stoichiometry vectors.""" @@ -657,12 +665,15 @@ Implements the in-place mapper callable API: which is called by [`fill_scaled_rates!`](@ref) during aggregator initialization and reinitialization to populate the working rate vector from parameters. -`dest` should be filled with the current rates for each reaction. If -`maj.rescale_rates_on_update` is `true`, the mapper should also apply -stoichiometric scaling via [`scalerates!`](@ref). +`dest` must contain the final scaled coefficients for each reaction. The +built-in mapper extracts parameter values and applies stoichiometric scaling +via `scalerates!` when `maj.rescale_rates_on_update` is `true`. Custom mappers (e.g. ModelingToolkitBase's `JumpSysMajParamMapper`) should implement this 3-arg callable to support parameterized `MassActionJump`s. +They own their scaling policy: `fill_scaled_rates!` does not scale their output +again. A mapper producing pre-scaled coefficients should use `scale_rates = false` +and avoid applying the stoichiometric factors a second time. """ struct MassActionJumpParamMapper{U} param_idxs::U diff --git a/src/problem.jl b/src/problem.jl index 29e23b6d2..98861d59a 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -28,7 +28,7 @@ transition rates are constant *between* jumps (called [`ConstantRateJump`](@ref) [`DiscreteProblem`](https://docs.sciml.ai/DiffEqDocs/stable/types/discrete_types/) ```julia -prob = DiscreteProblem(u0, p, tspan) +prob = DiscreteProblem(u0, tspan, p) ``` where `u0` is the initial condition, `p` the parameters and `tspan` the time span. If we @@ -75,10 +75,12 @@ examples and commonly asked questions. `JumpProblem` contains mutable state (aggregator data, callbacks) and is **not thread-safe**. A single `JumpProblem` instance must not be solved concurrently from multiple threads or tasks without first creating independent copies via `deepcopy`. - When running ensemble simulations via `EnsembleProblem`, this is handled automatically - — the `SciMLBase` ensemble layer provides per-task isolation and per-trajectory RNG - seeding. This warning only applies to manually parallelized `solve` calls outside the - ensemble interface. + With `EnsembleProblem`, `safetycopy = true` requests a separate problem copy + per trajectory. If using `safetycopy = false`, the caller's + `prob_func` and solver settings must provide the required isolation. CPU ensemble + methods can assign separate RNGs to trajectories; a custom `rng_func` must not + return one shared mutable RNG. Separate RNGs alone do not isolate mutable jump + state, and kernel ensembles have their own backend RNG rules. """ mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggregator}, J1, J2, J3, J4, K} <: SciMLBase.AbstractJumpProblem{P, J} diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index a235f6d32..5a5a9889a 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -1123,12 +1123,20 @@ Ensemble algorithm marker for GPU execution of tau-leaping ensemble simulations. ## Random number generation -Kernel ensembles use backend-local RNGs and reject explicit host `rng` and -`rng_func` keywords. -The `seed` solve keyword is supported with the KernelAbstractions CPU backend; -other backends require their own seeding API and reject `seed`. With -`trajectories = 1`, execution uses the serial CPU solver, which accepts its -usual `rng`, `seed`, and ensemble `rng_func` keywords. +For multiple trajectories, kernel ensembles use backend-local RNGs and reject +explicit host `rng` and `rng_func` keywords. The `seed` solve keyword is supported +with the KernelAbstractions CPU backend and seeds Julia's task-local generator. +CPU-kernel replay with the same seed requires the same thread count, workload +(including trajectory count), CPU backend settings, and Julia and package +versions. Identical trajectories across thread counts are not guaranteed. These +conditions apply to CPU kernel execution; CPU ensemble algorithms that explicitly +assign RNGs to trajectories have their own RNG behavior. + +Device backends reject `seed` with `ArgumentError`. Seed the device RNG through +the backend's API, such as `CUDA.seed!`, before calling `solve`. + +With `trajectories = 1`, execution delegates to `EnsembleSerial()` on the CPU +and accepts its usual `rng`, `seed`, and ensemble `rng_func` keywords. ## Returns diff --git a/test/ensemble_uniqueness.jl b/test/ensemble_uniqueness.jl index 9878fc8b7..0390e8bca 100644 --- a/test/ensemble_uniqueness.jl +++ b/test/ensemble_uniqueness.jl @@ -11,16 +11,16 @@ dprob = DiscreteProblem(u0, (0.0, 100.0)) jump_prob = JumpProblem(dprob, Direct(), j1, j2) ensemble_prob = EnsembleProblem(jump_prob) sol = solve(ensemble_prob, FunctionMap(), trajectories = 3; rng = StableRNG(12345)) -@test Array(sol.u[1]) !== Array(sol.u[2]) -@test Array(sol.u[1]) !== Array(sol.u[3]) -@test Array(sol.u[2]) !== Array(sol.u[3]) +@test Array(sol.u[1]) != Array(sol.u[2]) +@test Array(sol.u[1]) != Array(sol.u[3]) +@test Array(sol.u[2]) != Array(sol.u[3]) @test eltype(sol.u[1].u[1]) == Int # Test with SSAStepper - pass rng to solve so trajectories get unique sequences jump_prob = JumpProblem(dprob, Direct(), j1, j2) ensemble_prob2 = EnsembleProblem(jump_prob) sol = solve(ensemble_prob2, SSAStepper(), trajectories = 3; rng = StableRNG(12345)) -@test Array(sol.u[1]) !== Array(sol.u[2]) -@test Array(sol.u[1]) !== Array(sol.u[3]) -@test Array(sol.u[2]) !== Array(sol.u[3]) +@test Array(sol.u[1]) != Array(sol.u[2]) +@test Array(sol.u[1]) != Array(sol.u[3]) +@test Array(sol.u[2]) != Array(sol.u[3]) @test eltype(sol.u[1].u[1]) == Int diff --git a/test/thread_safety.jl b/test/thread_safety.jl index 5ca0f42d9..007025512 100644 --- a/test/thread_safety.jl +++ b/test/thread_safety.jl @@ -1,5 +1,6 @@ using DiffEqBase, SciMLBase, Test using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq +using StableRNGs sr = [1.0, 2.0, 50.0] maj = MassActionJump(sr, [[1 => 1], [1 => 1], [0 => 1]], [[1 => 1], [1 => -1], [1 => 1]]) @@ -9,24 +10,50 @@ u0 = [5] dprob = DiscreteProblem(u0, tspan, params) jprob = JumpProblem(dprob, Direct(), maj) -# Verify threaded solves complete and produce distinct trajectories. -# NOTE: We intentionally do NOT pass `rng` here. In threaded ensembles, passing a -# shared rng object via `solve(...; rng=...)` does not yet provide correct -# per-trajectory stream handling. Until SciMLBase's ensemble RNG updates land -# (master rng -> per-trajectory rng), correctness in threaded contexts relies on -# task-local `Random.default_rng()`. +# Keep coverage of default task-local RNGs as well as the explicit ensemble RNG +# controls below. SciMLBase uses an ensemble `rng` or `seed` to assign streams to +# individual trajectories; the master RNG is not shared by the jump integrators. sol = solve(EnsembleProblem(jprob), SSAStepper(), EnsembleThreads(); trajectories = 400) -@test length(sol) == 400 -firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol)] +@test length(sol.u) == 400 +firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol.u)] @test allunique(firstrx_time) sol2 = solve(EnsembleProblem(jprob; safetycopy = true), SSAStepper(), EnsembleThreads(); trajectories = 400) -@test length(sol2) == 400 -firstrx_time2 = [sol2.u[i].t[findfirst(>(sol2.u[i].t[1]), sol2.u[i].t)] for i in 1:length(sol2)] +@test length(sol2.u) == 400 +firstrx_time2 = [sol2.u[i].t[findfirst(>(sol2.u[i].t[1]), sol2.u[i].t)] for i in 1:length(sol2.u)] @test allunique(firstrx_time2) +@testset "Threaded ensemble RNG and seed controls" begin + paths(sol) = [(trajectory.t, trajectory.u) for trajectory in sol.u] + first_event_times(sol) = [trajectory.t[findfirst(>(first(trajectory.t)), trajectory.t)] + for trajectory in sol.u] + for safetycopy in (false, true), input in (:rng, :seed) + @testset "safetycopy=$safetycopy, $input" begin + # Construct a fresh master RNG for each ensemble solve. Only the + # ensemble layer uses this object; each trajectory receives its own RNG. + rng_options(seed) = input === :rng ? (; rng = StableRNG(seed)) : (; seed) + ensemble = EnsembleProblem(jprob; safetycopy) + serial = solve(ensemble, SSAStepper(), EnsembleSerial(); + trajectories = 16, rng_options(12345)...) + threaded = solve(ensemble, SSAStepper(), EnsembleThreads(); + trajectories = 16, rng_options(12345)...) + replay = solve(ensemble, SSAStepper(), EnsembleThreads(); + trajectories = 16, rng_options(12345)...) + changed = solve(ensemble, SSAStepper(), EnsembleThreads(); + trajectories = 16, rng_options(54321)...) + + @test length(threaded.u) == 16 + @test all(SciMLBase.successful_retcode, threaded.u) + @test paths(threaded) == paths(serial) + @test paths(threaded) == paths(replay) + @test first_event_times(threaded) != first_event_times(changed) + @test allunique(first_event_times(threaded)) + end + end +end + # test for https://github.com/SciML/JumpProcesses.jl/issues/472 let function f!(du, u, p, t) @@ -60,7 +87,7 @@ let prob = EnsembleProblem(jump_prob) sol = solve(prob, SRIW1(), EnsembleThreads(); trajectories = 400, save_everystep = false) - firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol)] + firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol.u)] @test allunique(firstrx_time) end end From 3b16f4a6c5582f0504c08bc64a9f0ec8292c9bbd Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 29 Sep 2026 14:53:10 -0400 Subject: [PATCH 34/52] Fix minimum dependency and downgrade test coverage --- .github/workflows/Downgrade.yml | 2 ++ Project.toml | 2 +- 2 files changed, 3 insertions(+), 1 deletion(-) diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml index 26c77b953..c296fee7d 100644 --- a/.github/workflows/Downgrade.yml +++ b/.github/workflows/Downgrade.yml @@ -16,4 +16,6 @@ jobs: # Runs strict (allow_reresolve=false); expected RED until the upstream fix # ForwardDiff/LogExpFunctions — red until ForwardDiff 1.3.4 (#811). Auto-greens on upstream release. uses: "SciML/.github/.github/workflows/downgrade.yml@v1" + with: + group: All secrets: "inherit" diff --git a/Project.toml b/Project.toml index 9c002e4fc..1762b57b1 100644 --- a/Project.toml +++ b/Project.toml @@ -52,7 +52,7 @@ OrdinaryDiffEq = "7.0.0, 7" OrdinaryDiffEqCore = "4.3.0, 4" OrdinaryDiffEqFunctionMap = "2.0.0, 2" Pkg = "1" -PoissonRandom = "0.4" +PoissonRandom = "0.4.8" PrecompileTools = "1.2.1" Random = "1" RecursiveArrayTools = "4.3.0, 4" From feaf85ee7bac30872f4975b60f9ad439cf6f4e5f Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 29 Sep 2026 15:38:37 -0400 Subject: [PATCH 35/52] Use latest Julia CI with the pinned stochastic initialization fix --- .github/workflows/Documentation.yml | 8 +++----- .github/workflows/Downgrade.yml | 5 +++-- .github/workflows/Tests.yml | 16 +++++++++++++++- .github/workflows/ThreadSafety.yml | 3 +-- .github/workflows/benchmark.yml | 3 ++- Project.toml | 8 +++++++- 6 files changed, 31 insertions(+), 12 deletions(-) diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml index fb51f43a6..6dec07596 100644 --- a/.github/workflows/Documentation.yml +++ b/.github/workflows/Documentation.yml @@ -1,11 +1,9 @@ name: Documentation on: - push: - branches: - - master - tags: '*' - pull_request: + # Restore push (master and tags) and pull_request after ModelingToolkitBase + # supports the v10 JumpProblem and mapper APIs used by the Catalyst tutorial. + workflow_dispatch: jobs: build: diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml index c296fee7d..8a630a96a 100644 --- a/.github/workflows/Downgrade.yml +++ b/.github/workflows/Downgrade.yml @@ -13,8 +13,9 @@ on: jobs: downgrade: name: "Downgrade" - # Runs strict (allow_reresolve=false); expected RED until the upstream fix - # ForwardDiff/LogExpFunctions — red until ForwardDiff 1.3.4 (#811). Auto-greens on upstream release. + # The downgrade action discards URL sources. Restore this job after the + # fixed StochasticDiffEqCore is released and the temporary pin is removed. + if: ${{ false }} uses: "SciML/.github/.github/workflows/downgrade.yml@v1" with: group: All diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index 45dc20491..6ce571b22 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -18,5 +18,19 @@ concurrency: jobs: tests: - uses: "SciML/.github/.github/workflows/grouped-tests.yml@v1" + # Restore grouped-tests.yml and test/test_groups.toml's full matrix after + # OrdinaryDiffEq.jl#4614 is released and the temporary Core source is removed. + strategy: + fail-fast: false + matrix: + group: ["QA", "InterfaceI", "InterfaceII"] + name: "${{ matrix.group }} (latest Julia)" + uses: "SciML/.github/.github/workflows/tests.yml@v1" + with: + project: "." + group: "${{ matrix.group }}" + julia-version: "1" + num-threads: "1" + coverage-directories: "src,ext" + timeout-minutes: 120 secrets: "inherit" diff --git a/.github/workflows/ThreadSafety.yml b/.github/workflows/ThreadSafety.yml index 36c7220e2..e79efb8f9 100644 --- a/.github/workflows/ThreadSafety.yml +++ b/.github/workflows/ThreadSafety.yml @@ -22,10 +22,9 @@ jobs: strategy: fail-fast: false matrix: + # Restore lts/pre after the temporary Core source pin is removed. version: - "1" - - "lts" - - "pre" runs-on: ubuntu-latest steps: - uses: actions/checkout@v7 diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 4f5b471ec..6e349481b 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -16,7 +16,8 @@ jobs: strategy: fail-fast: false matrix: - version: ["1", "lts"] + # Restore lts when the full version matrix is enabled again. + version: ["1"] env: # Force consistent Julia depot path for self-hosted runners JULIA_DEPOT_PATH: ~/.julia diff --git a/Project.toml b/Project.toml index 1762b57b1..d481a336c 100644 --- a/Project.toml +++ b/Project.toml @@ -64,10 +64,15 @@ StableRNGs = "1" StaticArrays = "1.9.18" Statistics = "1" StochasticDiffEq = "7.0.0, 7" +StochasticDiffEqCore = "2.2.4" SymbolicIndexingInterface = "0.3.48" Test = "1" julia = "1.10" +[sources] +# Temporary until OrdinaryDiffEq.jl#4614 is merged and released. +StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "77c07bffe2a225a95436190570eb8795f69342d0", subdir = "lib/StochasticDiffEqCore"} + [extras] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" @@ -83,7 +88,8 @@ SciMLTesting = "09d9d899-5365-40a9-917a-5f67fddea283" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" +StochasticDiffEqCore = "19c5a474-6cd1-4a5f-be79-46dc34e54d7f" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Adapt", "FastBroadcast", "ForwardDiff", "KernelAbstractions", "LinearAlgebra", "LinearSolve", "OrdinaryDiffEq", "OrdinaryDiffEqFunctionMap", "Pkg", "SafeTestsets", "SciMLTesting", "StableRNGs", "Statistics", "StochasticDiffEq", "Test"] +test = ["Adapt", "FastBroadcast", "ForwardDiff", "KernelAbstractions", "LinearAlgebra", "LinearSolve", "OrdinaryDiffEq", "OrdinaryDiffEqFunctionMap", "Pkg", "SafeTestsets", "SciMLTesting", "StableRNGs", "Statistics", "StochasticDiffEq", "StochasticDiffEqCore", "Test"] From 93a96523c340d181d14166d12344d9d812541949 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Tue, 29 Sep 2026 15:38:55 -0400 Subject: [PATCH 36/52] Apply repository formatting and clarify test helper names --- HISTORY.md | 16 +- LICENSE.md | 6 +- README.md | 6 +- benchmark/benchmarks.jl | 10 +- docs/make.jl | 6 +- docs/src/api.md | 39 +- .../applications/advanced_point_process.md | 3 + docs/src/index.md | 3 +- docs/src/jump_types.md | 25 +- .../tutorials/discrete_stochastic_example.md | 3 +- docs/src/tutorials/jump_diffusion.md | 2 +- ext/JumpProcessesForwardDiffExt.jl | 2 +- ext/JumpProcessesKernelAbstractionsExt.jl | 2 +- ext/JumpProcessesOrdinaryDiffEqCoreExt.jl | 2 +- ext/ssa_stepper.jl | 28 +- src/JumpProcesses.jl | 2 +- src/SSA_stepper.jl | 77 +++- src/aggregators/prioritytable.jl | 40 +- src/aggregators/rssa.jl | 6 +- src/aggregators/rssacr.jl | 6 +- src/aggregators/ssajump.jl | 1 + src/coupling.jl | 4 +- src/extended_jump_array.jl | 19 +- src/jumps.jl | 378 +++++++++++------- src/problem.jl | 52 ++- src/simple_regular_solve.jl | 74 ++-- src/solve.jl | 10 +- src/spatial/flatten.jl | 3 +- src/spatial/hop_rates.jl | 130 ++++-- src/spatial/reaction_rates.jl | 22 +- src/spatial/topology.jl | 12 +- src/spatial/utils.jl | 12 +- test/constant_rate.jl | 3 +- test/degenerate_rx_cases.jl | 6 +- test/ensemble_problems.jl | 7 +- test/explicit_tau_kernel_tests.jl | 1 + test/hawkes_test.jl | 3 +- test/kernel_contract_tests.jl | 6 +- test/massaction_state.jl | 7 +- test/qa.jl | 50 +-- test/regular_jumps.jl | 86 ++-- test/rng_kwarg_tests.jl | 4 +- test/runtests.jl | 2 +- test/saveat_regression.jl | 16 +- test/sortingdirect_reuse.jl | 3 +- test/spatial/hop_rates.jl | 12 +- test/spatial/reaction_rates.jl | 4 +- test/splitcoupled.jl | 3 +- test/ssa_aggregator_smoke.jl | 19 +- test/ssa_callback_test.jl | 15 +- test/ssa_kernel_tests.jl | 9 +- test/stochastic_dispatch.jl | 42 +- test/thread_safety.jl | 15 +- 53 files changed, 832 insertions(+), 482 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index a4ec2a767..34d53f597 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -12,7 +12,6 @@ changes are not intended for a 9.x release. # Before (no longer works): jprob = JumpProblem(dprob, Direct(), jump; rng = Xoshiro(1234)) sol = solve(jprob, SSAStepper()) - # After: jprob = JumpProblem(dprob, Direct(), jump) sol = solve(jprob, SSAStepper(); rng = Xoshiro(1234)) @@ -62,7 +61,6 @@ changes are not intended for a 9.x release. ```julia # Before (no longer works): jprob = JumpProblem(dprob, Direct(), maj; scale_rates = false) - # After: maj = MassActionJump(rates, reactant_stoch, net_stoch; scale_rates = false) jprob = JumpProblem(dprob, Direct(), maj) @@ -73,15 +71,15 @@ changes are not intended for a 9.x release. current parameters at initialization and reset. `remake` no longer mutates shared rate coefficients in the jump definition. Fixed-rate and symbolic definitions retain their supported `scaled_rates` representation. This means: - - `update_parameters!` has been removed. Mass action rates are now + + `update_parameters!` has been removed. Mass action rates are now automatically recomputed from the current parameter values whenever the aggregator reinitializes. After modifying parameters (e.g. in a callback), call `reset_aggregated_jumps!(integrator)` to trigger reinitialization with the updated parameter values. - - The `update_jump_params` keyword has been removed from + + The `update_jump_params` keyword has been removed from `reset_aggregated_jumps!`; supplying either `true` or `false` raises an explanatory error. Rate refresh is automatic during reset. - - Custom parameter mappers (e.g. ModelingToolkitBase's + + Custom parameter mappers (e.g. ModelingToolkitBase's `JumpSysMajParamMapper`) must replace the old extraction/mutation callables with `(mapper)(dest::AbstractVector, maj::MassActionJump, params)` and fill `dest` with the scaled rates. The built-in mapper honors @@ -118,7 +116,6 @@ changes are not intended for a 9.x release. ```julia using JumpProcesses - rate(u, p, t) = p[1] * u[1] / (1 + u[2]) function affect!(integrator) integrator.u[1] -= 1 @@ -131,10 +128,8 @@ changes are not intended for a 9.x release. urate = p[1] * uhigh[1] / (1 + ulow[2]) ) end - jump = ConstantRateJump(rate, affect!; bounds) prob = DiscreteProblem([10, 0], (0.0, 10.0), [1.0]) - # Each species affects the rate of jump 1; jump 1 changes both species. vartojumps_map = [[1], [1]] jumptovars_map = [[1, 2]] @@ -157,7 +152,7 @@ changes are not intended for a 9.x release. - Added a default aggregator selection algorithm based on the number of passed in jumps. i.e. the following now auto-selects an aggregator (`Direct` in this case): - + ```julia using JumpProcesses rate(u, p, t) = u[1] @@ -172,10 +167,11 @@ changes are not intended for a 9.x release. `ConstantRateJump`s, and bounded `VariableRateJump`s, one no longer needs to specify `SSAStepper()` when calling `solve`, i.e. the following now works for the previous example and is equivalent to manually passing `SSAStepper()`: - + ```julia sol = solve(jprob) ``` + - Plotting a solution generated with `save_positions = (false, false)` now uses piecewise linear plots between any saved time points specified via `saveat` instead (previously the plots appeared piecewise constant even though each diff --git a/LICENSE.md b/LICENSE.md index 6fb1ddfd4..dcc4017ef 100644 --- a/LICENSE.md +++ b/LICENSE.md @@ -1,17 +1,17 @@ The JumpDiffEq.jl package is licensed under the MIT "Expat" License: > Copyright (c) 2017: ChrisRackauckas. -> +> > Permission is hereby granted, free of charge, to any person obtaining a copy > of this software and associated documentation files (the "Software"), to deal > in the Software without restriction, including without limitation the rights > to use, copy, modify, merge, publish, distribute, sublicense, and/or sell > copies of the Software, and to permit persons to whom the Software is > furnished to do so, subject to the following conditions: -> +> > The above copyright notice and this permission notice shall be included in all > copies or substantial portions of the Software. -> +> > THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR > IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, > FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE diff --git a/README.md b/README.md index e3606dbb3..578131a2d 100644 --- a/README.md +++ b/README.md @@ -6,7 +6,9 @@ + + [![Build Status](https://github.com/SciML/JumpProcesses.jl/workflows/CI/badge.svg)](https://github.com/SciML/JumpProcesses.jl/actions?query=workflow%3ACI) [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor%27s%20Guide-blueviolet)](https://github.com/SciML/ColPrac) [![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle) @@ -38,7 +40,6 @@ The documentation includes ## Contributions welcomed! - Contact us in sciml-bridged on Slack to discuss where to get started, the [`Help wanted`](https://github.com/SciML/JumpProcesses.jl/issues/431) issue, or just open a PR to address an open issue or add new functionality. Contributions, no matter how small, are always welcome and appreciated, including documentation editing/writing. See also the [contribution section](#contributing-and-getting-help). @@ -167,8 +168,9 @@ plot(sol) for guidance on PRs, issues, and other matters relating to contributing to SciML. - See the [SciML Style Guide](https://github.com/SciML/SciMLStyle) for common coding practices and other style decisions. + - There are a few community forums for getting help and asking questions: - + + The #diffeq-bridged and #sciml-bridged channels in the [Julia Slack](https://julialang.org/slack/) + The #diffeq-bridged and #sciml-bridged channels in the diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 8d7167676..71e731c72 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -50,8 +50,8 @@ SUITE["solve"] = BenchmarkGroup() jprob_maj = JumpProblem(dprob, Direct(), maj) jprob_cr = JumpProblem(dprob, Direct(), jump1, jump2) -SUITE["solve"]["massaction"] = @benchmarkable solve($jprob_maj, SSAStepper(); rng) setup=(rng = StableRNG(12345)) evals=1 -SUITE["solve"]["constantrate"] = @benchmarkable solve($jprob_cr, SSAStepper(); rng) setup=(rng = StableRNG(12345)) evals=1 +SUITE["solve"]["massaction"] = @benchmarkable solve($jprob_maj, SSAStepper(); rng) setup=(rng=StableRNG(12345)) evals=1 +SUITE["solve"]["constantrate"] = @benchmarkable solve($jprob_cr, SSAStepper(); rng) setup=(rng=StableRNG(12345)) evals=1 # ============================================================================= # Aggregators @@ -63,8 +63,8 @@ jprob_rdirect = JumpProblem(dprob, RDirect(), maj) jprob_sorting = JumpProblem(dprob, SortingDirect(), maj) jprob_nrm = JumpProblem(dprob, NRM(), maj) -SUITE["aggregators"]["RDirect"] = @benchmarkable solve($jprob_rdirect, SSAStepper(); rng) setup=(rng = StableRNG(12345)) evals=1 +SUITE["aggregators"]["RDirect"] = @benchmarkable solve($jprob_rdirect, SSAStepper(); rng) setup=(rng=StableRNG(12345)) evals=1 SUITE["aggregators"]["SortingDirect"] = @benchmarkable solve( $jprob_sorting, SSAStepper(); rng -) setup=(rng = StableRNG(12345)) evals=1 -SUITE["aggregators"]["NRM"] = @benchmarkable solve($jprob_nrm, SSAStepper(); rng) setup=(rng = StableRNG(12345)) evals=1 +) setup=(rng=StableRNG(12345)) evals=1 +SUITE["aggregators"]["NRM"] = @benchmarkable solve($jprob_nrm, SSAStepper(); rng) setup=(rng=StableRNG(12345)) evals=1 diff --git a/docs/make.jl b/docs/make.jl index 9fd328ff4..92675e33e 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -11,7 +11,8 @@ cp(joinpath(docpath, "Project.toml"), joinpath(assetpath, "Project.toml"), force include("pages.jl") -mathengine = MathJax3(Dict(:loader => Dict("load" => ["[tex]/require", "[tex]/mathtools"]), :tex => Dict("inlineMath" => [["\$", "\$"], ["\\(", "\\)"]], +mathengine = MathJax3(Dict(:loader => Dict("load" => ["[tex]/require", "[tex]/mathtools"]), + :tex => Dict("inlineMath" => [["\$", "\$"], ["\\(", "\\)"]], "packages" => [ "base", "ams", @@ -20,7 +21,8 @@ mathengine = MathJax3(Dict(:loader => Dict("load" => ["[tex]/require", "[tex]/ma "require" ]))) -makedocs(sitename = "JumpProcesses.jl", authors = "Chris Rackauckas", modules = [JumpProcesses], +makedocs( + sitename = "JumpProcesses.jl", authors = "Chris Rackauckas", modules = [JumpProcesses], clean = true, linkcheck = true, checkdocs = :exports, format = Documenter.HTML(; assets = ["assets/favicon.ico"], canonical = "https://docs.sciml.ai/JumpProcesses/", diff --git a/docs/src/api.md b/docs/src/api.md index 1ca9eeed7..82ba9fd08 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -158,11 +158,11 @@ sol = solve(jprob, SSAStepper(); seed = 1234) For SSAStepper, OrdinaryDiffEq ODE/DAE solvers, and all five CPU tau-leaping algorithms, an explicit `rng` takes priority over `seed`: -| User provides | Result | -|---|---| -| `rng` via `solve`/`init` | Uses that `rng` | -| `seed` via `solve`/`init` | Creates `Xoshiro(seed)` | -| Nothing | Uses `Random.default_rng()` | +| User provides | Result | +|:------------------------- |:--------------------------- | +| `rng` via `solve`/`init` | Uses that `rng` | +| `seed` via `solve`/`init` | Creates `Xoshiro(seed)` | +| Nothing | Uses `Random.default_rng()` | StochasticDiffEq handles SDE/RODE RNGs for both `solve` and `init`. An explicit RNG other than `TaskLocalRNG` takes priority over seeds. With no RNG, or with @@ -175,15 +175,16 @@ solvers use `Xoshiro(0)`. ### Behavior by solver pathway -| Solver | Default RNG (nothing passed) | `rng` / `seed` support | -|---|---|---| -| `SSAStepper` | `Random.default_rng()` | Full support via `solve`/`init` kwargs | -| OrdinaryDiffEq ODE/DAE solvers (e.g., `Tsit5`, `DFBDF`) | `Random.default_rng()` | Full support via `solve`/`init` kwargs | -| StochasticDiffEq SDE/RODE solvers (e.g., `SRIW1`, `RandomEM`) | `Xoshiro` from the stored problem seed, or a random seed | Full support; `TaskLocalRNG` follows the seed policy above and is converted to `Xoshiro` | -| `SimpleTauLeaping` | `Random.default_rng()` | Full support via `solve` kwargs | -| `SimpleExplicitTauLeaping`, `SimpleImplicitTauLeaping`, `SimpleTrapezoidalLeaping`, `SimpleAdaptiveTauLeaping` | `Random.default_rng()` | Full support via `solve` kwargs | +| Solver | Default RNG (nothing passed) | `rng` / `seed` support | +|:-------------------------------------------------------------------------------------------------------------- |:-------------------------------------------------------- |:---------------------------------------------------------------------------------------- | +| `SSAStepper` | `Random.default_rng()` | Full support via `solve`/`init` kwargs | +| OrdinaryDiffEq ODE/DAE solvers (e.g., `Tsit5`, `DFBDF`) | `Random.default_rng()` | Full support via `solve`/`init` kwargs | +| StochasticDiffEq SDE/RODE solvers (e.g., `SRIW1`, `RandomEM`) | `Xoshiro` from the stored problem seed, or a random seed | Full support; `TaskLocalRNG` follows the seed policy above and is converted to `Xoshiro` | +| `SimpleTauLeaping` | `Random.default_rng()` | Full support via `solve` kwargs | +| `SimpleExplicitTauLeaping`, `SimpleImplicitTauLeaping`, `SimpleTrapezoidalLeaping`, `SimpleAdaptiveTauLeaping` | `Random.default_rng()` | Full support via `solve` kwargs | !!! note + Use an explicit `rng` or `seed` where the solver supports it, and start from identical model and mutable aggregator state. An explicit RNG is advanced by the solve; construct or copy it from the same starting state for replay. @@ -215,13 +216,13 @@ The following table documents which code handles `solve`/`init` for each solver type. This is relevant for developers working on JumpProcesses or its solver backends. -| Solver type | `__solve` handled by | `__init` handled by | Uses `__jump_init`? | -|---|---|---|---| -| `SSAStepper` | JumpProcesses (`solve.jl`) | JumpProcesses (`SSA_stepper.jl`) | No | -| OrdinaryDiffEq ODE/DAE (e.g., `Tsit5`, `DFBDF`) | JumpProcesses (`solve.jl`) | JumpProcesses' OrdinaryDiffEqCore extension → OrdinaryDiffEq | Yes | -| StochasticDiffEq SDE/RODE (e.g., `SRIW1`, `RandomEM`) | StochasticDiffEqCore | StochasticDiffEqCore's JumpProblem initializer | No | -| StochasticDiffEq jump algorithms (e.g., `TauLeaping`) | StochasticDiffEqCore | StochasticDiffEqCore's specialized jump-algorithm initializer | No | -| All five CPU tau-leaping algorithms | JumpProcesses (`simple_regular_solve.jl`, custom `DiffEqBase.solve`) | N/A | No | +| Solver type | `__solve` handled by | `__init` handled by | Uses `__jump_init`? | +|:----------------------------------------------------- |:-------------------------------------------------------------------- |:------------------------------------------------------------- |:------------------- | +| `SSAStepper` | JumpProcesses (`solve.jl`) | JumpProcesses (`SSA_stepper.jl`) | No | +| OrdinaryDiffEq ODE/DAE (e.g., `Tsit5`, `DFBDF`) | JumpProcesses (`solve.jl`) | JumpProcesses' OrdinaryDiffEqCore extension → OrdinaryDiffEq | Yes | +| StochasticDiffEq SDE/RODE (e.g., `SRIW1`, `RandomEM`) | StochasticDiffEqCore | StochasticDiffEqCore's JumpProblem initializer | No | +| StochasticDiffEq jump algorithms (e.g., `TauLeaping`) | StochasticDiffEqCore | StochasticDiffEqCore's specialized jump-algorithm initializer | No | +| All five CPU tau-leaping algorithms | JumpProcesses (`simple_regular_solve.jl`, custom `DiffEqBase.solve`) | N/A | No | For **SSAStepper**, `rng` is resolved via `resolve_rng` in `SSA_stepper.jl`'s `__init` and stored on the [`SSAIntegrator`](@ref). diff --git a/docs/src/applications/advanced_point_process.md b/docs/src/applications/advanced_point_process.md index 7813c83c6..cdc99de53 100644 --- a/docs/src/applications/advanced_point_process.md +++ b/docs/src/applications/advanced_point_process.md @@ -833,11 +833,14 @@ outside the scope this tutorial. [^1]: D. J. Daley and D. Vere-Jones, An Introduction to the Theory of Point Processes: Volume I: Elementary Theory and Methods, Springer-Verlag (2003). doi:10.1007/b97277. + [^2]: T., Björk, Point Processes and Jump Diffusions: An Introduction with Finance Applications, Cambridge University Press (2021). doi:10.1017/9781009002127. + [^3]: F. B. Hanson, Applied Stochastic Processes and Control for Jump-Diffusions: Modeling, Analysis and Computation, Society for Industrial and Applied Mathematics (2007). doi:10.1137/1.9780898718638. + [^4]: P. J. Laub, Y. Lee and T. Taimre, The Elements of Hawkes Processes, Springer International Publishing (2021). doi:10.1007/978-3-030-84639-8. diff --git a/docs/src/index.md b/docs/src/index.md index dbd5be6b3..7d0e03205 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -146,8 +146,9 @@ Pkg.add("JumpProcesses") for guidance on PRs, issues, and other matters relating to contributing to SciML. - See the [SciML Style Guide](https://github.com/SciML/SciMLStyle) for common coding practices and other style decisions. + - There are a few community forums for getting help and asking questions: - + + The #diffeq-bridged and #sciml-bridged channels in the [Julia Slack](https://julialang.org/slack/) + The #diffeq-bridged and #sciml-bridged channels in the diff --git a/docs/src/jump_types.md b/docs/src/jump_types.md index cf6f02941..1128b8a9f 100644 --- a/docs/src/jump_types.md +++ b/docs/src/jump_types.md @@ -75,7 +75,7 @@ their own special time integrators. The constructor for a [`ConstantRateJump`](@ref) is: ```julia -ConstantRateJump(rate, affect!; bounds=nothing) +ConstantRateJump(rate, affect!; bounds = nothing) ``` - `rate(u, p, t)` is a function which calculates the rate given the current @@ -134,7 +134,7 @@ scaling and the removed `JumpProblem` constructor keywords. `k*A*(A-1)*(A-2)/3!`. To *avoid* having the reaction rates rescaled (by `1/2` and `1/6` for these two examples), one can pass the `MassActionJump` constructor the optional named parameter `scale_rates = false`, i.e., use - + ```julia MassActionJump(reactant_stoich, net_stoich; scale_rates = false, param_idxs) ``` @@ -142,31 +142,32 @@ scaling and the removed `JumpProblem` constructor keywords. - Zero order reactions can be passed as `reactant_stoich`s in one of two ways. Consider the ``\varnothing \overset{k}{\rightarrow} A`` reaction with rate `k=1`: - + ```julia p = [1.0] reactant_stoich = [[0 => 1]] net_stoich = [[1 => 1]] jump = MassActionJump(reactant_stoich, net_stoich; param_idxs = [1]) ``` - + Alternatively, one can create an empty vector of pairs to represent the reaction: - + ```julia p = [1.0] reactant_stoich = [Vector{Pair{Int, Int}}()] net_stoich = [[1 => 1]] jump = MassActionJump(reactant_stoich, net_stoich; param_idxs = [1]) ``` + - For performance reasons, it is recommended to order species indices in stoichiometry vectors from smallest to largest. That is - + ```julia reactant_stoich = [[1 => 2, 3 => 1, 4 => 2], [2 => 2, 3 => 2]] ``` - + is preferred over - + ```julia reactant_stoich = [[3 => 1, 1 => 2, 4 => 2], [3 => 2, 2 => 2]] ``` @@ -492,29 +493,37 @@ as will trying to update either `p` or `tspan` while passing a new [^1]: V. Lemaire, M. Thieullen and N. Thomas, Exact Simulation of the Jump Times of a Class of Piecewise Deterministic Markov Processes, Journal of Scientific Computing, 75 (3), 1776-1807 (2018). doi:10.1007/s10915-017-0607-4. + [^2]: D. T. Gillespie, A general method for numerically simulating the stochastic time evolution of coupled chemical reactions, Journal of Computational Physics, 22 (4), 403–434 (1976). doi:10.1016/0021-9991(76)90041-3. + [^3]: A. Slepoy, A.P. Thompson and S.J. Plimpton, A constant-time kinetic Monte Carlo algorithm for simulation of large biochemical reaction networks, Journal of Chemical Physics, 128 (20), 205101 (2008). doi:10.1063/1.2919546. + [^4]: J. M. McCollum, G. D. Peterson, C. D. Cox, M. L. Simpson and N. F. Samatova, The sorting direct method for stochastic simulation of biochemical systems with varying reaction execution behavior, Computational Biology and Chemistry, 30 (1), 39049 (2006). doi:10.1016/j.compbiolchem.2005.10.007. + [^5]: V. H. Thanh, C. Priami and R. Zunino, Efficient rejection-based simulation of biochemical reactions with stochastic noise and delays, Journal of Chemical Physics, 141 (13), 134116 (2014). doi:10.1063/1.4896985. + [^6]: V. H. Thanh, R. Zunino and C. Priami, On the rejection-based algorithm for simulation and analysis of large-scale reaction networks, Journal of Chemical Physics, 142 (24), 244106 (2015). doi:10.1063/1.4922923. + [^7]: V. H. Thanh, R. Zunino, and C. Priami, Efficient constant-time complexity algorithm for stochastic simulation of large reaction networks, IEEE/ACM Transactions on Computational Biology and Bioinformatics, 14 (3), 657-667 (2017). doi:10.1109/TCBB.2016.2530066. + [^8]: M. A. Gibson and J. Bruck, Efficient exact stochastic simulation of chemical systems with many species and many channels, Journal of Physical Chemistry A, 104 (9), 1876-1889 (2000). doi:10.1021/jp993732q. + [^9]: M. Farajtabar, Y. Wang, M. Gomez-Rodriguez, S. Li, H. Zha, and L. Song, COEVOLVE: a joint point process model for information diffusion and network evolution, Journal of Machine Learning Research 18(1), 1305–1353 (2017). doi: diff --git a/docs/src/tutorials/discrete_stochastic_example.md b/docs/src/tutorials/discrete_stochastic_example.md index fb2389ad2..03ef1bc17 100644 --- a/docs/src/tutorials/discrete_stochastic_example.md +++ b/docs/src/tutorials/discrete_stochastic_example.md @@ -21,7 +21,7 @@ chemical kinetics (i.e., Gillespie) models. It is not necessary to have read the τ-leaping methods. !!! note - + This tutorial assumes you have read the [Ordinary Differential Equations tutorial](https://docs.sciml.ai/DiffEqDocs/stable/getting_started/) in [`DifferentialEquations.jl`](https://docs.sciml.ai/DiffEqDocs/stable). We begin by demonstrating how to build jump processes using @@ -931,6 +931,7 @@ u₀ = [1000.0, 50.0, 0.0] prob = DiscreteProblem(u₀, tspan, p) jump_prob = JumpProblem(prob, PureLeaping(), rj) ``` + We pass the `PureLeaping` aggregator to indicate we will use a τ-leaping algorithm to simulate the process. Note that when a `JumpProblem` has a `RegularJump`, τ-leaping algorithms are required for simulating it. This is detailed on the [jump solvers page](@ref diff --git a/docs/src/tutorials/jump_diffusion.md b/docs/src/tutorials/jump_diffusion.md index 8b73b8df6..1caacba4d 100644 --- a/docs/src/tutorials/jump_diffusion.md +++ b/docs/src/tutorials/jump_diffusion.md @@ -1,7 +1,7 @@ # [Piecewise Deterministic Markov Processes and Jump Diffusion Equations](@id jump_diffusion_tutorial) !!! note - + This tutorial assumes you have read the [Ordinary Differential Equations tutorial](https://docs.sciml.ai/DiffEqDocs/stable/getting_started/) in [`DifferentialEquations.jl`](https://docs.sciml.ai/DiffEqDocs/stable). Jump Diffusion equations are stochastic differential equations (SDEs) with discontinuous diff --git a/ext/JumpProcessesForwardDiffExt.jl b/ext/JumpProcessesForwardDiffExt.jl index f603ac57b..ebff4133f 100644 --- a/ext/JumpProcessesForwardDiffExt.jl +++ b/ext/JumpProcessesForwardDiffExt.jl @@ -6,7 +6,7 @@ import ForwardDiff @inline function DiffEqBase.ODE_DEFAULT_NORM( u::ExtendedJumpArray{<:ForwardDiff.Dual}, t::ForwardDiff.Dual - ) +) return invoke(DiffEqBase.ODE_DEFAULT_NORM, Tuple{ExtendedJumpArray, Any}, u, t) end diff --git a/ext/JumpProcessesKernelAbstractionsExt.jl b/ext/JumpProcessesKernelAbstractionsExt.jl index c5275252b..342cdf711 100644 --- a/ext/JumpProcessesKernelAbstractionsExt.jl +++ b/ext/JumpProcessesKernelAbstractionsExt.jl @@ -138,7 +138,7 @@ end # Get input/output arrays ts_view = @inbounds view(_ts, :, i) - us_view = @inbounds view(_us,:,:,i) + us_view = @inbounds view(_us, :, :, i) # Initialize first time step and state @inbounds ts_view[1] = tspan[1] diff --git a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl index cf9697b45..df907d6b0 100644 --- a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl +++ b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl @@ -4,7 +4,7 @@ using JumpProcesses import DiffEqBase import SciMLBase import OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, DAEAlgorithm, - StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm + StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm function _jump_init(_jump_prob, alg; merge_callbacks = true, kwargs...) kwargs = DiffEqBase.merge_problem_kwargs(_jump_prob; merge_callbacks, kwargs...) diff --git a/ext/ssa_stepper.jl b/ext/ssa_stepper.jl index e32a0127d..98b050fec 100644 --- a/ext/ssa_stepper.jl +++ b/ext/ssa_stepper.jl @@ -13,19 +13,33 @@ Rate constants are stored in the same floating point type as the problem's `tspan` so that rate and time arithmetic in the kernel share one type. """ struct GPUMassActionJump{R, S, C, O} - """Rate constants, already scaled by the mass action combinatoric prefactors.""" + """ + Rate constants, already scaled by the mass action combinatoric prefactors. + """ scaled_rates::R - """Flattened reactant species indices.""" + """ + Flattened reactant species indices. + """ rs_species::S - """Flattened reactant stoichiometric coefficients.""" + """ + Flattened reactant stoichiometric coefficients. + """ rs_coeffs::C - """CSR offsets into the reactant arrays, of length `numjumps + 1`.""" + """ + CSR offsets into the reactant arrays, of length `numjumps + 1`. + """ rs_offsets::O - """Flattened net stoichiometry species indices.""" + """ + Flattened net stoichiometry species indices. + """ ns_species::S - """Flattened net stoichiometric coefficients.""" + """ + Flattened net stoichiometric coefficients. + """ ns_coeffs::C - """CSR offsets into the net stoichiometry arrays, of length `numjumps + 1`.""" + """ + CSR offsets into the net stoichiometry arrays, of length `numjumps + 1`. + """ ns_offsets::O end diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl index 554b74325..bc37376db 100644 --- a/src/JumpProcesses.jl +++ b/src/JumpProcesses.jl @@ -179,7 +179,7 @@ export SSAStepper # leaping: include("simple_regular_solve.jl") export SimpleTauLeaping, SimpleExplicitTauLeaping, SimpleImplicitTauLeaping, - SimpleTrapezoidalLeaping, SimpleAdaptiveTauLeaping, EnsembleGPUKernel + SimpleTrapezoidalLeaping, SimpleAdaptiveTauLeaping, EnsembleGPUKernel # spatial: include("spatial/spatial_massaction_jump.jl") diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index 919ace32d..07c87b7d1 100644 --- a/src/SSA_stepper.jl +++ b/src/SSA_stepper.jl @@ -74,45 +74,81 @@ $(FIELDS) """ mutable struct SSAIntegrator{F, uType, tType, tdirType, P, S, CB, SA, OPT, TS, R} <: AbstractSSAIntegrator{SSAStepper, Nothing, uType, tType} - """The underlying `prob.f` function. Not currently used.""" + """ + The underlying `prob.f` function. Not currently used. + """ f::F - """The current solution values.""" + """ + The current solution values. + """ u::uType - """The current solution time.""" + """ + The current solution time. + """ t::tType - """The previous time a jump occurred.""" + """ + The previous time a jump occurred. + """ tprev::tType - """The direction time is changing in (must be positive, indicating time is increasing)""" + """ + The direction time is changing in (must be positive, indicating time is increasing) + """ tdir::tdirType - """The current parameters.""" + """ + The current parameters. + """ p::P - """The current solution object.""" + """ + The current solution object. + """ sol::S i::Int - """The next jump time.""" + """ + The next jump time. + """ tstop::tType - """The jump aggregator callback.""" + """ + The jump aggregator callback. + """ cb::CB - """Times to save the solution at.""" + """ + Times to save the solution at. + """ saveat::SA - """Whether to save every time a jump occurs.""" + """ + Whether to save every time a jump occurs. + """ save_everystep::Bool - """Whether to save at the final step.""" + """ + Whether to save at the final step. + """ save_end::Bool - """Index of the next `saveat` time.""" + """ + Index of the next `saveat` time. + """ cur_saveat::Int - """Tuple storing callbacks.""" + """ + Tuple storing callbacks. + """ opts::OPT - """User supplied times to step to, useful with callbacks.""" + """ + User supplied times to step to, useful with callbacks. + """ tstops::TS tstops_idx::Int u_modified::Bool keep_stepping::Bool # false if should terminate a simulation - """If true, will write tstops into the user-passed array""" + """ + If true, will write tstops into the user-passed array + """ alias_tstops::Bool - """If true indicates we have already allocated the tstops array""" + """ + If true indicates we have already allocated the tstops array + """ copied_tstops::Bool - """The random number generator.""" + """ + The random number generator. + """ rng::R end @@ -187,7 +223,7 @@ function DiffEqBase.solve!(integrator::SSAIntegrator) if integrator.saveat !== nothing && !isempty(integrator.saveat) # Split to help prediction while integrator.cur_saveat <= length(integrator.saveat) && - integrator.saveat[integrator.cur_saveat] < integrator.t + integrator.saveat[integrator.cur_saveat] < integrator.t push!(integrator.sol.t, integrator.saveat[integrator.cur_saveat]) push!(integrator.sol.u, copy(integrator.u)) integrator.cur_saveat += 1 @@ -402,7 +438,7 @@ function DiffEqBase.step!(integrator::SSAIntegrator) @inbounds if integrator.saveat !== nothing && !isempty(integrator.saveat) # Split to help prediction while integrator.cur_saveat <= length(integrator.saveat) && - integrator.saveat[integrator.cur_saveat] < integrator.t + integrator.saveat[integrator.cur_saveat] < integrator.t push!(integrator.sol.t, integrator.saveat[integrator.cur_saveat]) push!(integrator.sol.u, copy(integrator.u)) integrator.cur_saveat += 1 @@ -478,7 +514,6 @@ function DiffEqBase.terminate!(integrator::SSAIntegrator, retcode = ReturnCode.T nothing end - function SciMLBase.isdenseplot(sol::ODESolution{ T, N, uType, uType2, DType, tType, rateType, discType, P, SSAStepper}) where {T, N, uType, uType2, DType, tType, rateType, discType, P} diff --git a/src/aggregators/prioritytable.jl b/src/aggregators/prioritytable.jl index 0d96561be..5a6ccbaa4 100644 --- a/src/aggregators/prioritytable.jl +++ b/src/aggregators/prioritytable.jl @@ -20,13 +20,19 @@ the interval to the next group. i.e. maxpriority -> N+1 One group (i.e. bin) of priority ids within the table """ mutable struct PriorityGroup{T, W <: AbstractVector} - "(strict) upper bound for priorities in this group" + """ + (strict) upper bound for priorities in this group + """ maxpriority::T - "number of priority ids" + """ + number of priority ids + """ numpids::Int - "priority ids associated with group" + """ + priority ids associated with group + """ pids::W end PriorityGroup{U}(maxpriority::T) where {T, U} = PriorityGroup(maxpriority, 0, Vector{U}()) @@ -74,25 +80,39 @@ Table to store the groups. abstract type AbstractPriorityTable end mutable struct PriorityTable{F, S, T, U <: Function} <: AbstractPriorityTable - "non-zero values below this are binned together, static" + """ + non-zero values below this are binned together, static + """ minpriority::F - "values above this cause new groups to be added" + """ + values above this cause new groups to be added + """ maxpriority::F - "bins storing priority ids within a given range" + """ + bins storing priority ids within a given range + """ groups::Vector{PriorityGroup{F, Vector{S}}} - "stores the sums of the priorities within each group" + """ + stores the sums of the priorities within each group + """ gsums::Vector{F} - "stores the sum of the group sums" + """ + stores the sum of the group sums + """ gsum::F - "maps priority id to group and idx within the group" + """ + maps priority id to group and idx within the group + """ pidtogroup::Vector{Tuple{T, T}} - "mapping from priority value to group id that stores it" + """ + mapping from priority value to group id that stores it + """ priortogid::U end diff --git a/src/aggregators/rssa.jl b/src/aggregators/rssa.jl index 23b273bd6..caa3b7c27 100644 --- a/src/aggregators/rssa.jl +++ b/src/aggregators/rssa.jl @@ -159,7 +159,8 @@ Update rates # for each dependent jump, update jump rate brackets for jidx in p.vartojumps_map[uidx] sum_rate -= crhigh[jidx] - p.cur_rate_low[jidx], crhigh[jidx] = get_jump_brackets(jidx, p, u, params, t) + p.cur_rate_low[jidx], crhigh[jidx] = get_jump_brackets( + jidx, p, u, params, t) sum_rate += crhigh[jidx] end end @@ -185,7 +186,8 @@ end # for each dependent jump, update jump rate brackets for jidx in p.vartojumps_map[uidx] sum_rate -= crhigh[jidx] - p.cur_rate_low[jidx], crhigh[jidx] = get_jump_brackets(jidx, p, u, params, t) + p.cur_rate_low[jidx], crhigh[jidx] = get_jump_brackets( + jidx, p, u, params, t) sum_rate += crhigh[jidx] end end diff --git a/src/aggregators/rssacr.jl b/src/aggregators/rssacr.jl index 718730242..9bf3f034f 100644 --- a/src/aggregators/rssacr.jl +++ b/src/aggregators/rssacr.jl @@ -182,7 +182,8 @@ update bracketing for species that depend on the just executed jump # for each dependent jump, update jump rate brackets for jidx in p.vartojumps_map[uidx] oldrate = crhigh[jidx] - p.cur_rate_low[jidx], crhigh[jidx] = get_jump_brackets(jidx, p, u, params, t) + p.cur_rate_low[jidx], crhigh[jidx] = get_jump_brackets( + jidx, p, u, params, t) # update the priority table update!(p.rt, jidx, oldrate, crhigh[jidx]) @@ -210,7 +211,8 @@ end # for each dependent jump, update jump rate brackets for jidx in p.vartojumps_map[uidx] oldrate = crhigh[jidx] - p.cur_rate_low[jidx], crhigh[jidx] = get_jump_brackets(jidx, p, u, params, t) + p.cur_rate_low[jidx], crhigh[jidx] = get_jump_brackets( + jidx, p, u, params, t) # update the priority table update!(p.rt, jidx, oldrate, crhigh[jidx]) diff --git a/src/aggregators/ssajump.jl b/src/aggregators/ssajump.jl index 73ec7bd4d..cc6a74d51 100644 --- a/src/aggregators/ssajump.jl +++ b/src/aggregators/ssajump.jl @@ -13,6 +13,7 @@ An aggregator interface for SSA-like algorithms. - `rates` # vector of rate functions for ConstantRateJumps - `affects!` # vector of affect functions for ConstantRateJumps - `save_positions` # tuple for whether to save the jumps before and/or after event + ### Optional fields: - `dep_gr` # dependency graph, dep_gr[i] = indices of reactions that should be updated when rx i occurs. diff --git a/src/coupling.jl b/src/coupling.jl index b91fc3465..55bdb9344 100644 --- a/src/coupling.jl +++ b/src/coupling.jl @@ -159,7 +159,7 @@ function build_split_jumps(prob::SciMLBase.AbstractJumpProblem, # only prob new_affect! = affect! new_rate = (u, p, t) -> rate(u.u, p, t) - - min(rate(u.u, p, t), rate_control(u.u_control, p, t)) + min(rate(u.u, p, t), rate_control(u.u_control, p, t)) push!(jumps, ConstantRateJump(new_rate, new_affect!)) # only prob_control new_affect! = function (integrator) @@ -168,7 +168,7 @@ function build_split_jumps(prob::SciMLBase.AbstractJumpProblem, flip_u!(integrator.u) end new_rate = (u, p, t) -> rate_control(u.u_control, p, t) - - min(rate(u.u, p, t), rate_control(u.u_control, p, t)) + min(rate(u.u, p, t), rate_control(u.u_control, p, t)) push!(jumps, ConstantRateJump(new_rate, new_affect!)) end jumps diff --git a/src/extended_jump_array.jl b/src/extended_jump_array.jl index 82c2b1c68..075e8227a 100644 --- a/src/extended_jump_array.jl +++ b/src/extended_jump_array.jl @@ -54,17 +54,23 @@ sol = solve(jprob, Tsit5()) - In a system with `VariableRateJump`s all callback, `ConstantRateJump`, and `VariableRateJump` `affect!` functions will receive integrators with `integrator.u` an `ExtendedJumpArray`. + - As such, `affect!` functions that wish to modify the state via vector operations should use `ueja.u.u` to obtain the aliased state object. + - `:u` and `:jump_u` always refer to the `ExtendedJumpArray` fields. Any other property access (and `propertynames`) is forwarded to the wrapped state `ueja.u`, so e.g. `ueja.x` works when `ueja.u` is an `ArrayPartition`. """ struct ExtendedJumpArray{T3 <: Number, T1, T <: AbstractArray{T3, T1}, T2} <: AbstractArray{T3, 1} - """The current state.""" + """ + The current state. + """ u::T - """The current rate (i.e. hazard, intensity, or propensity) values for the `VariableRateJump`s.""" + """ + The current rate (i.e. hazard, intensity, or propensity) values for the `VariableRateJump`s. + """ jump_u::T2 end @@ -122,11 +128,14 @@ function _mul_extended_jump_array!(c::ExtendedJumpArray, A::AbstractVecOrMat, end end -LinearAlgebra.mul!(c::ExtendedJumpArray, A::AbstractVecOrMat, u::AbstractVector) = +function LinearAlgebra.mul!(c::ExtendedJumpArray, A::AbstractVecOrMat, u::AbstractVector) _mul_extended_jump_array!(c, A, u) +end -LinearAlgebra.mul!(c::ExtendedJumpArray, A::LinearAlgebra.AbstractTriangular, - u::AbstractVector) = _mul_extended_jump_array!(c, A, u) +function LinearAlgebra.mul!(c::ExtendedJumpArray, A::LinearAlgebra.AbstractTriangular, + u::AbstractVector) + _mul_extended_jump_array!(c, A, u) +end # Whole-array copies stay ExtendedJumpArrays. A subset getindex asks `similar` # for a container of the index shape, which no longer matches this array's axes; diff --git a/src/jumps.jl b/src/jumps.jl index 2f847e59e..ef1688dd9 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -17,13 +17,17 @@ $(FIELDS) ## Examples Increasing rate in one species but decreasing in another: + ```julia rate(u, p, t) = p[1] * u[1] / (1 + u[2]) -bounds(ulow, uhigh, u, p, t) = RateBounds(lrate = p[1] * ulow[1] / (1 + uhigh[2]), - urate = p[1] * uhigh[1] / (1 + ulow[2])) +function bounds(ulow, uhigh, u, p, t) + RateBounds(lrate = p[1] * ulow[1] / (1 + uhigh[2]), + urate = p[1] * uhigh[1] / (1 + ulow[2])) +end ``` Extremas lying inside the bracket: + ```julia rate(u, p, t) = p[1] * u[1] * (p[2] - u[1]) function bounds(ulow, uhigh, u, p, t) @@ -35,26 +39,33 @@ end ``` Rate that depends on time: + ```julia rate(u, p, t) = p[1] * u[1] * exp(-p[2] * t) function bounds(ulow, uhigh, u, p, t) Δ = p[3] RateBounds(lrate = p[1] * ulow[1] * exp(-p[2] * (t + Δ)), - urate = p[1] * uhigh[1] * exp(-p[2] * t), - rateinterval = Δ) + urate = p[1] * uhigh[1] * exp(-p[2] * t), + rateinterval = Δ) end ``` """ struct RateBounds{R, T} - """Lower bound of the rate.""" + """ + Lower bound of the rate. + """ lrate::R - """Upper bound of the rate.""" + """ + Upper bound of the rate. + """ urate::R - """Time window over which the bounds hold.""" + """ + Time window over which the bounds hold. + """ rateinterval::T end -function RateBounds(; lrate=nothing, urate=nothing, rateinterval = Inf) +function RateBounds(; lrate = nothing, urate = nothing, rateinterval = Inf) (lrate === nothing && urate === nothing) && error("`RateBounds` requires at least one of `lrate` or `urate`.") l = lrate === nothing ? zero(urate) : lrate @@ -78,16 +89,23 @@ $(FIELDS) - All fields are optional, but an aggregator errors at initialization if the form it consumes is missing. """ struct RateBoundFunctions{B, L, U} - """Computes both bounds in a single call.""" + """ + Computes both bounds in a single call. + """ bounds::B - """Computes the lower bound.""" + """ + Computes the lower bound. + """ lrate::L - """Computes the upper bound.""" + """ + Computes the upper bound. + """ urate::U end -RateBoundFunctions(; bounds = nothing, lrate = nothing, urate = nothing) = +function RateBoundFunctions(; bounds = nothing, lrate = nothing, urate = nothing) RateBoundFunctions(bounds, lrate, urate) +end hasbounds(::Nothing) = false hasbounds(::RateBoundFunctions{Nothing}) = false @@ -108,45 +126,58 @@ Defines a jump process with a rate (i.e. hazard, intensity, or propensity) that does not *explicitly* depend on time. More precisely, one where the rate function is constant *between* the occurrence of jumps. For detailed examples and usage information, see the -- [Tutorial](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) + + - [Tutorial](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) ## Fields $(FIELDS) ## Examples + Suppose `u[1]` gives the number of particles and `p[1]` the probability per time each particle can decay away. A corresponding `ConstantRateJump` for this jump process is + ```julia -rate(u,p,t) = p[1]*u[1] +rate(u, p, t) = p[1]*u[1] affect!(integrator) = integrator.u[1] -= 1 crj = ConstantRateJump(rate, affect!) ``` + Notice, here that `rate` changes in time, but is constant between the occurrence of jumps (when `u[1]` will decrease). Rate bounds may be supplied separately, for use with bracketing aggregators such as `RSSA`: + ```julia rate(u, p, t) = p[1] * u[1] / (1 + u[2]) affect!(integrator) = integrator.u[1] -= 1 -bounds(ulow, uhigh, u, p, t) = RateBounds(lrate=p[1]*ulow[1] / (1+uhigh[2]), - urate=p[1]*uhigh[1] / (1+ulow[2])) +function bounds(ulow, uhigh, u, p, t) + RateBounds(lrate = p[1]*ulow[1] / (1+uhigh[2]), + urate = p[1]*uhigh[1] / (1+ulow[2])) +end crj = ConstantRateJump(rate, affect!; bounds) ``` ## Notes -- When rate bounds are not supplied, they are computed by evaluating the rate at `ulow` - and `uhigh`. These bounds are only correct if the rate is monotonic with respect to - state, i.e. increasing in all species or decreasing in all species. + - When rate bounds are not supplied, they are computed by evaluating the rate at `ulow` + and `uhigh`. These bounds are only correct if the rate is monotonic with respect to + state, i.e. increasing in all species or decreasing in all species. """ struct ConstantRateJump{F1, F2, B <: Union{Nothing, RateBoundFunctions}} <: AbstractJump - """Function `rate(u,p,t)` that returns the jump's current rate.""" + """ + Function `rate(u,p,t)` that returns the jump's current rate. + """ rate::F1 - """Function `affect(integrator)` that updates the state for one occurrence of the jump.""" + """ + Function `affect(integrator)` that updates the state for one occurrence of the jump. + """ affect!::F2 - """Optional `RateBoundFunctions` holding user supplied bracketing functions.""" + """ + Optional `RateBoundFunctions` holding user supplied bracketing functions. + """ bounds::B end @@ -194,22 +225,21 @@ end # fallback assumes the rate is monotonic in the state; it would be cleaner to demand # user-specified bounds for decreasing rates and only handle the increasing case by default @inline function cjump_brackets(c::ConstantRateJump{F1, F2, Nothing}, - ulow, uhigh, u, p, t) where {F1, F2} + ulow, uhigh, u, p, t) where {F1, F2} rlow = c.rate(ulow, p, t) rhigh = c.rate(uhigh, p, t) rlow <= rhigh ? (rlow, rhigh) : (rhigh, rlow) end -@inline lower_rate_bound(c::ConstantRateJump, ulow, uhigh, u, p, t) = - c.bounds.lrate(ulow, uhigh, u, p, t).lrate +@inline lower_rate_bound(c::ConstantRateJump, ulow, uhigh, u, p, t) = c.bounds.lrate( + ulow, uhigh, u, p, t).lrate @inline lower_rate_bound(c::ConstantRateJump{F1, F2, Nothing}, ulow, uhigh, u, p, - t) where {F1, F2} = min(c.rate(ulow, p, t), c.rate(uhigh, p, t)) + t) where {F1, F2} = min(c.rate(ulow, p, t), c.rate(uhigh, p, t)) -@inline upper_rate_bound(c::ConstantRateJump, ulow, uhigh, u, p, t) = - c.bounds.urate(ulow, uhigh, u, p, t).urate +@inline upper_rate_bound(c::ConstantRateJump, ulow, uhigh, u, p, t) = c.bounds.urate( + ulow, uhigh, u, p, t).urate @inline upper_rate_bound(c::ConstantRateJump{F1, F2, Nothing}, ulow, uhigh, u, p, - t) where {F1, F2} = max(c.rate(ulow, p, t), c.rate(uhigh, p, t)) - + t) where {F1, F2} = max(c.rate(ulow, p, t), c.rate(uhigh, p, t)) """ $(TYPEDEF) @@ -217,94 +247,104 @@ $(TYPEDEF) Defines a jump process with a rate (i.e. hazard, intensity, or propensity) that may explicitly depend on time. More precisely, one where the rate function is allowed to change *between* the occurrence of jumps. For detailed examples and usage information, see the -- [Tutorial](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) + + - [Tutorial](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) Note that two types of `VariableRateJump`s are currently supported, with different performance characteritistics. -- A general `VariableRateJump` or `VariableRateJump` will refer to one in which only `rate` - and `affect` functions are specified. - - * These are the most general in what they can represent, but require the use of an - `ODEProblem` or `SDEProblem` whose underlying timestepper handles their evolution in - time (via the callback interface). - * This is the least performant jump type in simulations. - -- Bounded `VariableRateJump`s require passing the keyword arguments `urate` and - `rateinterval`, corresponding to functions `urate(u, p, t)` and `rateinterval(u, p, t)`, - see below. These must calculate a time window over which the rate function is bounded by a - constant. Note that it is ok if the rate bound would be violated within the time interval - due to a change in `u` arising from another `ConstantRateJump`, `MassActionJump` or - *bounded* `VariableRateJump` being executed, as the chosen aggregator will then handle - recalculating the rate bound and interval. *However, if the bound could be violated within - the time interval due to a change in `u` arising from continuous dynamics such as a - coupled ODE, SDE, or a general `VariableRateJump`, bounds should not be given.* This - ensures the jump is classified as a general `VariableRateJump` and properly handled. One - can also optionally provide a lower bound function, `lrate(u, p, t)`, via the `lrate` - keyword argument. This can lead to increased performance. The validity of the lower bound - should hold under the same conditions and rate interval as `urate`. - - * Bounded `VariableRateJump`s can currently be used in the `Coevolve` aggregator, and - can therefore be efficiently simulated in pure-jump `DiscreteProblem`s using the - `SSAStepper` time-stepper. - * These can be substantially more performant than general `VariableRateJump`s without - the rate bound functions. + + - A general `VariableRateJump` or `VariableRateJump` will refer to one in which only `rate` + and `affect` functions are specified. + + + These are the most general in what they can represent, but require the use of an + `ODEProblem` or `SDEProblem` whose underlying timestepper handles their evolution in + time (via the callback interface). + + This is the least performant jump type in simulations. + + - Bounded `VariableRateJump`s require passing the keyword arguments `urate` and + `rateinterval`, corresponding to functions `urate(u, p, t)` and `rateinterval(u, p, t)`, + see below. These must calculate a time window over which the rate function is bounded by a + constant. Note that it is ok if the rate bound would be violated within the time interval + due to a change in `u` arising from another `ConstantRateJump`, `MassActionJump` or + *bounded* `VariableRateJump` being executed, as the chosen aggregator will then handle + recalculating the rate bound and interval. *However, if the bound could be violated within + the time interval due to a change in `u` arising from continuous dynamics such as a + coupled ODE, SDE, or a general `VariableRateJump`, bounds should not be given.* This + ensures the jump is classified as a general `VariableRateJump` and properly handled. One + can also optionally provide a lower bound function, `lrate(u, p, t)`, via the `lrate` + keyword argument. This can lead to increased performance. The validity of the lower bound + should hold under the same conditions and rate interval as `urate`. + + + Bounded `VariableRateJump`s can currently be used in the `Coevolve` aggregator, and + can therefore be efficiently simulated in pure-jump `DiscreteProblem`s using the + `SSAStepper` time-stepper. + + These can be substantially more performant than general `VariableRateJump`s without + the rate bound functions. Reemphasizing, the additional user provided functions leveraged by bounded -`VariableRateJumps`, `urate(u, p, t)`, `rateinterval(u, p, t)`, and the optional `lrate(u, -p, t)` require that -- For `s` in `[t, t + rateinterval(u, p, t)]`, we have that `lrate(u, p, t) <= rate(u, p, s) - <= urate(u, p, t)`. -- It is ok if these bounds would be violated during the time window due to another - `ConstantRateJump`, `MassActionJump` or bounded `VariableRateJump` occurring. However, - they must remain valid if `u` changes for any other reason (for example, due to - continuous dynamics like ODEs, SDEs, or general `VariableRateJump`s). +`VariableRateJumps`, `urate(u, p, t)`, `rateinterval(u, p, t)`, and the optional `lrate(u, p, t)` require that + + - For `s` in `[t, t + rateinterval(u, p, t)]`, we have that `lrate(u, p, t) <= rate(u, p, s) <= urate(u, p, t)`. + - It is ok if these bounds would be violated during the time window due to another + `ConstantRateJump`, `MassActionJump` or bounded `VariableRateJump` occurring. However, + they must remain valid if `u` changes for any other reason (for example, due to + continuous dynamics like ODEs, SDEs, or general `VariableRateJump`s). ## Fields $(FIELDS) ## Examples + Suppose `u[1]` gives the number of particles and `t*p[1]` the probability per time each particle can decay away. A corresponding `VariableRateJump` for this jump process is + ```julia -rate(u,p,t) = t*p[1]*u[1] +rate(u, p, t) = t*p[1]*u[1] affect!(integrator) = integrator.u[1] -= 1 vrj = VariableRateJump(rate, affect!) ``` To define a bounded `VariableRateJump` that can be used with supporting aggregators such as `Coevolve`, we must define bounds and a rate interval: + ```julia -rateinterval(u,p,t) = (1 / p[1]) * 2 -rate(u,p,t) = t * p[1] * u[1] +rateinterval(u, p, t) = (1 / p[1]) * 2 +rate(u, p, t) = t * p[1] * u[1] lrate(u, p, t) = rate(u, p, t) -urate(u,p,t) = rate(u, p, t + rateinterval(u,p,t)) +urate(u, p, t) = rate(u, p, t + rateinterval(u, p, t)) affect!(integrator) = integrator.u[1] -= 1 vrj = VariableRateJump(rate, affect!; lrate = lrate, urate = urate, - rateinterval = rateinterval) + rateinterval = rateinterval) ``` ## Notes -- When using an aggregator that supports bounded `VariableRateJump`s, `DiscreteProblem` can - be used. Otherwise, `ODEProblem` or `SDEProblem` must be used. -- **When not using aggregators that support bounded `VariableRateJump`s, or when there are - general `VariableRateJump`s, `integrator`s store an effective state type that wraps the - main state vector.** See [`ExtendedJumpArray`](@ref) for details on using this object. In - this case all `ConstantRateJump`, `VariableRateJump` and callback `affect!` functions - receive an integrator with `integrator.u` an [`ExtendedJumpArray`](@ref). -- Salis H., Kaznessis Y., Accurate hybrid stochastic simulation of a system of coupled - chemical or biochemical reactions, Journal of Chemical Physics, 122 (5), - DOI:10.1063/1.1835951 is used for calculating jump times with `VariableRateJump`s within - ODE/SDE integrators. + + - When using an aggregator that supports bounded `VariableRateJump`s, `DiscreteProblem` can + be used. Otherwise, `ODEProblem` or `SDEProblem` must be used. + - **When not using aggregators that support bounded `VariableRateJump`s, or when there are + general `VariableRateJump`s, `integrator`s store an effective state type that wraps the + main state vector.** See [`ExtendedJumpArray`](@ref) for details on using this object. In + this case all `ConstantRateJump`, `VariableRateJump` and callback `affect!` functions + receive an integrator with `integrator.u` an [`ExtendedJumpArray`](@ref). + - Salis H., Kaznessis Y., Accurate hybrid stochastic simulation of a system of coupled + chemical or biochemical reactions, Journal of Chemical Physics, 122 (5), + DOI:10.1063/1.1835951 is used for calculating jump times with `VariableRateJump`s within + ODE/SDE integrators. """ struct VariableRateJump{R, F, R2, R3, R4, I, T, T2} <: AbstractJump - """Function `rate(u,p,t)` that returns the jump's current rate given state - `u`, parameters `p` and time `t`.""" + """ + Function `rate(u,p,t)` that returns the jump's current rate given state + `u`, parameters `p` and time `t`. + """ rate::R - """Function `affect!(integrator)` that updates the state for one occurrence - of the jump given `integrator`.""" + """ + Function `affect!(integrator)` that updates the state for one occurrence + of the jump given `integrator`. + """ affect!::F - """Optional function `lrate(u, p, t)` that computes a lower bound on the rate in the + """ + Optional function `lrate(u, p, t)` that computes a lower bound on the rate in the interval `t` to `t + rateinterval(u, p, t)` at time `t` given state `u` and parameters `p`. This bound must rigorously hold during the time interval as long as another `ConstantRateJump`, `MassActionJump`, or *bounded* `VariableRateJump` has not been @@ -312,21 +352,25 @@ struct VariableRateJump{R, F, R2, R3, R4, I, T, T2} <: AbstractJump `Coevolve`, providing a lower-bound can lead to improved performance. """ lrate::R2 - """Optional function `urate(u, p, t)` for general `VariableRateJump`s, but is required + """ + Optional function `urate(u, p, t)` for general `VariableRateJump`s, but is required to define a bounded `VariableRateJump`, which can be used with supporting aggregators, - currently only `Coevolve`, and offers improved computational performance. Computes an + currently only `Coevolve`, and offers improved computational performance. Computes an upper bound for the rate in the interval `t` to `t + rateinterval(u, p, t)` at time `t` given state `u` and parameters `p`. This bound must rigorously hold during the time interval as long as another `ConstantRateJump`, `MassActionJump`, or *bounded* - `VariableRateJump` has not been sampled. """ + `VariableRateJump` has not been sampled. + """ urate::R3 - """Optional function `rateinterval(u, p, t)` for general `VariableRateJump`s, but is + """ + Optional function `rateinterval(u, p, t)` for general `VariableRateJump`s, but is required to define a bounded `VariableRateJump`, which can be used with supporting - aggregators, currently only `Coevolve`, and offers improved computational performance. + aggregators, currently only `Coevolve`, and offers improved computational performance. Computes the time interval from time `t` over which the `urate` and `lrate` bounds will hold, `t` to `t + rateinterval(u, p, t)`, given state `u` and parameters `p`. This bound must rigorously hold during the time interval as long as another `ConstantRateJump`, - `MassActionJump`, or *bounded* `VariableRateJump` has not been sampled. """ + `MassActionJump`, or *bounded* `VariableRateJump` has not been sampled. + """ rateinterval::R4 idxs::I rootfind::Bool @@ -380,13 +424,15 @@ Representation for encoding rates and multiple simultaneous jumps for use in τ- methods. ### Constructors -- `RegularJump(rate, c, numjumps; mark_dist = nothing)` + + - `RegularJump(rate, c, numjumps; mark_dist = nothing)` ## Fields $(FIELDS) ## Examples + ```julia function rate!(out, u, p, t) out[1] = (0.1 / 1000.0) * u[1] * u[2] @@ -417,9 +463,13 @@ struct RegularJump{iip, R, C, MD} saving the output in `du[i]`. """ c::C - """ Number of jumps in the system.""" + """ + Number of jumps in the system. + """ numjumps::Int - """ A distribution for marks. Not currently used or supported. """ + """ + A distribution for marks. Not currently used or supported. + """ mark_dist::MD function RegularJump{iip}(rate, c, numjumps::Int; mark_dist = nothing) where {iip} new{iip, typeof(rate), typeof(c), typeof(mark_dist)}(rate, c, numjumps, mark_dist) @@ -448,12 +498,14 @@ $(TYPEDEF) Optimized representation for `ConstantRateJump`s that can be represented in mass action form, offering improved performance within jump algorithms compared to `ConstantRateJump`. For detailed examples and usage information, see the -- [Main - Docs](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Defining-a-Mass-Action-Jump) -- [Tutorial](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) + + - [Main + Docs](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Defining-a-Mass-Action-Jump) + - [Tutorial](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) ### Constructors -- `MassActionJump(reactant_stoich, net_stoich; scale_rates = true, param_idxs = nothing)` + + - `MassActionJump(reactant_stoich, net_stoich; scale_rates = true, param_idxs = nothing)` Here `reactant_stoich` denotes the reactant stoichiometry for each reaction and `net_stoich` the net stoichiometry for each reaction. @@ -463,55 +515,57 @@ Here `reactant_stoich` denotes the reactant stoichiometry for each reaction and $(FIELDS) ## Keyword Arguments -- `scale_rates = true`, whether to rescale the reaction rate constants according - to the stoichiometry. -- `nocopy = false`, whether the `MassActionJump` can alias the `scaled_rates` and - `reactant_stoch` from the input. Note, if `scale_rates=true` this will - potentially modify both of these. -- `param_idxs = nothing`, indexes in the parameter vector, `JumpProblem.prob.p`, - that correspond to each reaction's rate; a scalar index is accepted for one reaction. -- `param_mapper = nothing`, an alternative to `param_idxs`: a callable - `(dest, maj, params)` that fills the final scaled working rates in `dest` and - returns `nothing`. It must not mutate the definition or input parameters. -- `rescale_rates_on_update = scale_rates`, the scaling policy used by the built-in - parameter mapper. Custom mappers implement their own scaling and should honor - this policy when mapping unscaled coefficients; pre-scaled output must not be - scaled again. + + - `scale_rates = true`, whether to rescale the reaction rate constants according + to the stoichiometry. + - `nocopy = false`, whether the `MassActionJump` can alias the `scaled_rates` and + `reactant_stoch` from the input. Note, if `scale_rates=true` this will + potentially modify both of these. + - `param_idxs = nothing`, indexes in the parameter vector, `JumpProblem.prob.p`, + that correspond to each reaction's rate; a scalar index is accepted for one reaction. + - `param_mapper = nothing`, an alternative to `param_idxs`: a callable + `(dest, maj, params)` that fills the final scaled working rates in `dest` and + returns `nothing`. It must not mutate the definition or input parameters. + - `rescale_rates_on_update = scale_rates`, the scaling policy used by the built-in + parameter mapper. Custom mappers implement their own scaling and should honor + this policy when mapping unscaled coefficients; pre-scaled output must not be + scaled again. See the tutorial and main docs for details. ## Examples + An SIR model with `S + I --> 2I` at rate β as the first reaction and `I --> R` at rate ν as the second reaction can be encoded by + ```julia -p = (β=1e-4, ν=.01) -u0 = [999, 1, 0] # (S,I,R) -tspan = (0.0, 250.0) +p = (β = 1e-4, ν = 0.01) +u0 = [999, 1, 0] # (S,I,R) +tspan = (0.0, 250.0) rateidxs = [1, 2] # i.e. [β,ν] reactant_stoich = [ - [1 => 1, 2 => 1], # 1*S and 1*I - [2 => 1] # 1*I + [1 => 1, 2 => 1], # 1*S and 1*I + [2 => 1] # 1*I ] net_stoich = [ - [1 => -1, 2 => 1], # -1*S and 1*I - [2 => -1, 3 => 1] # -1*I and 1*R + [1 => -1, 2 => 1], # -1*S and 1*I + [2 => -1, 3 => 1] # -1*I and 1*R ] -maj = MassActionJump(reactant_stoich, net_stoich; param_idxs=rateidxs) +maj = MassActionJump(reactant_stoich, net_stoich; param_idxs = rateidxs) prob = DiscreteProblem(u0, tspan, p) jprob = JumpProblem(prob, Direct(), maj) ``` ## Notes -- By default, fixed reaction rates are rescaled when constructing the `MassActionJump`; - the built-in parameter mapper rescales rates when filling each solver's working - buffer at initialization/reset, as explained in the [main - docs](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Defining-a-Mass-Action-Jump). - Disable this with the kwarg `scale_rates=false`. -- Also see the [main - docs](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Defining-a-Mass-Action-Jump) - for how to specify reactions with no products or no reactants. - + - By default, fixed reaction rates are rescaled when constructing the `MassActionJump`; + the built-in parameter mapper rescales rates when filling each solver's working + buffer at initialization/reset, as explained in the [main + docs](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Defining-a-Mass-Action-Jump). + Disable this with the kwarg `scale_rates=false`. + - Also see the [main + docs](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Defining-a-Mass-Action-Jump) + for how to specify reactions with no products or no reactants. """ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump """ @@ -522,19 +576,28 @@ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump solves. Prefer parameter indices or a custom mapper for rates that change. """ scaled_rates::T - """The reactant stoichiometry vectors.""" + """ + The reactant stoichiometry vectors. + """ reactant_stoch::S - """The net stoichiometry vectors.""" + """ + The net stoichiometry vectors. + """ net_stoch::U - """Parameter mapping functor to identify reaction rate constants with parameters in `p` vectors.""" + """ + Parameter mapping functor to identify reaction rate constants with parameters in `p` vectors. + """ param_mapper::V - """Whether the built-in mapper callable should apply stoichiometric scaling to rates.""" + """ + Whether the built-in mapper callable should apply stoichiometric scaling to rates. + """ rescale_rates_on_update::Bool function MassActionJump{T, S, U, V}(rates::T, rs_in::S, ns::U, pmapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool, - rescale_rates_on_update::Bool = scale_rates) where {T <: AbstractVector, S, U, V} + rescale_rates_on_update::Bool = scale_rates) where { + T <: AbstractVector, S, U, V} sr = nocopy ? rates : copy(rates) rs = nocopy ? rs_in : copy(rs_in) for i in eachindex(rs) @@ -610,7 +673,8 @@ function MassActionJump(rs, ns; param_idxs = nothing, param_mapper = nothing, if param_mapper === nothing (param_idxs === nothing) && error("If no parameter indices are given via param_idxs, an explicit parameter mapping must be passed in via param_mapper.") - pmapper = MassActionJumpParamMapper(param_idxs isa Integer ? [param_idxs] : param_idxs) + pmapper = MassActionJumpParamMapper(param_idxs isa Integer ? [param_idxs] : + param_idxs) else (param_idxs !== nothing) && error("Only one of param_idxs and param_mapper should be passed.") @@ -688,7 +752,9 @@ function (mapper::MassActionJumpParamMapper{U})(dest::AbstractVector, nothing end -to_collection(ratemap::MassActionJumpParamMapper) = MassActionJumpParamMapper(copy(ratemap.param_idxs)) +function to_collection(ratemap::MassActionJumpParamMapper) + MassActionJumpParamMapper(copy(ratemap.param_idxs)) +end function Base.merge!(pmap1::MassActionJumpParamMapper{U}, pmap2::MassActionJumpParamMapper{U}) where {U <: AbstractVector} @@ -727,22 +793,24 @@ Defines a collection of jumps that should collectively be included in a simulati $(FIELDS) ## Examples + Here we construct two jumps, store them in a `JumpSet`, and then simulate the resulting process. + ```julia using JumpProcesses, OrdinaryDiffEq -rate1(u,p,t) = p[1] +rate1(u, p, t) = p[1] affect1!(integrator) = (integrator.u[1] += 1) crj = ConstantRateJump(rate1, affect1!) -rate2(u,p,t) = (t/(1+t))*p[2]*u[1] +rate2(u, p, t) = (t/(1+t))*p[2]*u[1] affect2!(integrator) = (integrator.u[1] -= 1) vrj = VariableRateJump(rate2, affect2!) jset = JumpSet(crj, vrj) -f!(du,u,p,t) = (du .= 0) +f!(du, u, p, t) = (du .= 0) u0 = [0.0] p = (20.0, 2.0) tspan = (0.0, 200.0) @@ -752,13 +820,21 @@ sol = solve(jprob, Tsit5()) ``` """ struct JumpSet{T1, T2, T3, T4} <: AbstractJump - """Collection of [`VariableRateJump`](@ref)s""" + """ + Collection of [`VariableRateJump`](@ref)s + """ variable_jumps::T1 - """Collection of [`ConstantRateJump`](@ref)s""" + """ + Collection of [`ConstantRateJump`](@ref)s + """ constant_jumps::T2 - """Collection of [`RegularJump`](@ref)s""" + """ + Collection of [`RegularJump`](@ref)s + """ regular_jump::T3 - """Collection of [`MassActionJump`](@ref)s""" + """ + Collection of [`MassActionJump`](@ref)s + """ massaction_jump::T4 end function JumpSet(vj, cj, rj, maj::MassActionJump{S, T, U, V}) where {S <: Number, T, U, V} @@ -880,7 +956,8 @@ end # if given containers of rates and stoichiometry directly create a jump # copy arrays since majump_merge! will mutate them in-place -function setup_majump_to_merge(sr::T, rs::AbstractVector{S}, ns::AbstractVector{U}, pmapper, +function setup_majump_to_merge( + sr::T, rs::AbstractVector{S}, ns::AbstractVector{U}, pmapper, rescale_rates_on_update::Bool) where {T <: AbstractVector, S <: AbstractArray, U <: AbstractArray} MassActionJump(copy(sr), copy(rs), copy(ns), pmapper; @@ -1032,8 +1109,9 @@ function get_jump_bracket_fwrappers(u, p, t, jumps, agg) end end -make_bracket_fn(c::ConstantRateJump) = +function make_bracket_fn(c::ConstantRateJump) (ulow, uhigh, u, p, t) -> cjump_brackets(c, ulow, uhigh, u, p, t) +end function get_jump_lrate_fwrappers(u, p, t, jumps, agg) BoundWrapper = FunctionWrappers.FunctionWrapper{typeof(t), @@ -1063,8 +1141,10 @@ function get_jump_urate_fwrappers(u, p, t, jumps, agg) end end -make_lrate_fn(c::ConstantRateJump) = +function make_lrate_fn(c::ConstantRateJump) (ulow, uhigh, u, p, t) -> lower_rate_bound(c, ulow, uhigh, u, p, t) +end -make_urate_fn(c::ConstantRateJump) = +function make_urate_fn(c::ConstantRateJump) (ulow, uhigh, u, p, t) -> upper_rate_bound(c, ulow, uhigh, u, p, t) +end diff --git a/src/problem.jl b/src/problem.jl index 98861d59a..5f987983a 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -72,6 +72,7 @@ the DifferentialEquations.jl [docs](https://docs.sciml.ai/JumpProcesses/stable/) examples and commonly asked questions. !!! warning "Thread Safety" + `JumpProblem` contains mutable state (aggregator data, callbacks) and is **not thread-safe**. A single `JumpProblem` instance must not be solved concurrently from multiple threads or tasks without first creating independent copies via `deepcopy`. @@ -83,24 +84,42 @@ examples and commonly asked questions. state, and kernel ensembles have their own backend RNG rules. """ mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggregator}, J1, - J2, J3, J4, K} <: SciMLBase.AbstractJumpProblem{P, J} - """The type of problem to couple the jumps to. For a pure jump process use `DiscreteProblem`, to couple to ODEs, `ODEProblem`, etc.""" + J2, J3, J4, K} <: SciMLBase.AbstractJumpProblem{P, J} + """ + The type of problem to couple the jumps to. For a pure jump process use `DiscreteProblem`, to couple to ODEs, `ODEProblem`, etc. + """ prob::P - """The aggregator algorithm that determines the next jump times and types for `ConstantRateJump`s and `MassActionJump`s. Examples include `Direct`.""" + """ + The aggregator algorithm that determines the next jump times and types for `ConstantRateJump`s and `MassActionJump`s. Examples include `Direct`. + """ aggregator::A - """The underlying state data associated with the chosen aggregator.""" + """ + The underlying state data associated with the chosen aggregator. + """ discrete_jump_aggregation::J - """`CallBackSet` with the underlying `ConstantRate` and `VariableRate` jumps.""" + """ + `CallBackSet` with the underlying `ConstantRate` and `VariableRate` jumps. + """ jump_callback::C - """The `ConstantRateJump`s.""" + """ + The `ConstantRateJump`s. + """ constant_jumps::J1 - """The `VariableRateJump`s.""" + """ + The `VariableRateJump`s. + """ variable_jumps::J2 - """The `RegularJump`s.""" + """ + The `RegularJump`s. + """ regular_jump::J3 - """The `MassActionJump`s.""" + """ + The `MassActionJump`s. + """ massaction_jump::J4 - """kwargs to pass on to solve call.""" + """ + kwargs to pass on to solve call. + """ kwargs::K end function JumpProblem(p::P, a::A, dj::J, jc::C, cj::J1, vj::J2, rj::J3, mj::J4, @@ -142,7 +161,8 @@ function DiffEqBase.remake(jprob::JumpProblem; u0 = missing, p = missing, final_u0 = u0 else # Resolve symbolic maps if needed (handles [:X => 3.0], [sys.X => 3.0], Dict, etc.) - resolved_u0, _ = SciMLBase.updated_u0_p(prob, u0, missing; interpret_symbolicmap, use_defaults) + resolved_u0, _ = SciMLBase.updated_u0_p( + prob, u0, missing; interpret_symbolicmap, use_defaults) # Extract state values and wrap with fresh jump_u (resampled) # This handles both: @@ -157,9 +177,11 @@ function DiffEqBase.remake(jprob::JumpProblem; u0 = missing, p = missing, final_u0 = remake_extended_u0(prob, state_vals) end - newprob = DiffEqBase.remake(prob; u0 = final_u0, p, interpret_symbolicmap, use_defaults, kwargs...) + newprob = DiffEqBase.remake( + prob; u0 = final_u0, p, interpret_symbolicmap, use_defaults, kwargs...) else - newprob = DiffEqBase.remake(prob; u0, p, interpret_symbolicmap, use_defaults, kwargs...) + newprob = DiffEqBase.remake( + prob; u0, p, interpret_symbolicmap, use_defaults, kwargs...) end else @@ -237,7 +259,6 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS (false, true) : (true, true), spatial_system = nothing, hopping_constants = nothing, callback = nothing, tstops = nothing, use_vrj_bounds = true, kwargs...) - if haskey(kwargs, :rng) throw(ArgumentError("`rng` is no longer a keyword argument for `JumpProblem`. Pass `rng` to `solve` or `init` instead, e.g. `solve(jprob, SSAStepper(); rng = my_rng)`.")) end @@ -317,7 +338,6 @@ function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; (false, true) : (true, true), spatial_system = nothing, hopping_constants = nothing, callback = nothing, tstops = nothing, kwargs...) - if haskey(kwargs, :rng) throw(ArgumentError("`rng` is no longer a keyword argument for `JumpProblem`. Pass `rng` to `solve` or `init` instead, e.g. `solve(jprob, SSAStepper(); rng = my_rng)`.")) end @@ -356,7 +376,7 @@ end aggregator(jp::JumpProblem{iip, P, A}) where {iip, P, A} = A -@inline function extend_tstops!(tstops, jp::JumpProblem) +@inline function extend_tstops!(tstops, jp::JumpProblem) !(jp.jump_callback.discrete_callbacks isa Tuple{}) && push!(tstops, jp.jump_callback.discrete_callbacks[1].condition.next_jump_time) end diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index 5a5a9889a..f8674702f 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -115,14 +115,17 @@ $(FIELDS) ```julia using JumpProcesses -maj = MassActionJump([1.0, 1.0], [[1 => 1], [2 => 1]], [[1 => -1, 2 => 1], [1 => 1, 2 => -1]]) +maj = MassActionJump([1.0, 1.0], [[1 => 1], [2 => 1]], [ + [1 => -1, 2 => 1], [1 => 1, 2 => -1]]) prob = DiscreteProblem([100, 100], (0.0, 10.0)) jprob = JumpProblem(prob, PureLeaping(), maj) sol = solve(jprob, SimpleImplicitTauLeaping()) ``` """ struct SimpleImplicitTauLeaping{T <: AbstractFloat} <: SciMLBase.AbstractDEAlgorithm - """Error control parameter used when selecting `tau`.""" + """ + Error control parameter used when selecting `tau`. + """ epsilon::T end @@ -159,14 +162,17 @@ $(FIELDS) ```julia using JumpProcesses -maj = MassActionJump([1.0, 1.0], [[1 => 1], [2 => 1]], [[1 => -1, 2 => 1], [1 => 1, 2 => -1]]) +maj = MassActionJump([1.0, 1.0], [[1 => 1], [2 => 1]], [ + [1 => -1, 2 => 1], [1 => 1, 2 => -1]]) prob = DiscreteProblem([100, 100], (0.0, 10.0)) jprob = JumpProblem(prob, PureLeaping(), maj) sol = solve(jprob, SimpleTrapezoidalLeaping()) ``` """ struct SimpleTrapezoidalLeaping{T <: AbstractFloat} <: SciMLBase.AbstractDEAlgorithm - """Error control parameter used when selecting `tau`.""" + """ + Error control parameter used when selecting `tau`. + """ epsilon::T end @@ -209,7 +215,8 @@ $(FIELDS) ```julia using JumpProcesses -maj = MassActionJump([1.0, 1.0], [[1 => 1], [2 => 1]], [[1 => -1, 2 => 1], [1 => 1, 2 => -1]]) +maj = MassActionJump([1.0, 1.0], [[1 => 1], [2 => 1]], [ + [1 => -1, 2 => 1], [1 => 1, 2 => -1]]) prob = DiscreteProblem([100, 100], (0.0, 10.0)) jprob = JumpProblem(prob, PureLeaping(), maj) @@ -219,15 +226,25 @@ sol = solve(jprob, SimpleAdaptiveTauLeaping(implicit_alg = SimpleTrapezoidalLeap """ struct SimpleAdaptiveTauLeaping{T <: AbstractFloat, A <: SciMLBase.AbstractDEAlgorithm} <: SciMLBase.AbstractDEAlgorithm - """Error control parameter used when selecting `tau`.""" + """ + Error control parameter used when selecting `tau`. + """ epsilon::T - """The algorithm used for steps that are taken implicitly.""" + """ + The algorithm used for steps that are taken implicitly. + """ implicit_alg::A - """Whether to judge stiffness from the eigenvalues of the drift Jacobian.""" + """ + Whether to judge stiffness from the eigenvalues of the drift Jacobian. + """ eigenvalue_check::Bool - """Eigenvalue ratio above which the system counts as stiff.""" + """ + Eigenvalue ratio above which the system counts as stiff. + """ stiffness_ratio_threshold::T - """Factor relaxing `epsilon` when a step is taken implicitly.""" + """ + Factor relaxing `epsilon` when a step is taken implicitly. + """ implicit_epsilon_factor::T end @@ -259,9 +276,9 @@ function validate_pure_leaping_inputs( jump_prob::JumpProblem, alg::Union{ SimpleExplicitTauLeaping, SimpleImplicitTauLeaping, - SimpleTrapezoidalLeaping, SimpleAdaptiveTauLeaping, + SimpleTrapezoidalLeaping, SimpleAdaptiveTauLeaping } - ) +) 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(),...). \ @@ -283,8 +300,9 @@ OrdinaryDiffEq conventions. Endpoint saving is controlled purely by the returned `save_start`/`save_end` flags. When the user passes `nothing` for these, defaults are: -- No saveat or saveat is a Number: `true` for both. -- saveat is a collection: `true` if the corresponding endpoint is `in` the collection. + + - No saveat or saveat is a Number: `true` for both. + - saveat is a collection: `true` if the corresponding endpoint is `in` the collection. """ function _process_saveat(saveat, tspan, save_start, save_end) t0, tf = tspan @@ -293,7 +311,7 @@ function _process_saveat(saveat, tspan, save_start, save_end) _save_start = something(save_start, true) _save_end = something(save_end, true) elseif saveat isa Number - saveat_vec = collect(t0 + saveat:saveat:tf) + saveat_vec = collect((t0 + saveat):saveat:tf) if !isempty(saveat_vec) && last(saveat_vec) == tf pop!(saveat_vec) end @@ -667,7 +685,7 @@ function implicit_equation!(resid, u_new, params) for j in 1:numjumps for spec_idx in axes(nu, 1) resid[spec_idx] -= nu[spec_idx, j] * half * - (rate_new[j] + rate_current[j]) * tau + (rate_new[j] + rate_current[j]) * tau end end end @@ -677,7 +695,7 @@ end function solve_implicit( u_current, rate_new, rate_current, nu, p, t, tau, rate, numjumps, alg - ) +) u_guess = convert(Vector{float(eltype(u_current))}, u_current) params = (; u_current, rate_new, rate_current, nu, p, t, tau, rate, numjumps, alg) prob = NonlinearProblem(implicit_equation!, u_guess, params) @@ -693,7 +711,7 @@ function simple_implicit_tau_leaping_loop!( rate, nu, hor, max_hor, max_stoich, numjumps, epsilon, dtmin, saveat_times, usave, tsave, du, counts, rate_cache, rate_current, maj, save_end - ) +) save_idx = 1 # Upper bound carried across iterations. Unlike the explicit loop, whose @@ -714,7 +732,7 @@ function simple_implicit_tau_leaping_loop!( ) tau = min(tau, tau_cap, t_end - t_current) if !isempty(saveat_times) && save_idx <= length(saveat_times) && - t_current + tau > saveat_times[save_idx] + t_current + tau > saveat_times[save_idx] tau = saveat_times[save_idx] - t_current end @@ -726,7 +744,7 @@ function simple_implicit_tau_leaping_loop!( if tau <= dtmin error( "$(nameof(typeof(alg))) failed to converge at t = $t_current " * - "with the smallest permitted step dtmin = $dtmin." + "with the smallest permitted step dtmin = $dtmin." ) end tau_cap = tau / 2 @@ -738,7 +756,7 @@ function simple_implicit_tau_leaping_loop!( for j in eachindex(counts) scaled = rate_cache[j] * tau counts[j] = scaled <= zero(scaled) ? zero(eltype(counts)) : - pois_rand(rng, scaled) + pois_rand(rng, scaled) end du .= 0 @@ -757,7 +775,7 @@ function simple_implicit_tau_leaping_loop!( t_new = t_current + tau if isempty(saveat_times) || - (save_idx <= length(saveat_times) && t_new >= saveat_times[save_idx]) + (save_idx <= length(saveat_times) && t_new >= saveat_times[save_idx]) push!(usave, copy(u_new)) push!(tsave, t_new) if !isempty(saveat_times) && t_new >= saveat_times[save_idx] @@ -783,7 +801,7 @@ function DiffEqBase.solve( seed = nothing, rng = nothing, dtmin = nothing, saveat = nothing, save_start = nothing, save_end = nothing - ) +) validate_pure_leaping_inputs(jump_prob, alg) || error("$(nameof(typeof(alg))) can only be used with PureLeaping JumpProblem with a MassActionJump.") @@ -923,7 +941,7 @@ function simple_adaptive_tau_leaping_loop!( dtmin, saveat_times, usave, tsave, du, counts, rate_cache, rate_effective, rate_current, maj, implicit_alg, eigenvalue_check, stiffness_ratio_threshold, implicit_epsilon_factor, save_end - ) +) save_idx = 1 # Upper bound carried across iterations. An explicit retry redraws the @@ -959,7 +977,7 @@ function simple_adaptive_tau_leaping_loop!( end tau = min(tau, tau_cap, t_end - t_current) if !isempty(saveat_times) && save_idx <= length(saveat_times) && - t_current + tau > saveat_times[save_idx] + t_current + tau > saveat_times[save_idx] tau = saveat_times[save_idx] - t_current end @@ -974,7 +992,7 @@ function simple_adaptive_tau_leaping_loop!( if !converged tau <= dtmin && error("SimpleAdaptiveTauLeaping failed to converge at t = $t_current " * - "with the smallest permitted step dtmin = $dtmin.") + "with the smallest permitted step dtmin = $dtmin.") tau_cap = tau / 2 continue end @@ -1008,7 +1026,7 @@ function simple_adaptive_tau_leaping_loop!( # Save state if at a saveat time or if saveat is empty if isempty(saveat_times) || - (save_idx <= length(saveat_times) && t_new >= saveat_times[save_idx]) + (save_idx <= length(saveat_times) && t_new >= saveat_times[save_idx]) push!(usave, copy(u_new)) push!(tsave, t_new) if !isempty(saveat_times) && t_new >= saveat_times[save_idx] @@ -1034,7 +1052,7 @@ function DiffEqBase.solve( seed = nothing, rng = nothing, dtmin = nothing, saveat = nothing, save_start = nothing, save_end = nothing - ) +) validate_pure_leaping_inputs(jump_prob, alg) || error("SimpleAdaptiveTauLeaping can only be used with PureLeaping JumpProblem with a MassActionJump.") diff --git a/src/solve.jl b/src/solve.jl index b249d7704..77fdcdaa6 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -15,8 +15,9 @@ function resolve_rng(rng, seed) end end -SciMLBase.supports_solve_rng(jprob::JumpProblem, alg::SciMLBase.AbstractDEAlgorithm) = +function SciMLBase.supports_solve_rng(jprob::JumpProblem, alg::SciMLBase.AbstractDEAlgorithm) SciMLBase.supports_solve_rng(jprob.prob, alg) +end function SciMLBase.__solve(jump_prob::JumpProblem{IIP, P}, alg::SciMLBase.AbstractDEAlgorithm; @@ -41,8 +42,9 @@ function SciMLBase.__solve(jump_prob::JumpProblem{IIP, P}, integrator.sol end -SciMLBase.supports_solve_rng(jprob::JumpProblem, ::Nothing) = +function SciMLBase.supports_solve_rng(jprob::JumpProblem, ::Nothing) jprob.prob isa SciMLBase.DiscreteProblem +end # if passed a JumpProblem over a DiscreteProblem, and no aggregator is selected use # SSAStepper @@ -56,7 +58,8 @@ function SciMLBase.__solve(jump_prob::JumpProblem; kwargs...) end function SciMLBase.__init(_jump_prob::JumpProblem{IIP, P}, - alg::SciMLBase.AbstractDEAlgorithm; merge_callbacks = true, kwargs...) where {IIP, P} + alg::SciMLBase.AbstractDEAlgorithm; merge_callbacks = true, kwargs...) where { + IIP, P} # Merge jump_prob.kwargs with passed kwargs kwargs = DiffEqBase.merge_problem_kwargs(_jump_prob; merge_callbacks, kwargs...) @@ -67,7 +70,6 @@ function __jump_init(_jump_prob::JumpProblem{IIP, P}, alg; callback = nothing, seed = nothing, rng = nothing, alias_jump = Threads.threadid() == 1, kwargs...) where {IIP, P} - _rng = resolve_rng(rng, seed) if alias_jump diff --git a/src/spatial/flatten.jl b/src/spatial/flatten.jl index 2b0dbc672..e412fe5a0 100644 --- a/src/spatial/flatten.jl +++ b/src/spatial/flatten.jl @@ -51,7 +51,8 @@ function flatten(netstoch::AbstractArray, reactstoch::AbstractArray, hop_constants = Matrix{Vector{F}}(undef, size(hopping_constants)) for ci in CartesianIndices(hop_constants) (species, site) = Tuple(ci) - hop_constants[ci] = hopping_constants[species, site] * ones(outdegree(spatial_system, site)) + hop_constants[ci] = hopping_constants[species, site] * + ones(outdegree(spatial_system, site)) end flatten(netstoch, reactstoch, rx_rates, spatial_system, u0, tspan, hop_constants; kwargs...) diff --git a/src/spatial/hop_rates.jl b/src/spatial/hop_rates.jl index d986d23c5..8154c8567 100644 --- a/src/spatial/hop_rates.jl +++ b/src/spatial/hop_rates.jl @@ -123,13 +123,19 @@ end ############## hopping rates of form D_s ################ struct HopRatesGraphDs{F} <: AbstractHopRates - "hopping_constants[i] is the hop constant of species i" + """ + hopping_constants[i] is the hop constant of species i + """ hopping_constants::Vector{F} - "rates[i,j] is total hopping rate of species i at site j" + """ + rates[i,j] is total hopping rate of species i at site j + """ rates::Matrix{F} - "sum_rates[j] is the sum of hopping rates at site j" + """ + sum_rates[j] is the sum of hopping rates at site j + """ sum_rates::Vector{F} end @@ -157,18 +163,25 @@ end return hopping rate of species at site """ function evalhoprate(hop_rates::HopRatesGraphDs, u, species, site, spatial_system) - @inbounds u[species, site] * hop_rates.hopping_constants[species] * outdegree(spatial_system, site) + @inbounds u[species, site] * hop_rates.hopping_constants[species] * + outdegree(spatial_system, site) end ############## hopping rates of form D_{s,i} ################ struct HopRatesGraphDsi{F} <: AbstractHopRates - "hopping_constants[i,j] is the hop constant of species i at site j" + """ + hopping_constants[i,j] is the hop constant of species i at site j + """ hopping_constants::Matrix{F} - "rates[i,j] is total hopping rate of species i at site j" + """ + rates[i,j] is total hopping rate of species i at site j + """ rates::Matrix{F} - "sum_rates[j] is the sum of hopping rates at site j" + """ + sum_rates[j] is the sum of hopping rates at site j + """ sum_rates::Vector{F} end @@ -196,18 +209,25 @@ end return hopping rate of species at site """ function evalhoprate(hop_rates::HopRatesGraphDsi, u, species, site, spatial_system) - @inbounds u[species, site] * hop_rates.hopping_constants[species, site] * outdegree(spatial_system, site) + @inbounds u[species, site] * hop_rates.hopping_constants[species, site] * + outdegree(spatial_system, site) end ############## hopping rates of form D_{s,i,j} ################ struct HopRatesGraphDsij{F} <: AbstractHopRates - "hop_const_cumulative_sums[s,i] is the vector of cumulative sums of hopping constants of species s at site i" + """ + hop_const_cumulative_sums[s,i] is the vector of cumulative sums of hopping constants of species s at site i + """ hop_const_cumulative_sums::Matrix{Vector{F}} - "rates[s,i] is the total hopping rate of species s at site i" + """ + rates[s,i] is the total hopping rate of species s at site i + """ rates::Matrix{F} - "sum_rates[i] is the total hopping rate out of site i" + """ + sum_rates[i] is the total hopping rate out of site i + """ sum_rates::Vector{F} end @@ -246,13 +266,19 @@ end Analogue of HopRatesGraphDsij, optimized for CartesianGrid """ struct HopRatesGridDsij{F} <: AbstractHopRates - "hop_const_cumulative_sums[:,s,i] is the vector of cumulative sums of hopping constants of species s at site i. Out-of-bounds neighbors are treated as having zero propensity." + """ + hop_const_cumulative_sums[:,s,i] is the vector of cumulative sums of hopping constants of species s at site i. Out-of-bounds neighbors are treated as having zero propensity. + """ hop_const_cumulative_sums::Array{F, 3} - "rates[s,i] is the total hopping rate of species s at site i" + """ + rates[s,i] is the total hopping rate of species s at site i + """ rates::Matrix{F} - "sum_rates[i] is the total hopping rate out of site i" + """ + sum_rates[i] is the total hopping rate out of site i + """ sum_rates::Vector{F} end @@ -299,16 +325,24 @@ end ############## hopping rates of form D_s * L_{i,j} ################ struct HopRatesGraphDsLij{F} <: AbstractHopRates - "hopping constants of species -- D_s" + """ + hopping constants of species -- D_s + """ species_hop_constants::Vector{F} - "nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors" + """ + nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors + """ hop_const_cumulative_sums::Vector{Vector{F}} - "rates[s,i] is the total hopping rate of species s at site i" + """ + rates[s,i] is the total hopping rate of species s at site i + """ rates::Matrix{F} - "sum_rates[i] is the total hopping rate out of site i" + """ + sum_rates[i] is the total hopping rate out of site i + """ sum_rates::Vector{F} end @@ -342,21 +376,30 @@ function sample_target_site(hop_rates::HopRatesGraphDsLij, site, species, rng, end function evalhoprate(hop_rates::HopRatesGraphDsLij, u, species, site, spatial_system) - @inbounds u[species, site] * hop_rates.species_hop_constants[species] * hop_rates.hop_const_cumulative_sums[site][end] + @inbounds u[species, site] * hop_rates.species_hop_constants[species] * + hop_rates.hop_const_cumulative_sums[site][end] end ############## hopping rates of form D_s * L_{i,j} optimized for cartesian grid ################ struct HopRatesGridDsLij{F} <: AbstractHopRates - "hopping constants of species -- D_s" + """ + hopping constants of species -- D_s + """ species_hop_constants::Vector{F} - "nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors" + """ + nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors + """ hop_const_cumulative_sums::Matrix{F} - "rates[s,i] is the total hopping rate of species s at site i" + """ + rates[s,i] is the total hopping rate of species s at site i + """ rates::Matrix{F} - "sum_rates[i] is the total hopping rate out of site i" + """ + sum_rates[i] is the total hopping rate out of site i + """ sum_rates::Vector{F} end @@ -395,21 +438,30 @@ function sample_target_site(hop_rates::HopRatesGridDsLij, site, species, rng, gr end function evalhoprate(hop_rates::HopRatesGridDsLij, u, species, site, grid) - @inbounds u[species, site] * hop_rates.species_hop_constants[species] * hop_rates.hop_const_cumulative_sums[end, site] + @inbounds u[species, site] * hop_rates.species_hop_constants[species] * + hop_rates.hop_const_cumulative_sums[end, site] end ############## hopping rates of form D_{s,i} * L_{i,j} ################ struct HopRatesGraphDsiLij{F} <: AbstractHopRates - "hopping constants of species -- D_s" + """ + hopping constants of species -- D_s + """ species_hop_constants::Matrix{F} - "nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors" + """ + nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors + """ hop_const_cumulative_sums::Vector{Vector{F}} - "rates[s,i] is the total hopping rate of species s at site i" + """ + rates[s,i] is the total hopping rate of species s at site i + """ rates::Matrix{F} - "sum_rates[i] is the total hopping rate out of site i" + """ + sum_rates[i] is the total hopping rate out of site i + """ sum_rates::Vector{F} end @@ -437,21 +489,30 @@ function sample_target_site(hop_rates::HopRatesGraphDsiLij, site, species, rng, end function evalhoprate(hop_rates::HopRatesGraphDsiLij, u, species, site, spatial_system) - @inbounds u[species, site] * hop_rates.species_hop_constants[species, site] * hop_rates.hop_const_cumulative_sums[site][end] + @inbounds u[species, site] * hop_rates.species_hop_constants[species, site] * + hop_rates.hop_const_cumulative_sums[site][end] end ############## hopping rates of form D_{s,i} * L_{i,j} optimized for cartesian grid ################ struct HopRatesGridDsiLij{F} <: AbstractHopRates - "hopping constants of species -- D_s" + """ + hopping constants of species -- D_s + """ species_hop_constants::Matrix{F} - "nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors" + """ + nbs_cumulative[i] is the vector of cumulative sums of hopping constants from site i to its neighbors + """ hop_const_cumulative_sums::Matrix{F} - "rates[s,i] is the total hopping rate of species s at site i" + """ + rates[s,i] is the total hopping rate of species s at site i + """ rates::Matrix{F} - "sum_rates[i] is the total hopping rate out of site i" + """ + sum_rates[i] is the total hopping rate out of site i + """ sum_rates::Vector{F} end @@ -492,5 +553,6 @@ function sample_target_site(hop_rates::HopRatesGridDsiLij, site, species, rng, g end function evalhoprate(hop_rates::HopRatesGridDsiLij, u, species, site, grid) - @inbounds u[species, site] * hop_rates.species_hop_constants[species, site] * hop_rates.hop_const_cumulative_sums[end, site] + @inbounds u[species, site] * hop_rates.species_hop_constants[species, site] * + hop_rates.hop_const_cumulative_sums[end, site] end diff --git a/src/spatial/reaction_rates.jl b/src/spatial/reaction_rates.jl index c06858519..b4854f44c 100644 --- a/src/spatial/reaction_rates.jl +++ b/src/spatial/reaction_rates.jl @@ -4,16 +4,24 @@ A file with structs and functions for sampling reactions and updating reaction r ### spatial rx rates ### struct RxRates{F, M} - "rx_rates[i,j] is rate of reaction i at site j" + """ + rx_rates[i,j] is rate of reaction i at site j + """ rates::Matrix{F} - "rx_rates_sum[j] is sum of reaction rates at site j" + """ + rx_rates_sum[j] is sum of reaction rates at site j + """ sum_rates::Vector{F} - "AbstractMassActionJump" + """ + AbstractMassActionJump + """ ma_jumps::M - "working copy of scaled mass action jump rates" + """ + working copy of scaled mass action jump rates + """ maj_rates::Vector{F} end @@ -95,9 +103,11 @@ function Base.show(io::IO, ::MIME"text/plain", rx_rates::RxRates) println(io, "RxRates with $num_rxs reactions and $num_sites sites") end -function eval_massaction_rate(u, rx, ma_jumps::M, site, maj_rates) where {M <: SpatialMassActionJump} +function eval_massaction_rate( + u, rx, ma_jumps::M, site, maj_rates) where {M <: SpatialMassActionJump} evalrxrate(u, rx, ma_jumps, site) end -function eval_massaction_rate(u, rx, ma_jumps::M, site, maj_rates) where {M <: MassActionJump} +function eval_massaction_rate( + u, rx, ma_jumps::M, site, maj_rates) where {M <: MassActionJump} evalrxrate((@view u[:, site]), rx, ma_jumps, maj_rates) end diff --git a/src/spatial/topology.jl b/src/spatial/topology.jl index c8ce986a1..ef21859cc 100644 --- a/src/spatial/topology.jl +++ b/src/spatial/topology.jl @@ -214,15 +214,21 @@ collect(neighbors(grid, 1)) == [2, 3] ``` """ struct CartesianGridRej{N, T} <: AbstractCartesianGrid - "dimensions (side lengths) of the grid" + """ + dimensions (side lengths) of the grid + """ dims::NTuple{N, Int} - "number of neighbor for each site" + """ + number of neighbor for each site + """ nums_neighbors::Vector{Int8} CI::CartesianIndices{N, T} LI::LinearIndices{N, T} - "offsets, e.g. [-1, 1] for 1D" + """ + offsets, e.g. [-1, 1] for 1D + """ offsets::Vector{CartesianIndex{N}} end diff --git a/src/spatial/utils.jl b/src/spatial/utils.jl index 3be2b7378..d6cdb51ff 100644 --- a/src/spatial/utils.jl +++ b/src/spatial/utils.jl @@ -6,13 +6,19 @@ A file with helper functions for spatial simulations stores info for a spatial jump """ struct SpatialJump{J} - "source location" + """ + source location + """ src::J - "index of jump as a hop or reaction" + """ + index of jump as a hop or reaction + """ jidx::Int - "destination location, equal to src for within-site reactions" + """ + destination location, equal to src for within-site reactions + """ dst::J end diff --git a/test/constant_rate.jl b/test/constant_rate.jl index 676014d8b..d8052fc65 100644 --- a/test/constant_rate.jl +++ b/test/constant_rate.jl @@ -1,4 +1,5 @@ -using JumpProcesses, DiffEqBase, OrdinaryDiffEq, OrdinaryDiffEqFunctionMap, SciMLBase, Statistics +using JumpProcesses, DiffEqBase, OrdinaryDiffEq, OrdinaryDiffEqFunctionMap, SciMLBase, + Statistics using Test using StableRNGs rng = StableRNG(12345) diff --git a/test/degenerate_rx_cases.jl b/test/degenerate_rx_cases.jl index 2539d27ea..c0eb89274 100644 --- a/test/degenerate_rx_cases.jl +++ b/test/degenerate_rx_cases.jl @@ -97,7 +97,8 @@ dep_graph = [ ] spec_to_dep_jumps = [[2]] jump_to_dep_specs = [[1], [1]] -namedpars = (dep_graph = dep_graph, vartojumps_map = spec_to_dep_jumps, jumptovars_map = jump_to_dep_specs) +namedpars = (dep_graph = dep_graph, vartojumps_map = spec_to_dep_jumps, + jumptovars_map = jump_to_dep_specs) for method in methods local jump_prob = JumpProblem(prob, method, jump, jump2; namedpars...) @@ -108,7 +109,8 @@ for method in methods end if doprint - println("Mix of constant and mass action jumps, method = ", typeof(method), ", sol[end] = ", sol[end, end]) + println("Mix of constant and mass action jumps, method = ", + typeof(method), ", sol[end] = ", sol[end, end]) end @test sol[end, end] > 200 end diff --git a/test/ensemble_problems.jl b/test/ensemble_problems.jl index 70b8d2c0e..d0ad24fff 100644 --- a/test/ensemble_problems.jl +++ b/test/ensemble_problems.jl @@ -189,9 +189,10 @@ end # Check jump_u thresholds directly via init (callback sets them during initialization) rng = StableRNG(12345) thresholds = [begin - integrator = init(jprob, Tsit5(); rng) - integrator.u.jump_u[1] - end for _ in 1:3] + integrator = init(jprob, Tsit5(); rng) + integrator.u.jump_u[1] + end + for _ in 1:3] @test allunique(thresholds) # From a full ensemble solve, check both first event times and the diff --git a/test/explicit_tau_kernel_tests.jl b/test/explicit_tau_kernel_tests.jl index 1bb9e6df4..ed8e641a1 100644 --- a/test/explicit_tau_kernel_tests.jl +++ b/test/explicit_tau_kernel_tests.jl @@ -27,6 +27,7 @@ function run_explicit_tau_kernel_tests(backend, nsims) ss = serial_solve(jp, nsims; kwargs...) @test sk.u[1].t == ss.u[1].t for k in eachindex(sk.u[1].t), s in 1:nspec + mk = mean(sk.u[i].u[k][s] for i in 1:nsims) ms = mean(ss.u[i].u[k][s] for i in 1:nsims) @test isapprox(mk, ms; rtol, atol) diff --git a/test/hawkes_test.jl b/test/hawkes_test.jl index 4ffdf0ed9..a468581ca 100644 --- a/test/hawkes_test.jl +++ b/test/hawkes_test.jl @@ -156,7 +156,8 @@ let alg = Coevolve() for vr_aggregator in (VR_FRM(), VR_Direct(), VR_DirectFW()) oprob = ODEProblem(f!, u0, tspan, p) jumps = hawkes_jump(u0, g, h) - jprob = JumpProblem(oprob, alg, jumps...; vr_aggregator, dep_graph = g, use_vrj_bounds = false) + jprob = JumpProblem( + oprob, alg, jumps...; vr_aggregator, dep_graph = g, use_vrj_bounds = false) @test length(jprob.variable_jumps) == 1 sols = Vector{ODESolution}(undef, Nsims) for n in 1:Nsims diff --git a/test/kernel_contract_tests.jl b/test/kernel_contract_tests.jl index ada4d5f15..170106187 100644 --- a/test/kernel_contract_tests.jl +++ b/test/kernel_contract_tests.jl @@ -43,9 +43,9 @@ function run_kernel_rng_tests(backend, alg, jump_prob; kwargs...) # A single trajectory still delegates all RNG controls to EnsembleSerial, # including the ensemble layer's precedence when both rng and seed occur. for rng_options in (() -> (; seed = 123), - () -> (; rng = Xoshiro(17)), - () -> (; rng = Xoshiro(17), seed = 123), - () -> (; seed = 123, rng_func = ctx -> Xoshiro(17))) + () -> (; rng = Xoshiro(17)), + () -> (; rng = Xoshiro(17), seed = 123), + () -> (; seed = 123, rng_func = ctx -> Xoshiro(17))) kernel = solve(ensemble_prob, alg, EnsembleGPUKernel(backend); trajectories = 1, kwargs..., rng_options()...) serial = solve(ensemble_prob, alg, EnsembleSerial(); diff --git a/test/massaction_state.jl b/test/massaction_state.jl index 83ad47d13..ea4a7e489 100644 --- a/test/massaction_state.jl +++ b/test/massaction_state.jl @@ -63,7 +63,8 @@ end @test jprob.massaction_jump.scaled_rates == [0.0, 0.0, 0.0, 0.0, 2.0, 2.0] @test jprob.prob.u0 == [6, 0, 9, 0] init(jprob, SSAStepper(); seed = 12345, alias_jump = true) - @test jprob.discrete_jump_aggregation.maj_rates == jprob.massaction_jump.scaled_rates + @test jprob.discrete_jump_aggregation.maj_rates == + jprob.massaction_jump.scaled_rates if parameterized @test maj.scaled_rates === nothing else @@ -74,7 +75,8 @@ end @testset "Mass-action merges preserve source arrays and parameter indices" begin @testset "$parameterized, $collection" for parameterized in (false, true), - collection in (false, true) + collection in (false, true) + rs1, rs2 = [[1 => 3]], [[2 => 1]] ns1, ns2 = [[1 => -3, 2 => 1]], [[2 => -1, 1 => 3]] idxs1, idxs2 = [1], [2] @@ -137,6 +139,7 @@ end maj = MassActionJump([[1 => 3]], [[1 => -3]]; param_idxs = 1) prob = DiscreteProblem([6], (0.0, 1.0), [6.0]) for algorithm in (Direct(), PureLeaping()), value in (false, true) + @test_throws ArgumentError JumpProblem(prob, algorithm, maj; scale_rates = value) @test_throws r"scale_rates.*no longer.*MassActionJump" JumpProblem(prob, algorithm, maj; scale_rates = value) diff --git a/test/qa.jl b/test/qa.jl index bad04d944..d923b7ab6 100644 --- a/test/qa.jl +++ b/test/qa.jl @@ -13,7 +13,7 @@ const REEXPORTED_API = ( :ReturnCode, :SDEFunction, :SDEProblem, :VectorContinuousCallback, :add_saveat!, :add_tstop!, :derivative_discontinuity!, :init, :neighbors, :outdegree, :reinit!, :remake, :savevalues!, :set_proposed_dt!, :set_t!, :set_u!, :solve, :solve!, :step!, - :successful_retcode, :terminate!, :u_modified!, + :successful_retcode, :terminate!, :u_modified! ) # The ExplicitImports ignore-lists below are names owned by other packages whose @@ -28,33 +28,33 @@ run_qa( # Names not (yet) declared public in their owner package's released API. all_qualified_accesses_are_public = (; ignore = ( - # Base / Base.Broadcast / Base.FastMath internals - Symbol("@pure"), :BroadcastStyle, :Broadcasted, :Cartesian, - :DefaultArrayStyle, :FastMath, :Unknown, :result_style, :sqrt_fast, - # SciMLBase non-public - :ConstantInterpolation, :DISCRETE_INPLACE_DEFAULT, - :__init, :__solve, :get_colorizers, :isdenseplot, - :parameterless_type, :plottable_indices, :save_discretes_if_enabled!, - :save_final_discretes!, :solution_new_retcode, :unwrapped_f, - :updated_u0_p, - # DiffEqBase non-public - :Stats, - # ForwardDiff: Dual is the AD number type used in - # ext/JumpProcessesForwardDiffExt.jl. It is exported but not - # declared `public` in ForwardDiff's released API. - :Dual, - # LinearAlgebra non-public - :AbstractQ, :AdjointQ, :QRPackedQ, - # FunctionWrappers non-public - :FunctionWrapper, - ), + # Base / Base.Broadcast / Base.FastMath internals + Symbol("@pure"), :BroadcastStyle, :Broadcasted, :Cartesian, + :DefaultArrayStyle, :FastMath, :Unknown, :result_style, :sqrt_fast, + # SciMLBase non-public + :ConstantInterpolation, :DISCRETE_INPLACE_DEFAULT, + :__init, :__solve, :get_colorizers, :isdenseplot, + :parameterless_type, :plottable_indices, :save_discretes_if_enabled!, + :save_final_discretes!, :solution_new_retcode, :unwrapped_f, + :updated_u0_p, + # DiffEqBase non-public + :Stats, + # ForwardDiff: Dual is the AD number type used in + # ext/JumpProcessesForwardDiffExt.jl. It is exported but not + # declared `public` in ForwardDiff's released API. + :Dual, + # LinearAlgebra non-public + :AbstractQ, :AdjointQ, :QRPackedQ, + # FunctionWrappers non-public + :FunctionWrapper + ), ), all_explicit_imports_are_public = (; ignore = ( - :add_fast, # Base.FastMath non-public - :gauss_points, :gauss_weights, # DiffEqCallbacks non-public - :plot_indices, # SciMLBase non-public - ), + :add_fast, # Base.FastMath non-public + :gauss_points, :gauss_weights, # DiffEqCallbacks non-public + :plot_indices # SciMLBase non-public ), + ) ) ) diff --git a/test/regular_jumps.jl b/test/regular_jumps.jl index eaeef4dc9..faba3a5c9 100644 --- a/test/regular_jumps.jl +++ b/test/regular_jumps.jl @@ -10,9 +10,9 @@ rng = StableRNG(12345) regular_jump = RegularJump(rate!, change!, 1) massaction_jump = MassActionJump([0.1], [[1 => 1]], [[1 => -1]]) for alg in ( - SimpleTauLeaping(), SimpleExplicitTauLeaping(), - SimpleImplicitTauLeaping(), SimpleTrapezoidalLeaping(), SimpleAdaptiveTauLeaping(), - ) + SimpleTauLeaping(), SimpleExplicitTauLeaping(), + SimpleImplicitTauLeaping(), SimpleTrapezoidalLeaping(), SimpleAdaptiveTauLeaping() + ) @testset "$(nameof(typeof(alg)))" begin jump = alg isa SimpleTauLeaping ? regular_jump : massaction_jump jp = JumpProblem(prob, PureLeaping(), jump) @@ -33,11 +33,11 @@ end ) massaction_jump = MassActionJump([20.0], [Pair{Int, Int}[]], [[1 => 1]]) for alg in ( - SimpleTauLeaping(), SimpleExplicitTauLeaping(), - SimpleImplicitTauLeaping(), SimpleTrapezoidalLeaping(), - SimpleAdaptiveTauLeaping(), - SimpleAdaptiveTauLeaping(implicit_alg = SimpleTrapezoidalLeaping()), - ) + SimpleTauLeaping(), SimpleExplicitTauLeaping(), + SimpleImplicitTauLeaping(), SimpleTrapezoidalLeaping(), + SimpleAdaptiveTauLeaping(), + SimpleAdaptiveTauLeaping(implicit_alg = SimpleTrapezoidalLeaping()) + ) @testset "$alg" begin jump = alg isa SimpleTauLeaping ? regular_jump : massaction_jump jp = JumpProblem(prob, PureLeaping(), jump) @@ -87,10 +87,10 @@ end custom_maj = MassActionJump(reactants, net; param_mapper = prescaled_mapper, scale_rates = false) for alg in ( - SimpleExplicitTauLeaping(), SimpleImplicitTauLeaping(), - SimpleTrapezoidalLeaping(), SimpleAdaptiveTauLeaping(), - SimpleAdaptiveTauLeaping(implicit_alg = SimpleTrapezoidalLeaping()), - ) + SimpleExplicitTauLeaping(), SimpleImplicitTauLeaping(), + SimpleTrapezoidalLeaping(), SimpleAdaptiveTauLeaping(), + SimpleAdaptiveTauLeaping(implicit_alg = SimpleTrapezoidalLeaping()) + ) @testset "$alg" begin prob = DiscreteProblem([30.0], (0.0, 2.0), copy(original_rates)) jp = JumpProblem(prob, PureLeaping(), maj) @@ -137,9 +137,9 @@ end rate = JumpProcesses.massaction_rate(maj, rates, 1) nu = reshape([-1.0], 1, 1) for (alg, expected) in ( - (SimpleImplicitTauLeaping(), 100 / 1.2), - (SimpleTrapezoidalLeaping(), 100 * 0.9 / 1.1), - ) + (SimpleImplicitTauLeaping(), 100 / 1.2), + (SimpleTrapezoidalLeaping(), 100 * 0.9 / 1.1) + ) predicted, converged = JumpProcesses.solve_implicit( [100], zeros(1), zeros(1), nu, p, 0.0, 0.5, rate, 1, alg) @test converged @@ -190,7 +190,8 @@ end 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[1] += 1; nothing) - jumps = (ConstantRateJump(rate1, affect1!), ConstantRateJump(rate2, affect2!), ConstantRateJump(rate3, affect3!)) + jumps = (ConstantRateJump(rate1, affect1!), ConstantRateJump(rate2, affect2!), + ConstantRateJump(rate3, affect3!)) u0 = [999.0, 10.0, 0.0] # S, I, R tspan = (0.0, 250.0) @@ -217,7 +218,8 @@ end jump_prob_tau = JumpProblem(prob_disc, PureLeaping(), rj) # Solve with SimpleTauLeaping (save only at t_compare times) - sol_simple = solve(EnsembleProblem(jump_prob_tau), SimpleTauLeaping(), EnsembleSerial(); + sol_simple = solve( + EnsembleProblem(jump_prob_tau), SimpleTauLeaping(), EnsembleSerial(); trajectories = Nsims, dt = 0.1, saveat = t_compare, rng) # MassActionJump formulation for adaptive tau-leaping algorithms @@ -228,7 +230,8 @@ end jump_prob_maj = JumpProblem(prob_disc, PureLeaping(), maj) # Solve with SimpleExplicitTauLeaping (save only at t_compare times) - sol_adaptive = solve(EnsembleProblem(jump_prob_maj), SimpleExplicitTauLeaping(), EnsembleSerial(); + sol_adaptive = solve( + EnsembleProblem(jump_prob_maj), SimpleExplicitTauLeaping(), EnsembleSerial(); trajectories = Nsims, saveat = t_compare, rng) # Compute mean I trajectories via direct indexing (I is index 2 in SIR) @@ -286,7 +289,8 @@ end 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!)) + jumps = (ConstantRateJump(rate1, affect1!), ConstantRateJump(rate2, affect2!), + ConstantRateJump(rate3, affect3!)) u0 = [999.0, 0.0, 10.0, 0.0] # S, E, I, R tspan = (0.0, 250.0) @@ -314,7 +318,8 @@ end jump_prob_tau = JumpProblem(prob_disc, PureLeaping(), rj) # Solve with SimpleTauLeaping (save only at t_compare times) - sol_simple = solve(EnsembleProblem(jump_prob_tau), SimpleTauLeaping(), EnsembleSerial(); + sol_simple = solve( + EnsembleProblem(jump_prob_tau), SimpleTauLeaping(), EnsembleSerial(); trajectories = Nsims, dt = 0.1, saveat = t_compare, rng) # MassActionJump formulation for adaptive tau-leaping algorithms @@ -325,7 +330,8 @@ end jump_prob_maj = JumpProblem(prob_disc, PureLeaping(), maj) # Solve with SimpleExplicitTauLeaping (save only at t_compare times) - sol_adaptive = solve(EnsembleProblem(jump_prob_maj), SimpleExplicitTauLeaping(), EnsembleSerial(); + sol_adaptive = solve( + EnsembleProblem(jump_prob_maj), SimpleExplicitTauLeaping(), EnsembleSerial(); trajectories = Nsims, saveat = t_compare, rng) # Compute mean I trajectories via direct indexing (I is index 3 in SEIR) @@ -372,7 +378,7 @@ end # Test zero-rate case for SimpleExplicitTauLeaping @testset "Zero Rates Test for SimpleExplicitTauLeaping" begin # SIR model: S + I -> 2I, I -> R - reactant_stoch = [[1=>1, 2=>1], [2=>1], Pair{Int,Int}[]] + reactant_stoch = [[1=>1, 2=>1], [2=>1], Pair{Int, Int}[]] net_stoch = [[1=>-1, 2=>1], [2=>-1, 3=>1], []] rates = [0.1/1000, 0.05, 0.0] # beta/N, gamma, dummy rate for empty reaction maj = MassActionJump(rates, reactant_stoch, net_stoch) @@ -381,7 +387,7 @@ end prob = DiscreteProblem(u0, tspan) jump_prob = JumpProblem(prob, PureLeaping(), maj) - sol = solve(jump_prob, SimpleExplicitTauLeaping(); dtmin = 0.1, saveat=1.0) + sol = solve(jump_prob, SimpleExplicitTauLeaping(); dtmin = 0.1, saveat = 1.0) # Check that solution completes and covers tspan @test sol.t[end] ≈ 250.0 atol=1e-6 @@ -441,62 +447,62 @@ end 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,), + 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 @@ -505,15 +511,15 @@ end @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) + spatial_system = spatial_sys) @test_throws ErrorException JumpProblem(prob, PureLeaping(), JumpSet(maj); - hopping_constants = hopping_consts) - + 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)) diff --git a/test/rng_kwarg_tests.jl b/test/rng_kwarg_tests.jl index a237cecdc..1728b03db 100644 --- a/test/rng_kwarg_tests.jl +++ b/test/rng_kwarg_tests.jl @@ -134,7 +134,7 @@ end @testset "SDE + VR: different seeds → different trajectories" begin jprob = make_sde_vr_jump_prob() sols = [solve(jprob, SRIW1(); save_everystep = false, - rng = StableRNG(s)) for s in (100, 200, 300)] + rng = StableRNG(s)) for s in (100, 200, 300)] times = [first_jump_time(s) for s in sols] @test allunique(times) end @@ -311,7 +311,7 @@ end @testset "SDE + VR" begin jprob = make_sde_vr_jump_prob() sols = [solve(jprob, SRIW1(); save_everystep = false, - seed = s) for s in (100, 200, 300)] + seed = s) for s in (100, 200, 300)] times = [first_jump_time(s) for s in sols] @test allunique(times) end diff --git a/test/runtests.jl b/test/runtests.jl index bd909d33a..18b26382f 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -53,7 +53,7 @@ end @time @safetestset "ExtendedJumpArray remake tests" begin include("extended_jump_array_remake.jl") end @time @safetestset "Symbol based problem indexing" begin include("jprob_symbol_indexing.jl") end @time @safetestset "Long time accuracy test" begin include("longtimes_test.jl") end - @time @safetestset "Hawkes process" begin include("hawkes_test.jl") end + @time @safetestset "Hawkes process" begin include("hawkes_test.jl") end @time @safetestset "Reaction rates" begin include("spatial/reaction_rates.jl") end @time @safetestset "Hop rates" begin include("spatial/hop_rates.jl") end @time @safetestset "Topology" begin include("spatial/topology.jl") end diff --git a/test/saveat_regression.jl b/test/saveat_regression.jl index e4c845bfa..e6f544a7a 100644 --- a/test/saveat_regression.jl +++ b/test/saveat_regression.jl @@ -77,16 +77,18 @@ end jprob = JumpProblem(prob, Direct(), jump; save_positions) # Each solve owns a fresh problem and RNG; keep access to its counters # even when this test is invoked on a worker thread. - sol = solve(jprob, SSAStepper(); saveat, save_start, save_end, rng, alias_jump = true) + sol = solve( + jprob, SSAStepper(); saveat, save_start, save_end, rng, alias_jump = true) sol, events, (rng.exponential_draws, rng.uniform_draws) end expected_events = [(0.25, 1), (0.75, 2)] expected_state(t) = t < 0.25 ? [0] : t < 0.75 ? [1] : [2] - @testset "event saves=$save_events, start=$save_start, end=$save_end" for - save_events in (false, true), save_start in (false, true), - save_end in (false, true) + @testset "event saves=$save_events, start=$save_start, end=$save_end" for save_events in (false, true), + save_start in (false, true), + save_end in (false, true) + save_positions = (false, save_events) for saveat in ([0.5], [0.125, 0.5], [0.5, 1.0], 0.5, [0.0, 0.5, 1.0]) sol, events, draws = simulate(saveat; save_positions, save_start, save_end) @@ -108,7 +110,8 @@ end end @testset "Finalization after the last event" for save_start in (false, true), - save_end in (false, true) + save_end in (false, true) + sol, events, draws = simulate([0.875]; save_start, save_end) expected_times = [0.875] save_start && pushfirst!(expected_times, 0.0) @@ -129,7 +132,8 @@ end end @testset "No observation requests" for save_start in (false, true), - save_end in (false, true) + save_end in (false, true) + expected_times = [0.25, 0.75] save_start && pushfirst!(expected_times, 0.0) save_end && push!(expected_times, 1.0) diff --git a/test/sortingdirect_reuse.jl b/test/sortingdirect_reuse.jl index 134461563..a34527a34 100644 --- a/test/sortingdirect_reuse.jl +++ b/test/sortingdirect_reuse.jl @@ -24,8 +24,7 @@ function sortingdirect_reuse_problem(representation) end @testset "SortingDirect preserves learned search order" begin - @testset "$representation" for representation in - (:parameterized_massaction, :fixed_massaction, :constant) + @testset "$representation" for representation in (:parameterized_massaction, :fixed_massaction, :constant) jprob = sortingdirect_reuse_problem(representation) sol = solve(jprob, SSAStepper(); seed = 12345, alias_jump = true) @test SciMLBase.successful_retcode(sol) diff --git a/test/spatial/hop_rates.jl b/test/spatial/hop_rates.jl index 2c1d93684..5aa07d7a3 100644 --- a/test/spatial/hop_rates.jl +++ b/test/spatial/hop_rates.jl @@ -121,7 +121,8 @@ for hop_rates in hop_rates_structs target_propensities[target] = sum([hop_constants[species, site][i] for species in 1:num_species]) end - statistical_test(hop_rates, spec_propensities, target_propensities, num_species, u, site, g, rng, rel_tol) + statistical_test(hop_rates, spec_propensities, target_propensities, + num_species, u, site, g, rng, rel_tol) end end test_reset(hop_rates, num_nodes) @@ -152,7 +153,8 @@ for hop_rates in hop_rates_structs site_hop_constants[site][i] for species in 1:num_species]) end - statistical_test(hop_rates, spec_propensities, target_propensities, num_species, u, site, g, rng, rel_tol) + statistical_test(hop_rates, spec_propensities, target_propensities, + num_species, u, site, g, rng, rel_tol) end end test_reset(hop_rates, num_nodes) @@ -175,14 +177,16 @@ hop_rates_structs = [ for hop_rates in hop_rates_structs show(io, "text/plain", hop_rates) for site in 1:num_nodes - spec_propensities = [species_hop_constants[species, site] * sum(site_hop_constants[site]) for species in 1:num_species] + spec_propensities = [species_hop_constants[species, site] * + sum(site_hop_constants[site]) for species in 1:num_species] target_propensities = Dict{Int, Float64}() for (i, target) in enumerate(JP.neighbors(g, site)) target_propensities[target] = sum([species_hop_constants[species, site] * site_hop_constants[site][i] for species in 1:num_species]) end - statistical_test(hop_rates, spec_propensities, target_propensities, num_species, u, site, g, rng, rel_tol) + statistical_test(hop_rates, spec_propensities, target_propensities, + num_species, u, site, g, rng, rel_tol) end end test_reset(hop_rates, num_nodes) diff --git a/test/spatial/reaction_rates.jl b/test/spatial/reaction_rates.jl index 2e6b5ac39..4dadca11a 100644 --- a/test/spatial/reaction_rates.jl +++ b/test/spatial/reaction_rates.jl @@ -42,7 +42,9 @@ for rx_rates in rx_rates_list for site in 1:num_nodes JP.update_rx_rates!(rx_rates, 1:num_rxs, integrator, site) @test JP.total_site_rx_rate(rx_rates, site) == 1.1 - rx_props = [JP.eval_massaction_rate(u, rx, rx_rates.ma_jumps, site, rx_rates.maj_rates) for rx in 1:num_rxs] + rx_props = [JP.eval_massaction_rate( + u, rx, rx_rates.ma_jumps, site, rx_rates.maj_rates) + for rx in 1:num_rxs] rx_probs = rx_props / sum(rx_props) d = Dict{Int, Int}() for i in 1:num_samples diff --git a/test/splitcoupled.jl b/test/splitcoupled.jl index 49c81fd4a..4526e1f38 100644 --- a/test/splitcoupled.jl +++ b/test/splitcoupled.jl @@ -1,4 +1,5 @@ -using JumpProcesses, DiffEqBase, OrdinaryDiffEq, OrdinaryDiffEqFunctionMap, SciMLBase, StochasticDiffEq, Statistics +using JumpProcesses, DiffEqBase, OrdinaryDiffEq, OrdinaryDiffEqFunctionMap, SciMLBase, + StochasticDiffEq, Statistics using Test using StableRNGs rng = StableRNG(12345) diff --git a/test/ssa_aggregator_smoke.jl b/test/ssa_aggregator_smoke.jl index 8674772f8..49952e5d2 100644 --- a/test/ssa_aggregator_smoke.jl +++ b/test/ssa_aggregator_smoke.jl @@ -40,18 +40,19 @@ function ssa_smoke_problem(aggregator, representation) end rate_ab(u, p, t) = p[1] * u[1] - rate_ba(u, p, t) = p[2] * u[2] + rate_b_to_a(u, p, t) = p[2] * u[2] function affect_ab!(integrator) integrator.u[1] -= 1 integrator.u[2] += 1 nothing end - function affect_ba!(integrator) + function affect_b_to_a!(integrator) integrator.u[1] += 1 integrator.u[2] -= 1 nothing end - jumps = (ConstantRateJump(rate_ab, affect_ab!), ConstantRateJump(rate_ba, affect_ba!)) + jumps = (ConstantRateJump(rate_ab, affect_ab!), + ConstantRateJump(rate_b_to_a, affect_b_to_a!)) dep_graph = [[1, 2], [1, 2]] vartojumps_map = [[1], [2]] jumptovars_map = [[1, 2], [1, 2]] @@ -59,9 +60,9 @@ function ssa_smoke_problem(aggregator, representation) end @testset "Nonspatial SSA aggregator smoke" begin - @testset "$(typeof(aggregator)): $representation" for - aggregator in JumpProcesses.JUMP_AGGREGATORS, - representation in (:parameterized_massaction, :constant) + @testset "$(typeof(aggregator)): $representation" for aggregator in JumpProcesses.JUMP_AGGREGATORS, + representation in (:parameterized_massaction, :constant) + check_ssa_smoke_problem(ssa_smoke_problem(aggregator, representation)) end @@ -73,9 +74,9 @@ end @testset "Automatic SSA selection smoke" begin # Duplicate A -> B channels keep the total rate fixed while crossing the # reaction-count thresholds for Direct, RSSA, and RSSACR. - @testset "$num_reactions reactions" for - (num_reactions, expected_aggregator) in ((19, Direct), (20, RSSA), - (99, RSSA), (100, RSSACR)) + @testset "$num_reactions reactions" for (num_reactions, expected_aggregator) in (( + 19, Direct), (20, RSSA), + (99, RSSA), (100, RSSACR)) p = fill(0.7 / num_reactions, num_reactions) prob = DiscreteProblem([20, 0], (0.0, 0.5), p) reactant_stoch = [[1 => 1] for _ in 1:num_reactions] diff --git a/test/ssa_callback_test.jl b/test/ssa_callback_test.jl index 7e7d8ad22..535e028ad 100644 --- a/test/ssa_callback_test.jl +++ b/test/ssa_callback_test.jl @@ -69,7 +69,8 @@ function paffect!(integrator) integrator.p[2] = 1.0 reset_aggregated_jumps!(integrator) end -sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], + callback = DiscreteCallback(pcondit, paffect!)) @test all(p .== [0.0, 1.0]) @test sol[1, end] == 100 @@ -77,7 +78,8 @@ p .= [1.0, 0.0] maj1 = MassActionJump([1 => 1], [1 => -1, 2 => 1]; param_idxs = 1) maj2 = MassActionJump([2 => 1], [1 => 1, 2 => -1]; param_idxs = 2) jprob = JumpProblem(dprob, Direct(), maj1, maj2, save_positions = (false, false)) -sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], + callback = DiscreteCallback(pcondit, paffect!)) @test all(p .== [0.0, 1.0]) @test sol[1, end] == 100 @@ -85,14 +87,16 @@ p2 = [1.0, 0.0, 0.0] maj3 = MassActionJump([1 => 1], [1 => -1, 2 => 1]; param_idxs = 3) dprob = DiscreteProblem(u₀, tspan, p2) jprob = JumpProblem(dprob, Direct(), maj1, maj2, maj3, save_positions = (false, false)) -sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], + callback = DiscreteCallback(pcondit, paffect!)) @test all(p2 .== [0.0, 1.0, 0.0]) @test sol[1, end] == 100 p2 .= [1.0, 0.0, 0.0] jprob = JumpProblem(dprob, Direct(), JumpSet(; massaction_jumps = [maj1, maj2, maj3]), save_positions = (false, false)) -sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], + callback = DiscreteCallback(pcondit, paffect!)) @test all(p2 .== [0.0, 1.0, 0.0]) @test sol[1, end] == 100 @@ -101,7 +105,8 @@ dprob = DiscreteProblem(u₀, tspan, p) maj4 = MassActionJump([[1 => 1], [2 => 1]], [[1 => -1, 2 => 1], [1 => 1, 2 => -1]]; param_idxs = [1, 2]) jprob = JumpProblem(dprob, Direct(), maj4, save_positions = (false, false)) -sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], callback = DiscreteCallback(pcondit, paffect!)) +sol = solve(jprob, SSAStepper(); rng, tstops = [1000.0], + callback = DiscreteCallback(pcondit, paffect!)) @test all(p .== [0.0, 1.0]) @test sol[1, end] == 100 diff --git a/test/ssa_kernel_tests.jl b/test/ssa_kernel_tests.jl index 58beb3bcc..f778b1fe0 100644 --- a/test/ssa_kernel_tests.jl +++ b/test/ssa_kernel_tests.jl @@ -61,6 +61,7 @@ function run_ssa_kernel_tests(backend, nsims) @test sol_kernel.u[1].t == sol_serial.u[1].t for k in eachindex(sol_kernel.u[1].t), s in 1:3 + mean_kernel = mean(sol_kernel.u[i].u[k][s] for i in 1:nsims) mean_serial = mean(sol_serial.u[i].u[k][s] for i in 1:nsims) @test isapprox(mean_kernel, mean_serial, rtol = 0.05, atol = 0.5) @@ -70,7 +71,8 @@ function run_ssa_kernel_tests(backend, nsims) # Second order reaction 2A -> B. Exercises the falling factorial A(A-1) in the # propensity and the combinatoric prefactor folded into `scaled_rates`. let - jump_prob = majump_prob([0.01], [[1 => 2]], [[1 => -2, 2 => 1]], [100, 0], (0.0, 5.0)) + jump_prob = majump_prob( + [0.01], [[1 => 2]], [[1 => -2, 2 => 1]], [100, 0], (0.0, 5.0)) sol_kernel = solve(EnsembleProblem(jump_prob), SSAStepper(), EnsembleGPUKernel(backend); trajectories = nsims, saveat = 1.0) @@ -78,6 +80,7 @@ function run_ssa_kernel_tests(backend, nsims) trajectories = nsims, saveat = 1.0) for k in eachindex(sol_kernel.u[1].t), s in 1:2 + mean_kernel = mean(sol_kernel.u[i].u[k][s] for i in 1:nsims) mean_serial = mean(sol_serial.u[i].u[k][s] for i in 1:nsims) @test isapprox(mean_kernel, mean_serial, rtol = 0.05, atol = 0.5) @@ -89,7 +92,8 @@ function run_ssa_kernel_tests(backend, nsims) # Third order reaction 3A -> B, checking the deeper falling factorial. let - jump_prob = majump_prob([1e-4], [[1 => 3]], [[1 => -3, 2 => 1]], [60, 0], (0.0, 5.0)) + jump_prob = majump_prob( + [1e-4], [[1 => 3]], [[1 => -3, 2 => 1]], [60, 0], (0.0, 5.0)) sol_kernel = solve(EnsembleProblem(jump_prob), SSAStepper(), EnsembleGPUKernel(backend); trajectories = nsims, saveat = 2.5) @@ -115,6 +119,7 @@ function run_ssa_kernel_tests(backend, nsims) trajectories = nsims, saveat = 5.0) for k in eachindex(sol_kernel.u[1].t), s in 1:3 + mean_kernel = mean(sol_kernel.u[i].u[k][s] for i in 1:nsims) mean_serial = mean(sol_serial.u[i].u[k][s] for i in 1:nsims) @test isapprox(mean_kernel, mean_serial, rtol = 0.05, atol = 0.5) diff --git a/test/stochastic_dispatch.jl b/test/stochastic_dispatch.jl index 89419036c..9f9243351 100644 --- a/test/stochastic_dispatch.jl +++ b/test/stochastic_dispatch.jl @@ -6,21 +6,22 @@ function stochastic_jump_problem(kind; seed = 0, variable = false, callback = no g!(du, u, p, t) = (du .= 0.2) f_rode!(du, u, p, t, W) = (du .= -0.1 .* u .+ 0.2 .* W) prob = kind === :SDE ? - SDEProblem(f!, g!, [10.0], (0.0, 0.5), zeros(Int, 2); seed) : - RODEProblem(f_rode!, [10.0], (0.0, 0.5), zeros(Int, 2); seed) + SDEProblem(f!, g!, [10.0], (0.0, 0.5), zeros(Int, 2); seed) : + RODEProblem(f_rode!, [10.0], (0.0, 0.5), zeros(Int, 2); seed) rate(u, p, t) = 10.0 affect!(integrator) = (integrator.u[1] += 1) jump = variable ? VariableRateJump(rate, affect!) : ConstantRateJump(rate, affect!) JumpProblem(prob, Direct(), jump; callback) end -stochastic_rng_options(kind) = +function stochastic_rng_options(kind) kind === :default ? (;) : kind === :seed ? (; seed = 456) : kind === :zero_seed ? (; seed = 0) : kind === :task_rng_seed ? (; rng = Random.default_rng(), seed = 456) : kind === :task_rng ? (; rng = Random.default_rng()) : (; rng = StableRNG(789), seed = 456) +end @testset "Stochastic initialization dispatches to the backend" begin for (kind, alg) in ((:SDE, EM()), (:RODE, RandomEM())) @@ -31,12 +32,13 @@ stochastic_rng_options(kind) = end @testset "Stochastic solve/init use the same RNG policy" begin - @testset "$kind stored seed=$stored_seed options=$options_kind variable=$variable" for - (kind, alg, variable) in ((:SDE, EM(), false), (:SDE, EM(), true), - (:RODE, RandomEM(), false)), - stored_seed in (0, 123), - options_kind in (:default, :seed, :zero_seed, :task_rng_seed, :task_rng, - :explicit_rng_seed) + @testset "$kind stored seed=$stored_seed options=$options_kind variable=$variable" for (kind, alg, variable) in (( + :SDE, EM(), false), (:SDE, EM(), true), + (:RODE, RandomEM(), false)), + stored_seed in (0, 123), + options_kind in (:default, :seed, :zero_seed, :task_rng_seed, :task_rng, + :explicit_rng_seed) + Random.seed!(999) sol = solve(stochastic_jump_problem(kind; seed = stored_seed, variable), alg; dt = 0.01, stochastic_rng_options(options_kind)...) @@ -57,7 +59,7 @@ end @test sol.u == integrator.sol.u @test sol.seed == integrator.sol.seed expected_seed = options_kind === :explicit_rng_seed ? 0 : - options_kind in (:seed, :task_rng_seed) ? 456 : stored_seed + options_kind in (:seed, :task_rng_seed) ? 456 : stored_seed if expected_seed != 0 || options_kind === :explicit_rng_seed @test sol.seed == expected_seed end @@ -65,14 +67,16 @@ end end @testset "Stochastic initialization preserves callback merging" begin - @testset "merge_callbacks=$merge_callbacks entry=$entry" for - merge_callbacks in (true, false), entry in (:solve, :init) + @testset "merge_callbacks=$merge_callbacks entry=$entry" for merge_callbacks in (true, false), + entry in (:solve, :init) + cb1 = DiscreteCallback((u, t, integrator) -> t == 0.25, integrator -> (integrator.p[1] += 1); save_positions = (false, false)) cb2 = DiscreteCallback((u, t, integrator) -> t == 0.25, integrator -> (integrator.p[2] += 1); save_positions = (false, false)) jprob = stochastic_jump_problem(:SDE; callback = cb1) - options = (; seed = 123, dt = 0.01, tstops = [0.25], callback = cb2, merge_callbacks) + options = (; + seed = 123, dt = 0.01, tstops = [0.25], callback = cb2, merge_callbacks) sol = if entry === :solve solve(jprob, EM(); options...) else @@ -101,14 +105,16 @@ end @testset "Stochastic initialization preserves native alias settings" begin for (kind, alg) in ((:SDE, EM()), (:RODE, RandomEM())), - alias_mode in (:default, :alias, :copy, :specifier_alias, :specifier_copy) + alias_mode in (:default, :alias, :copy, :specifier_alias, :specifier_copy) + jprob = stochastic_jump_problem(kind) - constructor = kind === :SDE ? SciMLBase.SDEAliasSpecifier : SciMLBase.RODEAliasSpecifier + constructor = kind === :SDE ? SciMLBase.SDEAliasSpecifier : + SciMLBase.RODEAliasSpecifier alias_jumps = alias_mode === :default ? Threads.threadid() == 1 : - alias_mode in (:alias, :specifier_alias) + alias_mode in (:alias, :specifier_alias) alias = alias_mode === :default ? nothing : - alias_mode in (:alias, :copy) ? alias_jumps : - constructor(; alias_jumps, alias_u0 = false) + alias_mode in (:alias, :copy) ? alias_jumps : + constructor(; alias_jumps, alias_u0 = false) integrator = init(jprob, alg; dt = 0.01, seed = 123, alias) aliases_aggregation = any(integrator.opts.callback.discrete_callbacks) do callback callback.condition === jprob.discrete_jump_aggregation diff --git a/test/thread_safety.jl b/test/thread_safety.jl index 007025512..473c8ad4c 100644 --- a/test/thread_safety.jl +++ b/test/thread_safety.jl @@ -16,20 +16,23 @@ jprob = JumpProblem(dprob, Direct(), maj) sol = solve(EnsembleProblem(jprob), SSAStepper(), EnsembleThreads(); trajectories = 400) @test length(sol.u) == 400 -firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol.u)] +firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] + for i in 1:length(sol.u)] @test allunique(firstrx_time) sol2 = solve(EnsembleProblem(jprob; safetycopy = true), SSAStepper(), EnsembleThreads(); trajectories = 400) @test length(sol2.u) == 400 -firstrx_time2 = [sol2.u[i].t[findfirst(>(sol2.u[i].t[1]), sol2.u[i].t)] for i in 1:length(sol2.u)] +firstrx_time2 = [sol2.u[i].t[findfirst(>(sol2.u[i].t[1]), sol2.u[i].t)] + for i in 1:length(sol2.u)] @test allunique(firstrx_time2) @testset "Threaded ensemble RNG and seed controls" begin paths(sol) = [(trajectory.t, trajectory.u) for trajectory in sol.u] first_event_times(sol) = [trajectory.t[findfirst(>(first(trajectory.t)), trajectory.t)] - for trajectory in sol.u] + for trajectory in sol.u] for safetycopy in (false, true), input in (:rng, :seed) + @testset "safetycopy=$safetycopy, $input" begin # Construct a fresh master RNG for each ensemble solve. Only the # ensemble layer uses this object; each trajectory receives its own RNG. @@ -69,7 +72,8 @@ let prob = EnsembleProblem(jump_prob) sol = solve(prob, Tsit5(), EnsembleThreads(), trajectories = 400, save_everystep = false) - firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol.u)] + firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] + for i in 1:length(sol.u)] @test allunique(firstrx_time) end end @@ -87,7 +91,8 @@ let prob = EnsembleProblem(jump_prob) sol = solve(prob, SRIW1(), EnsembleThreads(); trajectories = 400, save_everystep = false) - firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] for i in 1:length(sol.u)] + firstrx_time = [sol.u[i].t[findfirst(>(sol.u[i].t[1]), sol.u[i].t)] + for i in 1:length(sol.u)] @test allunique(firstrx_time) end end From 3cb6de7b4bfaf0cadb229a8c8362b437903736cc Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Fri, 9 Oct 2026 16:16:25 -0400 Subject: [PATCH 37/52] Stop copying jump state and add SSAStepper alias controls Solvers owned by JumpProcesses (SSAStepper and the ODE/DAE initialization path) now always reuse the JumpProblem's jump aggregator and callbacks, re-initializing them at every init, instead of copying them on threads other than the first. Callers isolate concurrent solves with deepcopy or an EnsembleProblem. The alias_jump keyword is removed: any value raises a migration error from solve and init (including when stored on the wrapped problem), from the JumpProblem constructor, and from multi-trajectory EnsembleGPUKernel solves. SSAStepper supports the common alias keyword for u0, p, and tstops (a Bool, a DiscreteAliasSpecifier, or an ODEAliasSpecifier) with OrdinaryDiffEq's defaults. It copies non-Vector tstops containers into a Vector at init, and saves recursive snapshots of the state as OrdinaryDiffEq does. Raise the SciMLBase compat floor to 3.34. Document the ownership and safety rules in HISTORY, in the JumpProblem, remake, SSAStepper, SortingDirect, and EnsembleGPUKernel docstrings, on the API, migration, and FAQ pages, and in a new tutorial on ensembles and problem reuse. Add jump-state ownership tests and kernel alias tests, and migrate the existing alias_jump uses in the tests. Validation on Julia 1.12.7 with the pinned StochasticDiffEqCore: InterfaceI (1,751 assertions) and InterfaceII (2,727) pass, as do four-thread ThreadSafety (333), strict QA (20, plus 5 pin checks), formatting, spelling and whitespace checks, all 12 tutorial blocks, and a Documenter build of the changed pages. Benchmark allocations are unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) --- HISTORY.md | 41 +- Project.toml | 2 +- docs/pages.jl | 3 +- docs/src/api.md | 97 ++++- docs/src/faq.md | 17 +- docs/src/jump_solve.md | 3 + docs/src/jump_types.md | 17 +- docs/src/migration.md | 68 +++- .../tutorials/ensembles_and_problem_reuse.md | 270 ++++++++++++ ext/JumpProcessesKernelAbstractionsExt.jl | 16 + ext/explicit_tau_leaping.jl | 1 + ext/ssa_stepper.jl | 1 + src/JumpProcesses.jl | 1 + src/SSA_stepper.jl | 128 ++++-- src/aggregators/aggregators.jl | 8 + src/problem.jl | 56 ++- src/simple_regular_solve.jl | 7 + src/solve.jl | 49 ++- test/explicit_tau_kernel_tests.jl | 1 + test/jump_state_ownership.jl | 384 ++++++++++++++++++ test/kernel_contract_tests.jl | 45 ++ test/massaction_state.jl | 10 +- test/runtests.jl | 2 + test/saveat_regression.jl | 8 +- test/sortingdirect_reuse.jl | 8 +- test/ssa_kernel_tests.jl | 1 + 26 files changed, 1132 insertions(+), 112 deletions(-) create mode 100644 docs/src/tutorials/ensembles_and_problem_reuse.md create mode 100644 test/jump_state_ownership.jl diff --git a/HISTORY.md b/HISTORY.md index 34d53f597..d033fafe4 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -36,12 +36,41 @@ changes are not intended for a 9.x release. `Xoshiro`. `TaskLocalRNG` is converted to `Xoshiro`; `seed = 0` means no seed override on this pathway, whereas the JumpProcesses RNG resolver uses `Xoshiro(0)`. - - **Breaking**: The singular `alias_jump` keyword is no longer consumed on - the StochasticDiffEq `init` pathway. For SDE jump-state copying, pass - `alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)` to `solve` or - `init`; use `alias_jumps = true` in the specifier to reuse jump state. RODE - problems use `SciMLBase.RODEAliasSpecifier`. Other JumpProcesses-owned - pathways retain their existing `alias_jump` behavior. + - **Breaking**: Solvers owned by JumpProcesses (`SSAStepper` and the + OrdinaryDiffEq ODE/DAE pathways) no longer copy a `JumpProblem`'s jump state. + On every thread they reuse the problem's aggregator and jump callbacks, + re-initializing them at each `init`; previously, solves on threads other than + thread 1 copied this state. A `JumpProblem` therefore supports one active solve + or integrator at a time. Concurrent solves from threads or tasks, and + integrators alive at the same time, must each use an independent copy, such as + `deepcopy(jprob)`. Problems created with `remake` share the original's jump + state, so they are not independent either. `EnsembleProblem` provides the + needed copies: `EnsembleThreads` copies the problem once per spawned task, and + `safetycopy = true` copies it for every trajectory. See the new tutorial on + ensembles and problem reuse for patterns that balance performance and safety. + - **Breaking**: The `alias_jump` keyword argument has been removed. Passing it to + `solve`, `init`, `JumpProblem`, or a multi-trajectory `EnsembleGPUKernel` + solve raises an error; remove it, and solve a copy of the problem when + independent jump state is needed. StochasticDiffEq SDE/RODE solvers ignore the + keyword and keep their own policy: pass + `alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)` (or + `SciMLBase.RODEAliasSpecifier` for RODEs) to copy jump state, or + `alias_jumps = true` to reuse it. When unspecified, those solvers reuse jump + state on thread 1 and copy it on other threads. + - `SSAStepper` now supports the common `alias` keyword argument for its `u0`, `p`, + and `tstops` inputs, given as a `Bool`, a `SciMLBase.DiscreteAliasSpecifier`, or + an `ODEAliasSpecifier` (which also controls `tstops`). The defaults match + OrdinaryDiffEq: `u0` is copied, `p` is reused, and a caller's `tstops` array is + never modified. `alias` never controls jump state. + - `SSAStepper` now saves states, and evaluates `integrator(t)`, with + `recursivecopy` as OrdinaryDiffEq does, so nested state types are saved as + independent snapshots. `tstops` given as a number, a tuple, or an + `AbstractVector` other than a `Vector` of the time type (such as a range or + view) is copied into a `Vector` at `init`, so `add_tstop!` works with it. + - **Breaking**: Multi-trajectory `EnsembleGPUKernel` solves reject `alias` values + other than `nothing`, and the `alias_jumps` keyword, since kernels always build + device-owned state. + - The minimum supported SciMLBase version is now 3.34. - `SSAIntegrator` now supports the `SciMLBase` RNG interface (`has_rng`, `get_rng`, `set_rng!`). - **Breaking**: Passing `seed` to a multi-trajectory `EnsembleGPUKernel` diff --git a/Project.toml b/Project.toml index d481a336c..627d1de7f 100644 --- a/Project.toml +++ b/Project.toml @@ -57,7 +57,7 @@ PrecompileTools = "1.2.1" Random = "1" RecursiveArrayTools = "4.3.0, 4" SafeTestsets = "0.1" -SciMLBase = "3.18.0, 3.1" +SciMLBase = "3.34" SciMLTesting = "2.8" SimpleNonlinearSolve = "2, 3" StableRNGs = "1" diff --git a/docs/pages.jl b/docs/pages.jl index 761708c58..53cbdd932 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -5,7 +5,8 @@ pages = ["index.md", "tutorials/discrete_stochastic_example.md", "tutorials/point_process_simulation.md", "tutorials/jump_diffusion.md", - "tutorials/spatial.md"], + "tutorials/spatial.md", + "tutorials/ensembles_and_problem_reuse.md"], "Applications" => Any["applications/advanced_point_process.md"], "Type Documentation" => Any["Jumps, JumpProblem, and Aggregators" => "jump_types.md", "Jump solvers" => "jump_solve.md"], diff --git a/docs/src/api.md b/docs/src/api.md index 82ba9fd08..2ac1a58e8 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -13,6 +13,7 @@ PureLeaping SSAStepper SplitCoupledJumpProblem reset_aggregated_jumps! +remake(::JumpProblem) ``` ## Jump Types @@ -123,6 +124,80 @@ Anything else from SciMLBase -- the BVP, DAE, DDE, nonlinear and optimization pr classes, the SciML operators, and the internals -- is not re-exported here; import it from SciMLBase directly. +## [Jump state ownership and problem reuse](@id jump_state_ownership) + +A [`JumpProblem`](@ref) stores mutable jump state: the jump aggregator and its +callbacks. Solvers owned by JumpProcesses ([`SSAStepper`](@ref) and the +OrdinaryDiffEq ODE/DAE pathways) reuse this state directly and re-initialize it +from the current state and parameters at every `init`. They never copy it. +Consequently: + + - Repeated serial solves of one problem are safe and avoid any copying. + - One `JumpProblem` supports one active solve or integrator at a time. + Concurrent solves from threads or tasks, and integrators that are alive at the + same time, each need an independent copy, such as `deepcopy(jprob)`. + - [`remake`](@ref remake(::JumpProblem)) shares the original's jump state, so + remade problems are not independent for concurrent use. + - `EnsembleProblem` supplies isolation. `EnsembleSerial` reuses the problem + sequentially. `EnsembleThreads` with `safetycopy = false` copies the problem + once per spawned task, with a serial fallback for one thread or a single + trajectory. `safetycopy = true` copies the problem before every trajectory, + before `prob_func` runs. + +The [ensembles and problem reuse tutorial](@ref ensembles_problem_reuse) shows +patterns that balance performance and safety, including ensemble hooks, stateful +callbacks, manual threading, and distributed ensembles. + +The `alias_jump` keyword argument has been removed; passing it raises an error. +The plural `alias_jumps` is not a solver keyword argument either: `solve` and +`init` reject it through SciMLBase's keyword validation, and multi-trajectory +`EnsembleGPUKernel` solves, which bypass that validation, reject it explicitly. + +### Aliasing `SSAStepper` inputs + +[`SSAStepper`](@ref) supports the common SciML `alias` keyword argument for its +`u0`, `p`, and `tstops` inputs. Pass a `Bool` to apply one choice to all three, a +`SciMLBase.DiscreteAliasSpecifier` to control `u0` and `p`, or an +`ODEAliasSpecifier` to also control `tstops`. A field set to `true` permits +aliasing, `false` requests a copy, and `nothing` selects the default: + +| Input | Default (`nothing`) | `true` | `false` | +|:-------- |:------------------------------------ |:---------------------------------------------------- |:--------------------------- | +| `u0` | copied with `recursivecopy` | the integrator uses and mutates `u0` | copied | +| `p` | reused | reused | copied with `recursivecopy` | +| `tstops` | the caller's array is never modified | new stops may be inserted into the caller's `Vector` | copied at `init` | + +`tstops` may be a number, a tuple, or an `AbstractVector` of times (such as a +range or view), or a callable `(p, tspan) -> times`. Only a `Vector` of the time +type can be aliased; other supported inputs are copied into a new `Vector` at +`init`. `alias` never controls jump state; for example, +`alias = false` does not give a solve its own jump aggregator. + +```julia +sol = solve(jprob, SSAStepper(); alias = SciMLBase.DiscreteAliasSpecifier(alias_p = false)) +``` + +### Coupled ODE/DAE problems + +For a `JumpProblem` wrapping an ODE or DAE problem, `alias` is passed unchanged to +the underlying solver. OrdinaryDiffEq solvers, including its DAE solvers, require +an `ODEAliasSpecifier`. The jump state is always reused, as described above. + +### StochasticDiffEq SDE/RODE problems + +StochasticDiffEq owns jump-state copying for SDE and RODE problems. Pass +`alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)`, or +`SciMLBase.RODEAliasSpecifier` for RODE problems, to copy the jump state, or set +`alias_jumps = true` to reuse it. When the field is unspecified, the backend +reuses the jump state on thread 1 and copies it on other threads. This backend +ignores the removed `alias_jump` keyword. + +### Kernel ensembles + +Multi-trajectory [`EnsembleGPUKernel`](@ref) solves always build device-owned +state. They reject the removed `alias_jump` keyword and any `alias` value other +than `nothing`. With `trajectories = 1`, the serial solver's rules apply. + ## Random Number Generator Control JumpProcesses supports controlling the random number generator (RNG) used for @@ -237,20 +312,14 @@ inputs. Both routes use the backend's JumpProblem initializer, including its RNG policy, jump-state copying, and callback setup. DiffEqBase merges stored problem keywords and user callbacks before dispatching `init`. -Control stochastic jump-state copying through the backend's alias specifier: -`alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)` for SDE problems, -or `SciMLBase.RODEAliasSpecifier` for RODE problems. Set the `alias_jumps` field -to `true` to reuse the original jump state. When that field is unspecified, the -backend aliases on thread 1 and copies on other threads. Both `solve` and `init` -follow this policy; `alias_jumps` is a field of the specifier, not a standalone -keyword argument. - -The singular `alias_jump` keyword is no longer consumed on the SDE `init` -path; use the backend alias specifier above. JumpProcesses-owned SSAStepper -and ODE/DAE initialization retain `alias_jump`, defaulting to `true` on -thread 1 and `false` on other threads. Independent mutable jump state is -still required for concurrent solves, including when using -`EnsembleProblem(...; safetycopy = false)`. +Jump-state copying on these pathways follows the backend's alias specifier; see +[Jump state ownership and problem reuse](@ref jump_state_ownership). In contrast, +`SSAStepper` and `__jump_init` never copy jump state: they build the integrator +around the problem's own jump callbacks, which `init` re-initializes. `__jump_init` +raises the removed-`alias_jump` error and forwards `alias` unchanged to the +underlying solver. JumpProcesses still defines the transitional +`resetted_jump_problem` and `reset_jump_problem!` helpers for released +StochasticDiffEq versions, but no longer calls them itself. For **tau-leaping**, JumpProcesses defines a custom `DiffEqBase.solve` that bypasses the standard `__solve`/`__init` pathway. It calls `resolve_rng` diff --git a/docs/src/faq.md b/docs/src/faq.md index 07ae35307..1c1c3b7ff 100644 --- a/docs/src/faq.md +++ b/docs/src/faq.md @@ -22,7 +22,22 @@ instead of an ODE/SDE time stepper can give a significant performance boost. In [Remaking `JumpProblem`s](@ref) we show how to modify parameters, the initial condition, and other components of a generated `JumpProblem`. This can be useful when trying to call `solve` many times while avoiding reallocations of the -internal aggregators for each new parameter value or initial condition. +internal aggregators for each new parameter value or initial condition. Remade +problems share the original's jump state, so they can be solved one after +another but not concurrently. The +[ensembles and problem reuse tutorial](@ref ensembles_problem_reuse) shows how +to combine `remake` with `EnsembleProblem`s safely. + +## Can I solve the same `JumpProblem` from multiple threads or tasks? + +Not directly. Solvers owned by JumpProcesses reuse a `JumpProblem`'s jump +aggregator and callbacks rather than copying them, so one problem supports one +active solve or integrator at a time. For concurrent solves, give each task its +own `deepcopy` of the problem, made once per task and reused for that task's +solves, or use an `EnsembleProblem`, which makes these copies for you. Problems +created with `remake` share the original's jump state and are not independent. +See [Jump state ownership and problem reuse](@ref jump_state_ownership) and the +[ensembles and problem reuse tutorial](@ref ensembles_problem_reuse). ## How can I define collections of many different jumps and pass them to `JumpProblem`? diff --git a/docs/src/jump_solve.md b/docs/src/jump_solve.md index 41ac44124..2fa1a18b9 100644 --- a/docs/src/jump_solve.md +++ b/docs/src/jump_solve.md @@ -62,6 +62,9 @@ algorithms are optimized for pure jump problems. [`SSAStepper`](@ref) generated-solution uses piecewise constant interpolation, and can therefore exactly evaluate the sampled solution path at any time when only saving the post-jump state for each jump. + `SSAStepper` supports the common `alias` keyword argument for `u0`, `p`, + and `tstops`, and reuses the problem's jump state rather than copying it; + see [Jump state ownership and problem reuse](@ref jump_state_ownership). ## RegularJump Compatible Methods diff --git a/docs/src/jump_types.md b/docs/src/jump_types.md index 1128b8a9f..65b42a4c1 100644 --- a/docs/src/jump_types.md +++ b/docs/src/jump_types.md @@ -428,13 +428,16 @@ When running many simulations, it can often be convenient to update the initial condition or simulation parameters without having to create and initialize a new `JumpProblem`. In such situations `remake` can be used to change the initial condition, time span, and the parameter vector. **Note,** the new `JumpProblem` -will alias internal data structures from the old problem, including core -components of the SSA aggregators. Sequential solves refresh parameter-mapped -mass-action working rates from each problem's current parameters; `remake` -does not overwrite rates in a shared parameter-mapped definition. For concurrent -solves, create independent problem copies. `SortingDirect` also retains its -learned search order, so a reused problem can have different initial sampling -state even with the same seed. +shares the old problem's jump state, including the SSA aggregator and jump +callbacks. Sequential solves are safe: each `init` re-initializes that state, +including parameter-mapped mass-action working rates, from the problem being +solved, and `remake` does not overwrite rates in a shared parameter-mapped +definition. The original and remade problems must not be solved concurrently, or +used for integrators that are alive at the same time; create independent copies +with `deepcopy` instead. See the +[ensembles and problem reuse tutorial](@ref ensembles_problem_reuse). +`SortingDirect` also retains its learned search order, so a reused problem can +have different initial sampling state even with the same seed. As an example, consider the following SIR model: diff --git a/docs/src/migration.md b/docs/src/migration.md index c3824a0a0..addcadfa2 100644 --- a/docs/src/migration.md +++ b/docs/src/migration.md @@ -137,26 +137,58 @@ the `SSAStepper`, `SimpleTauLeaping`, and `SimpleExplicitTauLeaping` kernels. With `trajectories = 1`, execution uses the serial CPU fallback and its usual RNG/seed controls. -## Copying and concurrent solves +## Jump state ownership, `alias_jump`, and concurrent solves + +JumpProcesses 10 never copies a `JumpProblem`'s jump state in the solvers it owns +(`SSAStepper` and the OrdinaryDiffEq ODE/DAE pathways). Every solve and `init` +reuses the problem's jump aggregator and callbacks, re-initializing them from the +current state and parameters. Previously, solves on threads other than thread 1 +copied this state. + +This makes repeated serial solves cheap, but one `JumpProblem` now supports only +one active solve or integrator at a time. Code that solves one shared problem from +several threads or tasks must give each task its own copy: + +```julia +# Before: relied on copies made off thread 1 (no longer made) +Threads.@threads for i in 1:n + sols[i] = solve(jprob, SSAStepper()) +end + +# After: copy once per task and reuse the copy within the task +chunks = Iterators.partition(1:n, cld(n, Threads.nthreads())) +tasks = map(chunks) do idxs + Threads.@spawn begin + local_prob = deepcopy(jprob) + [solve(local_prob, SSAStepper()) for _ in idxs] + end +end +sols = reduce(vcat, fetch.(tasks)) +``` + +An `EnsembleProblem` already provides this isolation: `EnsembleThreads` with +`safetycopy = false` copies the problem once per spawned task, and +`safetycopy = true` copies it before every trajectory. Problems created with +`remake` share the original's jump state, so they are not independent copies; +use `deepcopy` instead. The same applies to integrators from `init` that are +alive at the same time. See the +[ensembles and problem reuse tutorial](@ref ensembles_problem_reuse) for details. + +The `alias_jump` keyword argument has been removed, and passing it raises an +error. Delete it from `solve` and `init` calls; `alias_jump = true` is now the +behavior everywhere, and `alias_jump = false` is replaced by solving a copy of +the problem. StochasticDiffEq SDE/RODE solvers keep their own policy, controlled +by `alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)` (or +`SciMLBase.RODEAliasSpecifier`), and ignore `alias_jump`. + +`SSAStepper` now supports the common `alias` keyword argument for `u0`, `p`, and +`tstops`; see [Jump state ownership and problem reuse](@ref jump_state_ownership). +Its defaults match the previous behavior: `u0` is copied, `p` is reused, and a +caller's `tstops` array is never modified. Integrator RNGs and separate working-rate buffers do not make a shared -`JumpProblem` safe for concurrent solves: callbacks and aggregators remain -mutable. Use independent problem copies for manually launched tasks. With -`EnsembleProblem`, `safetycopy = true` requests a separate problem copy per -trajectory. If using `safetycopy = false`, ensure your `prob_func` and solver -settings preserve the required isolation; do not assume arbitrary shared -mutable objects are safe. -Do not return one shared mutable RNG from a custom ensemble `rng_func`. - -The existing `alias_jump` keyword remains available for JumpProcesses-owned -SSAStepper and ODE/DAE initialization. Its default is `true` on thread 1 and -`false` on other threads. The SDE/RODE backend instead uses -`alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)` or -`SciMLBase.RODEAliasSpecifier(; alias_jumps = false)` to request copied jump -state. The singular `alias_jump` keyword is no longer consumed on the SDE -`init` path; use the backend specifier. `alias_jumps` is a field of that -specifier, not a standalone keyword. A common JumpProcesses alias API is -not introduced by this migration. +`JumpProblem` safe for concurrent solves on their own. Do not return one shared +mutable RNG from a custom ensemble `rng_func`. `SortingDirect` retains its learned search order across solves and resets. For identical seeded trajectories, start from identical learned state as well diff --git a/docs/src/tutorials/ensembles_and_problem_reuse.md b/docs/src/tutorials/ensembles_and_problem_reuse.md new file mode 100644 index 000000000..5c0ff634a --- /dev/null +++ b/docs/src/tutorials/ensembles_and_problem_reuse.md @@ -0,0 +1,270 @@ +# [Ensembles, `remake`, and problem reuse](@id ensembles_problem_reuse) + +Stochastic simulations are usually run many times, so it is worth knowing how to reuse +a `JumpProblem` efficiently and how to run simulations in parallel safely. This tutorial +explains the rules that JumpProcesses follows and shows patterns that balance +performance and safety. + +The key fact: a [`JumpProblem`](@ref) stores **mutable jump state**, namely the jump +aggregator and its callbacks. Solvers owned by JumpProcesses ([`SSAStepper`](@ref) +and the ODE/DAE pathways) reuse this state directly and re-initialize it at the start +of every solve. They never copy it. This makes repeated solves fast, but it means that +**one `JumpProblem` supports one active solve or integrator at a time**: + + - Repeated serial solves of one problem are safe and fast. + - Concurrent solves, from threads or tasks, need independent copies of the problem. + - Problems created with `remake` share the original's jump state. + - `EnsembleProblem`s create the copies they need for you. + +We use a simple birth-death process, ``\emptyset \to X`` at rate ``p_1`` and +``X \to \emptyset`` at rate ``p_2 X``, throughout: + +```@example reuse +using JumpProcesses, SciMLBase + +maj = MassActionJump([[0 => 1], [1 => 1]], [[1 => 1], [1 => -1]]; param_idxs = [1, 2]) +dprob = DiscreteProblem([10], (0.0, 100.0), [1.0, 0.1]) +# only save the initial and final states +jprob = JumpProblem(dprob, Direct(), maj; save_positions = (false, false)) +nothing # hide +``` + +## Repeated serial solves + +Solving the same problem repeatedly in a loop is the fastest way to generate many +samples: each solve re-initializes the problem's jump state from the problem's +current initial condition and parameters, so nothing needs to be copied or rebuilt. + +```@example reuse +final_counts = [solve(jprob, SSAStepper(); seed).u[end][1] for seed in 1:1000] +sum(final_counts) / length(final_counts) +``` + +## Parameter sweeps with `remake` + +`remake` creates a problem with a new initial condition, parameters, or time span. +The new problem **shares** the original's jump state rather than copying it, which +makes `remake` cheap: + +```@example reuse +remade = remake(jprob; p = [2.0, 0.1]) +remade.discrete_jump_aggregation === jprob.discrete_jump_aggregation +``` + +Since every solve re-initializes the shared state from the problem being solved, remade +problems can be solved one after another: + +```@example reuse +mean_counts = map([0.5, 1.0, 2.0]) do birth_rate + prob = remake(jprob; p = [birth_rate, 0.1]) + sum(solve(prob, SSAStepper(); seed).u[end][1] for seed in 1:500) / 500 +end +``` + +However, the original and remade problems must not be solved concurrently, and must +not be used for integrators that are alive at the same time. Use `deepcopy` when you +need an independent problem. + +## Ensembles + +An [`EnsembleProblem`](https://docs.sciml.ai/DiffEqDocs/stable/features/ensemble/) +handles the copying needed to run trajectories safely: + + - `EnsembleSerial()` with `safetycopy = false` (the default when no `prob_func` is + given) reuses the problem sequentially and makes no copies. + - `EnsembleThreads()` with `safetycopy = false` copies the problem once per spawned + task, and each task reuses its copy for its trajectories. With one thread, or a + batch of one trajectory, it falls back to the serial path and uses the original + problem. + - With `safetycopy = true` (the default when a `prob_func` is given), every + trajectory solves its own copy, made **before** `prob_func` is called. + +```@example reuse +ensemble = EnsembleProblem(jprob) +esol = solve(ensemble, SSAStepper(), EnsembleThreads(); trajectories = 1000, seed = 1) +length(esol.u) # the trajectories +``` + +### Custom `prob_func`s + +A `prob_func` lets each trajectory use a different problem. A `prob_func` that only +modifies the problem it is given, for example with `remake`, is safe without safety +copies, because with `EnsembleThreads()` the problem it is given is already the copy +owned by the current task. Turning off the safety copies then avoids one copy per +trajectory: + +```@example reuse +prob_func(prob, ctx) = remake(prob; p = [0.5 + 0.001 * ctx.sim_id, 0.1]) +with_copies = EnsembleProblem(jprob; prob_func) # safetycopy = true by default +without_copies = EnsembleProblem(jprob; prob_func, safetycopy = false) + +function run_ensemble(ensemble) + solve(ensemble, SSAStepper(), EnsembleThreads(); + trajectories = 2000) +end +run_ensemble(with_copies), run_ensemble(without_copies) # compile before timing +(; with_copies = @elapsed(run_ensemble(with_copies)), + without_copies = @elapsed(run_ensemble(without_copies))) +``` + +A `prob_func` must only derive its result from the problem it receives. Returning +some other, shared problem defeats the ensemble's copies: + +```julia +# Unsafe with EnsembleThreads: uses the global `jprob`, not the task's copy `prob` +bad_prob_func(prob, ctx) = remake(jprob; p = [0.5 + 0.001 * ctx.sim_id, 0.1]) +``` + +### Stateful callbacks + +Callbacks passed as keyword arguments to `solve` are shared by every trajectory of an +ensemble, since ensembles copy only the problem. A callback that stores its own mutable +state should instead be passed to the `JumpProblem`, so that the ensemble's problem +copies include it, and should reset its state in its `initialize` function, so that +reusing it for the next trajectory starts fresh. For example, a callback that stops a +simulation after a given number of jumps: + +```@example reuse +function stop_after_jumps(n) + njumps = Ref(0) # callback state + reset!(cb, u, t, integrator) = (njumps[] = 0) + function affect!(integrator) + njumps[] += 1 + njumps[] >= n && terminate!(integrator) + end + DiscreteCallback((u, t, integrator) -> true, affect!; initialize = reset!) +end + +birth = MassActionJump([[0 => 1]], [[1 => 1]]; param_idxs = [1]) +birth_prob = DiscreteProblem([0], (0.0, 100.0), [1.0]) +counting_jprob = JumpProblem(birth_prob, Direct(), birth; callback = stop_after_jumps(5)) +esol = solve(EnsembleProblem(counting_jprob), SSAStepper(), EnsembleThreads(); + trajectories = 100) +all(sol.u[end] == [5] for sol in esol.u) +``` + +Passing `callback = stop_after_jumps(5)` to `solve` instead would share one counter +between all threads. + +## Manual threading + +When launching tasks yourself, give each task its own copy of the problem, made once +per task and reused for all of that task's solves: + +```@example reuse +nsims = 1000 +chunks = Iterators.partition(1:nsims, cld(nsims, Threads.nthreads())) +tasks = map(chunks) do seeds + Threads.@spawn begin + local_prob = deepcopy(jprob) # one independent copy per task + [solve(local_prob, SSAStepper(); seed).u[end][1] for seed in seeds] + end +end +final_counts = reduce(vcat, fetch.(tasks)) +length(final_counts) +``` + +Never solve one problem, or problems created from it with `remake`, from several tasks +at the same time. For the same reason, integrators that are alive at the same time need +independent problems: + +```@example reuse +integrator1 = init(deepcopy(jprob), SSAStepper(); seed = 1) +integrator2 = init(deepcopy(jprob), SSAStepper(); seed = 2) +step!(integrator1) +step!(integrator2) +integrator1.t, integrator2.t +``` + +## Distributed ensembles + +`EnsembleDistributed()` sends the problem to each worker. With `safetycopy = false` and +`pmap_batch_size > 1`, which is the default for 200 or more trajectories, a worker may +run several trajectories as concurrent tasks that share its copy of the problem. To be +safe, use `safetycopy = true` or `pmap_batch_size = 1`: + +```julia +using Distributed +esol = solve(EnsembleProblem(jprob; safetycopy = true), SSAStepper(), + EnsembleDistributed(); trajectories = 10_000) +``` + +## Aliasing `SSAStepper` inputs + +Independently of the jump state, `SSAStepper` supports the common SciML `alias` keyword +argument for its `u0`, `p`, and `tstops` inputs. By default it copies `u0`, reuses `p`, +and never modifies a `tstops` array you pass. For example, a callback that modifies +parameters changes the problem's own parameter vector by default. Passing +`alias_p = false` gives the solve its own copy: + +```@example reuse +stop_births = DiscreteCallback((u, t, integrator) -> t == 50.0, + function (integrator) + integrator.p[1] = 0.0 + reset_aggregated_jumps!(integrator) + end) +sol = solve(jprob, SSAStepper(); callback = stop_births, tstops = [50.0], + alias = SciMLBase.DiscreteAliasSpecifier(alias_p = false)) +jprob.prob.p +``` + +Permitting `u0` to be aliased avoids a copy, but the solve then mutates the problem's +`u0`, so the next solve starts from the final state of the previous one: + +```@example reuse +prob_u0 = deepcopy(jprob) +solve(prob_u0, SSAStepper(); seed = 1, alias = SciMLBase.DiscreteAliasSpecifier(alias_u0 = true)) +prob_u0.prob.u0 +``` + +Use an `ODEAliasSpecifier` to also control `tstops`. `alias` never controls the jump +state; even `alias = false` reuses the problem's jump aggregator. + +## SDE and RODE problems + +StochasticDiffEq's SDE and RODE solvers control jump-state copying themselves, through +the `alias_jumps` field of their alias specifier. When it is unspecified they reuse the +jump state on thread 1 and copy it on other threads. To request a copy explicitly: + +```julia +using StochasticDiffEq +sol = solve(sde_jprob, SRIW1(); alias = SciMLBase.SDEAliasSpecifier(alias_jumps = false)) +``` + +The guidance above still applies: do not rely on these copies to make shared problems +safe for concurrent use. + +## Reproducibility with `SortingDirect` + +[`SortingDirect`](@ref) learns a search order for the jumps during a solve and keeps it +for later solves of the same problem (including remade problems), which speeds up many +consecutive simulations. Sampling remains exact, but a given seed reproduces a +trajectory only from the same learned order. To replay a simulation exactly, start from +a copy of a problem that has never been solved: + +```@example reuse +sd_jprob = JumpProblem(dprob, SortingDirect(), maj; save_positions = (false, false)) +fresh = deepcopy(sd_jprob) # never solved +replay(seed) = solve(deepcopy(fresh), SSAStepper(); seed).u +replay(42) == replay(42) +``` + +Seeded ensembles using `SortingDirect` are also affected: exact replay requires the same +starting learned order, seeds, `batch_size`, thread count, and ensemble algorithm, and +serial and threaded runs generally give different trajectories. + +## Summary + +| Context | Recommended pattern | Copies of the jump state | Safe? | +|:-------------------------------------------- |:------------------------------------------------ |:--------------------------------- |:----------------------- | +| Many serial solves of one problem | Reuse the problem | None | Yes | +| Serial parameter sweep | `remake` | None (shared) | Yes | +| `EnsembleSerial`, `safetycopy = false` | Default `prob_func` | None | Yes | +| `EnsembleThreads`, `safetycopy = false` | Default, or a `prob_func` that remakes its input | One per spawned task | Yes | +| Any ensemble, `safetycopy = true` | Any `prob_func` deriving from its input | One per trajectory | Yes | +| `prob_func` returning a shared problem | Avoid | n/a | No | +| Stateful callbacks in ensembles | Pass to `JumpProblem`, reset in `initialize` | Copied with the problem | Yes | +| Stateful callbacks passed to `solve` | Avoid with threads | Shared | No | +| Manual threads or tasks | `deepcopy` once per task | One per task | Yes | +| Simultaneously alive integrators | `deepcopy` each problem | One per integrator | Yes | +| `EnsembleDistributed`, `pmap_batch_size > 1` | `safetycopy = true` or `pmap_batch_size = 1` | One per trajectory, or per worker | Yes with these settings | diff --git a/ext/JumpProcessesKernelAbstractionsExt.jl b/ext/JumpProcessesKernelAbstractionsExt.jl index 342cdf711..5e212792d 100644 --- a/ext/JumpProcessesKernelAbstractionsExt.jl +++ b/ext/JumpProcessesKernelAbstractionsExt.jl @@ -25,6 +25,21 @@ function seed_kernel_backend!(backend, seed, rng, kwargs) return nothing end +# Kernels always build device-owned state, so the removed `alias_jump` keyword and +# `alias` controls are rejected rather than silently ignored. These paths bypass +# SciMLBase's keyword validation, so the never-supported plural spelling is rejected here. +function reject_kernel_alias_kwargs(kwargs) + haskey(kwargs, :alias_jump) && + throw(ArgumentError(JumpProcesses.ALIAS_JUMP_REMOVED_MSG)) + haskey(kwargs, :alias_jumps) && throw(ArgumentError( + "`alias_jumps` is not a keyword argument; it is a field of some alias " * + "specifiers. EnsembleGPUKernel always builds device-owned jump state.")) + get(kwargs, :alias, nothing) === nothing || throw(ArgumentError( + "EnsembleGPUKernel does not support the `alias` keyword with multiple " * + "trajectories, since kernels always build device-owned state.")) + return nothing +end + include("ssa_stepper.jl") include("explicit_tau_leaping.jl") @@ -43,6 +58,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, ensemblealg.backend === nothing ? backend = CPU() : backend = ensemblealg.backend + reject_kernel_alias_kwargs(kwargs) seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob diff --git a/ext/explicit_tau_leaping.jl b/ext/explicit_tau_leaping.jl index 72dffdc63..3cb70a797 100644 --- a/ext/explicit_tau_leaping.jl +++ b/ext/explicit_tau_leaping.jl @@ -251,6 +251,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, end ensemblealg.backend === nothing ? backend = CPU() : backend = ensemblealg.backend + reject_kernel_alias_kwargs(kwargs) seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob diff --git a/ext/ssa_stepper.jl b/ext/ssa_stepper.jl index 98b050fec..1d5750205 100644 --- a/ext/ssa_stepper.jl +++ b/ext/ssa_stepper.jl @@ -320,6 +320,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, would have to run inside the GPU kernel.") ensemblealg.backend === nothing ? backend = CPU() : backend = ensemblealg.backend + reject_kernel_alias_kwargs(kwargs) seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl index bc37376db..be553da70 100644 --- a/src/JumpProcesses.jl +++ b/src/JumpProcesses.jl @@ -32,6 +32,7 @@ import SciMLBase: plot_indices import DataStructures: update! import Graphs: neighbors, outdegree import RecursiveArrayTools: recursivecopy! +using RecursiveArrayTools: RecursiveArrayTools, copyat_or_push!, recursivecopy import SymbolicIndexingInterface as SII # Import additional types and functions from DiffEqBase and SciMLBase diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index 07c87b7d1..35c35c542 100644 --- a/src/SSA_stepper.jl +++ b/src/SSA_stepper.jl @@ -16,6 +16,23 @@ Highly efficient integrator for pure jump problems that involve only `ConstantRa - Supports `rng` and `seed` keyword arguments in `solve`/`init` to control the random number generator used for jump sampling. `rng` accepts any `AbstractRNG`, while `seed` creates a `Xoshiro` generator. `rng` takes priority over `seed`. + - Reuses the `JumpProblem`'s jump state (aggregator and jump callback), re-initializing + it at every `init`; it is never copied. One `JumpProblem` therefore supports one active + solve or integrator at a time. For concurrent solves use independent copies, for + example `deepcopy(jprob)`, or an `EnsembleProblem`. See the `JumpProblem` docstring. + - Supports the common `alias` keyword argument for its `u0`, `p`, and `tstops` inputs, + following SciMLBase's alias-specifier convention: pass `nothing` (the default), a + `Bool` that applies to all three, a `SciMLBase.DiscreteAliasSpecifier` (`alias_u0`, + `alias_p`), or an `ODEAliasSpecifier` (which also has `alias_tstops`). `true` permits + aliasing and `false` requests a copy. The defaults copy `u0` (with `recursivecopy`), + reuse `p`, and never modify a caller's `tstops` array. With `alias_u0 = true` the + integrator mutates the problem's `u0`, so a later solve starts from its final state. + `alias_f` and `alias_du0` have no effect, and `alias` never controls jump state. + - `tstops` may be a number, a tuple, or an `AbstractVector` of times (such as a range or + view), or a callable `(p, tspan) -> times`. Only a `Vector` of the time type can be + aliased; other supported inputs are copied into a new `Vector` at `init`. + - Saved states are independent snapshots of the integrator state made with + `recursivecopy`, as in OrdinaryDiffEq. - As when using jumps with ODEs and SDEs, saving controls for whether to save each time a jump occurs are via the `save_positions` keyword argument to `JumpProblem`. Note that when choosing `SSAStepper` as the timestepper, `save_positions = (true,true)`, `(true,false)`, @@ -167,8 +184,26 @@ function SciMLBase.set_rng!(integrator::SSAIntegrator, rng) nothing end -(integrator::SSAIntegrator)(t) = copy(integrator.u) -(integrator::SSAIntegrator)(out, t) = (out .= integrator.u) +(integrator::SSAIntegrator)(t) = recursivecopy(integrator.u) +# Copy the current state into `out` and return it. Nested arrays are copied recursively, +# like saved states; other states, including scalars, are broadcast into `out`. +function (integrator::SSAIntegrator)(out, t) + u = integrator.u + if u isa AbstractArray && eltype(u) <: AbstractArray + recursivecopy!(out, u) + else + out .= u + end + return out +end + +# Save a snapshot of the current state, matching OrdinaryDiffEq's save semantics +# (`copyat_or_push!` stores a `recursivecopy` of the live state). +@inline function save_current_state!(integrator::SSAIntegrator, t) + push!(integrator.sol.t, t) + copyat_or_push!(integrator.sol.u, length(integrator.sol.u) + 1, integrator.u) + nothing +end # SciMLBase v3 / DiffEqBase v7 renamed the integrator's `u_modified` field to # `derivative_discontinuity` and internal callback code now reads/writes the @@ -224,15 +259,13 @@ function DiffEqBase.solve!(integrator::SSAIntegrator) # Split to help prediction while integrator.cur_saveat <= length(integrator.saveat) && integrator.saveat[integrator.cur_saveat] < integrator.t - push!(integrator.sol.t, integrator.saveat[integrator.cur_saveat]) - push!(integrator.sol.u, copy(integrator.u)) + save_current_state!(integrator, integrator.saveat[integrator.cur_saveat]) integrator.cur_saveat += 1 end end if integrator.save_end && integrator.sol.t[end] != end_time - push!(integrator.sol.t, end_time) - push!(integrator.sol.u, copy(integrator.u)) + save_current_state!(integrator, end_time) end end @@ -273,20 +306,58 @@ function check_continuous_callback_error(callback) return nothing end +# Normalize SSAStepper's `alias` keyword to `(alias_u0, alias_p, alias_tstops)`, each +# `nothing` (solver default), `true` (aliasing permitted), or `false`. `alias_f` and +# `alias_du0` have no effect: SSAStepper never evaluates `f` and has no `du0`. +ssa_alias_choices(::Nothing) = (nothing, nothing, nothing) +ssa_alias_choices(alias::Bool) = (alias, alias, alias) +function ssa_alias_choices(alias::SciMLBase.DiscreteAliasSpecifier) + (alias.alias_u0, alias.alias_p, nothing) +end +function ssa_alias_choices(alias::SciMLBase.ODEAliasSpecifier) + (alias.alias_u0, alias.alias_p, alias.alias_tstops) +end +function ssa_alias_choices(alias) + throw(ArgumentError("SSAStepper's `alias` keyword accepts `nothing`, a `Bool`, a " * + "`SciMLBase.DiscreteAliasSpecifier`, or a " * + "`SciMLBase.ODEAliasSpecifier`, but received a $(typeof(alias)). " * + "`alias` controls `u0`, `p`, and `tstops`; it never controls " * + "jump state.")) +end + +# Returns `(tstops, alias_tstops, copied_tstops)` for the integrator. `alias_tstops` +# means the integrator may insert into `tstops`, and `copied_tstops` means it already +# owns them. Only a caller's `Vector` of the time type can be aliased; other containers +# are materialized, as aliasing is permitted but not required. +function ssa_init_tstops(tstops, ::Type{T}, alias_tstops) where {T} + if tstops isa Vector{T} + if alias_tstops === true + return tstops, true, false + elseif alias_tstops === false + return copy(tstops), true, true + else + # Default: never mutate the caller's array; copy before the first insertion. + return tstops, false, false + end + end + owned = (tstops isa Number) ? T[tstops] : collect(T, tstops) + return owned, true, true +end + function SciMLBase.__init(jump_prob::JumpProblem, alg::SSAStepper; save_start = true, save_end = true, seed = nothing, rng = nothing, - alias_jump = Threads.threadid() == 1, + alias = nothing, + alias_jump = KeywordNotPassed(), saveat = nothing, callback = nothing, tstops = nothing, numsteps_hint = 100) - - # hack until alias system is in place - alias_tstops = false + check_alias_jump_removed(jump_prob, alias_jump) + alias_u0, alias_p, alias_tstops = ssa_alias_choices(alias) if !(jump_prob.prob isa DiscreteProblem) error("SSAStepper only supports DiscreteProblems.") @@ -303,16 +374,18 @@ function SciMLBase.__init(jump_prob::JumpProblem, _rng = resolve_rng(rng, seed) - if alias_jump - cb = jump_prob.jump_callback.discrete_callbacks[end] - else - cb = deepcopy(jump_prob.jump_callback.discrete_callbacks[end]) - end + # The problem's jump state is always reused and re-initialized below; callers + # isolate problems used concurrently (e.g. via `deepcopy` or an `EnsembleProblem`). + cb = jump_prob.jump_callback.discrete_callbacks[end] opts = (callback = CallbackSet(callback),) + # Defaults match OrdinaryDiffEq: copy `u0` and reuse `p`. + u0 = (alias_u0 === true) ? prob.u0 : recursivecopy(prob.u0) + p = (alias_p === false) ? recursivecopy(prob.p) : prob.p + if save_start t = [prob.tspan[1]] - u = [copy(prob.u0)] + u = [recursivecopy(u0)] else t = typeof(prob.tspan[1])[] u = typeof(prob.u0)[] @@ -346,22 +419,23 @@ function SciMLBase.__init(jump_prob::JumpProblem, (tdir <= 0) && error("The time interval to solve over is non-increasing, i.e. tspan[2] <= tspan[1]. This is not allowed for pure jump problem.") - # Stash callable tstops (e.g. SymbolicTstops); use empty vector for init. + # Stash callable tstops (e.g. SymbolicTstops); use an owned empty vector for init. + tType = eltype(prob.tspan) if tstops === nothing - alias_tstops = true callable_tstops = nothing - _tstops = eltype(jump_prob.prob.tspan)[] + _tstops, _alias_tstops, copied_tstops = tType[], true, true elseif tstops isa AbstractArray || tstops isa Tuple || tstops isa Number callable_tstops = nothing - _tstops = tstops + _tstops, _alias_tstops, copied_tstops = ssa_init_tstops(tstops, tType, alias_tstops) else callable_tstops = tstops - _tstops = eltype(jump_prob.prob.tspan)[] + _tstops, _alias_tstops, copied_tstops = tType[], true, true end - integrator = SSAIntegrator(prob.f, copy(prob.u0), prob.tspan[1], prob.tspan[1], tdir, - prob.p, sol, 1, prob.tspan[1], cb, _saveat, save_everystep, - save_end, cur_saveat, opts, _tstops, 1, false, true, alias_tstops, false, _rng) + integrator = SSAIntegrator(prob.f, u0, prob.tspan[1], prob.tspan[1], tdir, + p, sol, 1, prob.tspan[1], cb, _saveat, save_everystep, + save_end, cur_saveat, opts, _tstops, 1, false, true, _alias_tstops, + copied_tstops, _rng) cb.initialize(cb, integrator.u, prob.tspan[1], integrator) DiffEqBase.initialize!(opts.callback, integrator.u, prob.tspan[1], integrator) if save_start @@ -439,8 +513,7 @@ function DiffEqBase.step!(integrator::SSAIntegrator) # Split to help prediction while integrator.cur_saveat <= length(integrator.saveat) && integrator.saveat[integrator.cur_saveat] < integrator.t - push!(integrator.sol.t, integrator.saveat[integrator.cur_saveat]) - push!(integrator.sol.u, copy(integrator.u)) + save_current_state!(integrator, integrator.saveat[integrator.cur_saveat]) integrator.cur_saveat += 1 end end @@ -477,8 +550,7 @@ function DiffEqBase.savevalues!(integrator::SSAIntegrator, force = false) if integrator.save_everystep || force saved = true savedexactly = true - push!(integrator.sol.t, integrator.t) - push!(integrator.sol.u, copy(integrator.u)) + save_current_state!(integrator, integrator.t) end saved, savedexactly diff --git a/src/aggregators/aggregators.jl b/src/aggregators/aggregators.jl index ebaa13fc6..b7f2f32b3 100644 --- a/src/aggregators/aggregators.jl +++ b/src/aggregators/aggregators.jl @@ -48,6 +48,14 @@ The Sorting Direct method. Often the fastest algorithm for smaller to moderate sized systems (tens of jumps), or systems where a few jumps occur much more frequently than others. +The search order learned during a solve is kept in the problem's jump state and is not +reset by `init`, so repeated solves of the same problem (including problems created with +`remake`) start from the order learned so far. This speeds up many consecutive +simulations. Sampling remains exact, but a given seed reproduces a trajectory only from +the same learned order: for exact replay, solve a `deepcopy` of a problem that has never +been solved. Ensemble results additionally depend on the seeds, batching, thread count, +and ensemble algorithm. + J. M. McCollum, G. D. Peterson, C. D. Cox, M. L. Simpson and N. F. Samatova, The sorting direct method for stochastic simulation of biochemical systems with varying reaction execution behavior, Computational Biology and Chemistry, 30 diff --git a/src/problem.jl b/src/problem.jl index 5f987983a..66468b81d 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -71,17 +71,36 @@ page](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_e the DifferentialEquations.jl [docs](https://docs.sciml.ai/JumpProcesses/stable/) for usage examples and commonly asked questions. -!!! warning "Thread Safety" - - `JumpProblem` contains mutable state (aggregator data, callbacks) and is **not - thread-safe**. A single `JumpProblem` instance must not be solved concurrently from - multiple threads or tasks without first creating independent copies via `deepcopy`. - With `EnsembleProblem`, `safetycopy = true` requests a separate problem copy - per trajectory. If using `safetycopy = false`, the caller's - `prob_func` and solver settings must provide the required isolation. CPU ensemble - methods can assign separate RNGs to trajectories; a custom `rng_func` must not - return one shared mutable RNG. Separate RNGs alone do not isolate mutable jump - state, and kernel ensembles have their own backend RNG rules. +!!! warning "Thread safety and problem reuse" + + A `JumpProblem` stores mutable jump state: the jump aggregator and its callbacks. + Solvers owned by JumpProcesses, [`SSAStepper`](@ref) and the OrdinaryDiffEq ODE/DAE + pathways, reuse this state directly and re-initialize it at every `init`; they never + copy it. A `JumpProblem` is therefore **not thread-safe** and supports only one active + solve or integrator at a time: + + - Repeated serial solves of one problem are safe, and avoid copying. + - Concurrent solves from multiple threads or tasks, and integrators that are alive at + the same time, each need an independent copy, such as `deepcopy(jprob)`. + - `remake(jprob; ...)` shares the original's jump state, so remade problems are not + independent for concurrent use. + - `EnsembleProblem` supplies isolation. `EnsembleSerial` reuses the problem + sequentially. `EnsembleThreads` with `safetycopy = false` copies the problem once + per spawned task. `safetycopy = true` copies it before every trajectory, before + `prob_func` is called. A `prob_func` that returns some other, shared problem + defeats these copies, and callbacks passed as `solve` keyword arguments are shared + by all trajectories. + - `EnsembleDistributed` with `pmap_batch_size > 1` and `safetycopy = false` can run + concurrent tasks on one worker's copy of the problem; use `safetycopy = true` or + `pmap_batch_size = 1`. + - StochasticDiffEq SDE/RODE solvers apply their own jump-state copying policy. + + CPU ensemble methods can assign separate RNGs to trajectories; a custom `rng_func` must + not return one shared mutable RNG. Separate RNGs alone do not isolate mutable jump + state, and kernel ensembles have their own backend RNG rules. See the + [ensembles and problem reuse + tutorial](https://docs.sciml.ai/JumpProcesses/stable/tutorials/ensembles_and_problem_reuse/) + for patterns that balance performance and safety. """ mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggregator}, J1, J2, J3, J4, K} <: SciMLBase.AbstractJumpProblem{P, J} @@ -141,6 +160,19 @@ function remake_extended_u0(prob, newu0) end Base.@pure remaker_of(prob::T) where {T <: JumpProblem} = SciMLBase.parameterless_type(T) +""" + remake(jprob::JumpProblem; u0 = missing, p = missing, tspan, prob) + +Construct a `JumpProblem` with a new initial condition `u0`, parameters `p`, time span +`tspan`, or wrapped problem `prob`. + +The new problem shares the original's jump aggregator and jump callbacks rather than +copying them. Solvers owned by JumpProcesses re-initialize this state from the problem being +solved at every `init`, so the original and remade problems can be solved one after +another. They must not be solved concurrently, or used for integrators that are alive at +the same time; use `deepcopy` to obtain an independent problem. `SortingDirect`'s learned +search order is also shared. +""" function DiffEqBase.remake(jprob::JumpProblem; u0 = missing, p = missing, interpret_symbolicmap = true, use_defaults = false, kwargs...) T = remaker_of(jprob) @@ -268,6 +300,7 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS if haskey(kwargs, :useiszero) throw(ArgumentError("`useiszero` is no longer a keyword argument for `JumpProblem`. Set `useiszero` on the `MassActionJump` directly instead.")) end + haskey(kwargs, :alias_jump) && throw(ArgumentError(ALIAS_JUMP_REMOVED_MSG)) # keep the original MAJ (including {Nothing} parameterized MAJs); # fill_scaled_rates! in each aggregator's initialize! handles rate setup @@ -347,6 +380,7 @@ function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; if haskey(kwargs, :useiszero) throw(ArgumentError("`useiszero` is no longer a keyword argument for `JumpProblem`. Set `useiszero` on the `MassActionJump` directly instead.")) end + haskey(kwargs, :alias_jump) && throw(ArgumentError(ALIAS_JUMP_REMOVED_MSG)) # Validate no spatial systems (not currently supported) (spatial_system !== nothing || hopping_constants !== nothing) && diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index f8674702f..4053a91b9 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -1156,6 +1156,13 @@ the backend's API, such as `CUDA.seed!`, before calling `solve`. With `trajectories = 1`, execution delegates to `EnsembleSerial()` on the CPU and accepts its usual `rng`, `seed`, and ensemble `rng_func` keywords. +## Aliasing + +Kernels always build device-owned state. For multiple trajectories, the removed +`alias_jump` keyword raises its migration error, and the `alias_jumps` keyword and any +`alias` value other than `nothing` raise an `ArgumentError`. With `trajectories = 1`, the serial solver's rules +apply: `SSAStepper` supports `alias`, while the tau-leaping solvers reject both keywords. + ## Returns - A `SciMLBase.EnsembleAlgorithm` value for use as the ensemble algorithm argument to diff --git a/src/solve.jl b/src/solve.jl index 77fdcdaa6..58a304654 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -66,26 +66,51 @@ function SciMLBase.__init(_jump_prob::JumpProblem{IIP, P}, __jump_init(_jump_prob, alg; kwargs...) end -function __jump_init(_jump_prob::JumpProblem{IIP, P}, alg; - callback = nothing, seed = nothing, rng = nothing, - alias_jump = Threads.threadid() == 1, - kwargs...) where {IIP, P} - _rng = resolve_rng(rng, seed) +const ALIAS_JUMP_REMOVED_MSG = """ + The `alias_jump` keyword argument was removed in JumpProcesses v10. Solvers \ + owned by JumpProcesses now always reuse, and re-initialize, the `JumpProblem`'s \ + jump state instead of copying it. Remove the keyword. To obtain independent jump \ + state, for example for concurrent solves, solve `deepcopy(jump_prob)` or use an \ + `EnsembleProblem`. See the JumpProcesses documentation on ensembles and problem \ + reuse.""" + +# Default for a removed keyword, so that any supplied value, including `nothing`, is +# rejected. +struct KeywordNotPassed end + +function check_alias_jump_removed(alias_jump) + alias_jump isa KeywordNotPassed || throw(ArgumentError(ALIAS_JUMP_REMOVED_MSG)) + nothing +end - if alias_jump - jump_prob = _jump_prob - else - jump_prob = resetted_jump_problem(_jump_prob) +# Also reject the keyword when it is stored on the wrapped problem, whose keywords are +# not passed through `solve`/`init`. +function check_alias_jump_removed(jump_prob::JumpProblem, alias_jump) + check_alias_jump_removed(alias_jump) + prob = jump_prob.prob + if hasproperty(prob, :kwargs) && haskey(prob.kwargs, :alias_jump) + throw(ArgumentError(ALIAS_JUMP_REMOVED_MSG)) end + nothing +end + +# JumpProcesses-owned initialization never copies jump state: the integrator uses the +# problem's jump callbacks, which are re-initialized by `init`. Callers isolate problems +# that are used concurrently. `alias` is forwarded unchanged to the underlying solver. +function __jump_init(jump_prob::JumpProblem{IIP, P}, alg; + callback = nothing, seed = nothing, rng = nothing, + alias_jump = KeywordNotPassed(), kwargs...) where {IIP, P} + check_alias_jump_removed(jump_prob, alias_jump) + _rng = resolve_rng(rng, seed) init(jump_prob.prob, alg; callback = CallbackSet(jump_prob.jump_callback, callback), rng = _rng, kwargs...) end -# Keep function signatures for StochasticDiffEq backward compatibility. -# The seed argument is accepted but no longer used to reseed aggregator RNGs -# (RNG state is now managed by the integrator). +# Keep function signatures for StochasticDiffEq backward compatibility; JumpProcesses +# itself no longer calls these. The seed argument is accepted but no longer used to +# reseed aggregator RNGs (RNG state is now managed by the integrator). function resetted_jump_problem(_jump_prob, seed = nothing) deepcopy(_jump_prob) end diff --git a/test/explicit_tau_kernel_tests.jl b/test/explicit_tau_kernel_tests.jl index ed8e641a1..b7b8d8a5e 100644 --- a/test/explicit_tau_kernel_tests.jl +++ b/test/explicit_tau_kernel_tests.jl @@ -164,5 +164,6 @@ function run_explicit_tau_kernel_tests(backend, nsims) run_kernel_massaction_rate_tests(backend, SimpleExplicitTauLeaping(), PureLeaping()) rng_prob = leaping_prob([5.0], [Pair{Int, Int}[]], [[1 => 1]], [0.0], (0.0, 4.0)) run_kernel_rng_tests(backend, SimpleExplicitTauLeaping(), rng_prob; saveat = 1.0) + run_kernel_alias_tests(backend, SimpleExplicitTauLeaping(), rng_prob; saveat = 1.0) return nothing end diff --git a/test/jump_state_ownership.jl b/test/jump_state_ownership.jl new file mode 100644 index 000000000..bcc550965 --- /dev/null +++ b/test/jump_state_ownership.jl @@ -0,0 +1,384 @@ +using JumpProcesses, OrdinaryDiffEq, SciMLBase, Test +using StaticArrays: SVector + +# Ownership rules for JumpProcesses v10: JumpProcesses-owned solvers reuse (and +# re-initialize) the problem's jump state instead of copying it, `alias_jump` is +# removed, and SSAStepper supports the common `alias` controls for `u0`, `p` and +# `tstops`. + +function bd_maj() + MassActionJump([[0 => 1], [1 => 1]], [[1 => 1], [1 => -1]]; param_idxs = [1, 2]) +end + +function bd_jprob(; aggregator = Direct(), u0 = [10], p = [1.0, 0.1], + tspan = (0.0, 10.0), kwargs...) + JumpProblem(DiscreteProblem(u0, tspan, p), aggregator, bd_maj(); kwargs...) +end + +const COUPLED_JUMP = ConstantRateJump((u, p, t) -> 1.0, + integrator -> (integrator.u[1] += 1; nothing)) + +function ode_jprob() + oprob = ODEProblem((du, u, p, t) -> (du .= 0; nothing), [10.0], (0.0, 1.0)) + JumpProblem(oprob, Direct(), COUPLED_JUMP) +end + +function uses_problem_callback(integrator, jprob) + integrator.cb.condition === jprob.discrete_jump_aggregation +end +function uses_problem_aggregation(integrator, jprob) + any(cb -> cb.condition === jprob.discrete_jump_aggregation, + integrator.opts.callback.discrete_callbacks) +end + +const REMOVED_MSG = "`alias_jump` keyword argument was removed in JumpProcesses v10" + +# The OrdinaryDiffEq umbrella loads, but does not export, its DAE solvers. +function loaded_module(name) + only(filter(m -> nameof(m) === name, collect(values(Base.loaded_modules)))) +end +const DFBDF = loaded_module(:OrdinaryDiffEqBDF).DFBDF + +# Saved states of the nested test model are independent when no two share an inner array. +independent_snapshots(us) = allunique(objectid(u[1]) for u in us) + +@testset "Removed alias_jump keyword" begin + jprob = bd_jprob() + ojprob = ode_jprob() + dae(out, du, u, p, t) = (out .= du; nothing) + daejprob = JumpProblem(DAEProblem(dae, [0.0], [10.0], (0.0, 1.0)), Direct(), + COUPLED_JUMP) + for value in (true, false) + calls = (() -> solve(jprob, SSAStepper(); alias_jump = value), + () -> init(jprob, SSAStepper(); alias_jump = value), + () -> solve(ojprob, Tsit5(); alias_jump = value), + () -> init(ojprob, Tsit5(); alias_jump = value), + () -> init(daejprob, DFBDF(); alias_jump = value), + () -> bd_jprob(; alias_jump = value)) + for call in calls + @test_throws ArgumentError call() + @test_throws REMOVED_MSG call() + end + end + + # Any supplied value is rejected, including `nothing`. + @test_throws REMOVED_MSG init(jprob, SSAStepper(); alias_jump = nothing) + @test_throws REMOVED_MSG solve(ojprob, Tsit5(); alias_jump = nothing) + @test_throws REMOVED_MSG bd_jprob(; alias_jump = nothing) + + # The keyword is also rejected when stored on the wrapped problem. + stored_dprob = DiscreteProblem([10], (0.0, 1.0), [1.0, 0.1]; alias_jump = false) + stored_jprob = JumpProblem(stored_dprob, Direct(), bd_maj()) + @test_throws REMOVED_MSG init(stored_jprob, SSAStepper()) + @test_throws REMOVED_MSG solve(stored_jprob, SSAStepper()) + stored_oprob = ODEProblem((du, u, p, t) -> (du .= 0; nothing), [10.0], (0.0, 1.0); + alias_jump = false) + stored_ojprob = JumpProblem(stored_oprob, Direct(), COUPLED_JUMP) + @test_throws REMOVED_MSG init(stored_ojprob, Tsit5()) + @test_throws REMOVED_MSG solve(stored_ojprob, Tsit5()) + + # The plural spelling was never a solve keyword; SciMLBase's keyword validation + # rejects it before JumpProcesses code runs. + @test_throws SciMLBase.CommonKwargError solve(jprob, SSAStepper(); alias_jumps = false) + @test_throws SciMLBase.CommonKwargError init(jprob, SSAStepper(); alias_jumps = false) + @test_throws SciMLBase.CommonKwargError solve(ojprob, Tsit5(); alias_jumps = false) + @test_throws SciMLBase.CommonKwargError init(ojprob, Tsit5(); alias_jumps = false) +end + +@testset "CPU tau-leaping solvers reject alias keywords" begin + # These solvers build fresh working state, so neither keyword applies; their strict + # keyword signatures reject both rather than silently ignoring them. + leaping = JumpProblem(DiscreteProblem([10], (0.0, 1.0), [1.0, 0.1]), PureLeaping(), + bd_maj()) + for alg in (SimpleExplicitTauLeaping(), SimpleImplicitTauLeaping(), + SimpleTrapezoidalLeaping(), SimpleAdaptiveTauLeaping()) + @test_throws MethodError solve(leaping, alg; alias_jump = true) + @test_throws MethodError solve(leaping, alg; alias = true) + @test SciMLBase.successful_retcode(solve(leaping, alg; seed = 1)) + end +end + +@testset "JumpProcesses-owned solvers reuse the problem's jump state" begin + jprob = bd_jprob() + @test uses_problem_callback(init(jprob, SSAStepper()), jprob) + @test uses_problem_callback(fetch(Threads.@spawn init(jprob, SSAStepper())), jprob) + @test uses_problem_callback(init(jprob, SSAStepper(); alias = false), jprob) + + ojprob = ode_jprob() + @test uses_problem_aggregation(init(ojprob, Tsit5()), ojprob) + @test uses_problem_aggregation(fetch(Threads.@spawn init(ojprob, Tsit5())), ojprob) + + # `alias` is forwarded unchanged to the coupled solver. + integrator = init(ojprob, Tsit5(); alias = ODEAliasSpecifier(alias_u0 = true)) + @test integrator.u === ojprob.prob.u0 + @test uses_problem_aggregation(integrator, ojprob) +end + +@testset "Reuse re-initializes jump state from the current state and parameters" begin + jprob = bd_jprob() + aggregation = jprob.discrete_jump_aggregation + sol = solve(jprob, SSAStepper(); seed = 1) + @test SciMLBase.successful_retcode(sol) + @test length(sol.t) > 2 + + # Default `alias_p` reuses the problem's parameters, so in-place parameter changes + # are seen by the next initialization of the same jump state. + integrator = init(jprob, SSAStepper(); seed = 2) + @test integrator.u == [10] + @test aggregation.sum_rate ≈ 1.0 + 0.1 * 10 + jprob.prob.p[1] = 3.0 + integrator = init(jprob, SSAStepper(); seed = 2) + @test aggregation.maj_rates == [3.0, 0.1] + @test aggregation.sum_rate ≈ 3.0 + 0.1 * 10 +end + +@testset "remake shares jump state with the original problem" begin + jprob = bd_jprob() + remade = remake(jprob; p = [2.0, 0.2]) + @test remade.discrete_jump_aggregation === jprob.discrete_jump_aggregation + @test remade.jump_callback === jprob.jump_callback + + # Serial use stays correct: every init re-initializes the shared state from the + # problem being solved. + aggregation = jprob.discrete_jump_aggregation + init(remade, SSAStepper()) + @test aggregation.sum_rate ≈ 2.0 + 0.2 * 10 + init(jprob, SSAStepper()) + @test aggregation.sum_rate ≈ 1.0 + 0.1 * 10 +end + +@testset "SSAStepper alias controls for u0 and p" begin + specs = Any[nothing, true, false] + for alias_u0 in (nothing, true, false), alias_p in (nothing, true, false) + + push!(specs, ODEAliasSpecifier(; alias_u0, alias_p)) + push!(specs, SciMLBase.DiscreteAliasSpecifier(; alias_u0, alias_p)) + end + choice(spec::Union{Nothing, Bool}, field) = spec + choice(spec, field) = getfield(spec, field) + + for spec in specs + jprob = bd_jprob() + u0, p = jprob.prob.u0, jprob.prob.p + integrator = init(jprob, SSAStepper(); alias = spec, seed = 3) + u0_aliased = choice(spec, :alias_u0) === true + p_aliased = choice(spec, :alias_p) !== false + @test (integrator.u === u0) == u0_aliased + @test (integrator.p === p) == p_aliased + @test integrator.p == p + # `alias` never controls jump state. + @test uses_problem_callback(integrator, jprob) + + solve!(integrator) + @test integrator.sol.u[1] == [10] + @test u0 == (u0_aliased ? integrator.u : [10]) + end + + @test_throws ArgumentError init(bd_jprob(), SSAStepper(); + alias = SciMLBase.SDEAliasSpecifier()) +end + +@testset "Parameter-changing callbacks use the selected parameters" begin + for alias_p in (nothing, true, false) + jprob = bd_jprob() + p = jprob.prob.p + stop_jumps = DiscreteCallback((u, t, integrator) -> t == 1.0, + function (integrator) + integrator.p[1] = 0.0 + integrator.p[2] = 0.0 + reset_aggregated_jumps!(integrator) + end) + sol = solve(jprob, SSAStepper(); alias = ODEAliasSpecifier(; alias_p), + callback = stop_jumps, tstops = [1.0], seed = 4) + @test SciMLBase.successful_retcode(sol) + @test all(==(sol(1.0)), (sol(t) for t in 1.0:0.5:10.0)) + @test p == (alias_p === false ? [1.0, 0.1] : [0.0, 0.0]) + end +end + +@testset "SSAStepper tstops containers and aliasing" begin + for tstops in ([2.0, 5.0], (2.0, 5.0), 2.0:3.0:5.0, [2, 5], view([0.0, 2.0, 5.0], 2:3)) + integrator = init(bd_jprob(), SSAStepper(); tstops) + @test integrator.tstops isa Vector{Float64} + add_tstop!(integrator, 3.0) + @test integrator.tstops == [2.0, 3.0, 5.0] + solve!(integrator) + @test SciMLBase.successful_retcode(integrator.sol) + end + integrator = init(bd_jprob(), SSAStepper(); tstops = 4.0) + add_tstop!(integrator, 6.0) + @test integrator.tstops == [4.0, 6.0] + + for alias in (nothing, true, false, ODEAliasSpecifier(alias_tstops = true), + ODEAliasSpecifier(alias_tstops = false), SciMLBase.DiscreteAliasSpecifier()) + user = [2.0, 5.0] + integrator = init(bd_jprob(), SSAStepper(); tstops = user, alias) + aliased = (alias === true) || + (alias isa ODEAliasSpecifier && alias.alias_tstops === true) + copied_at_init = (alias === false) || + (alias isa ODEAliasSpecifier && alias.alias_tstops === false) + add_tstop!(integrator, 3.0) + @test integrator.tstops == [2.0, 3.0, 5.0] + @test user == (aliased ? [2.0, 3.0, 5.0] : [2.0, 5.0]) + @test (integrator.tstops === user) == aliased + @test !copied_at_init || integrator.tstops !== user + end + + # Without aliasing permission the caller's array is never mutated, and an + # explicit `false` copies at `init` so later caller changes are not seen. + user = [2.0, 5.0] + integrator = init(bd_jprob(), SSAStepper(); tstops = user, + alias = ODEAliasSpecifier(alias_tstops = false)) + user[2] = 7.0 + @test integrator.tstops == [2.0, 5.0] + + # Callable tstops are evaluated at init and unaffected by alias controls. + integrator = init(bd_jprob(), SSAStepper(); tstops = (p, tspan) -> [3.0], + alias = true) + @test integrator.tstops == [3.0] +end + +@testset "Saved states are independent recursive snapshots" begin + # A nested state type, where a shallow `copy` would share inner arrays. + jump = ConstantRateJump((u, p, t) -> 1.0, + integrator -> (integrator.u[1][1] += 1; nothing)) + nested_prob(save_positions) = JumpProblem( + DiscreteProblem([[0], [0]], (0.0, 10.0)), Direct(), jump; save_positions) + + for alias_u0 in (nothing, true, false) + jprob = nested_prob((false, true)) + sol = solve(jprob, SSAStepper(); saveat = 2.0, seed = 5, + alias = ODEAliasSpecifier(; alias_u0)) + counts = [u[1][1] for u in sol.u] + @test counts[1] == 0 + @test issorted(counts) + @test length(unique(counts)) > 1 + @test independent_snapshots(sol.u) + + # Every-step saves and `save_end`. + jprob = nested_prob((false, true)) + sol = solve(jprob, SSAStepper(); seed = 6, alias = ODEAliasSpecifier(; alias_u0)) + counts = [u[1][1] for u in sol.u] + # One save per jump, plus a final `save_end` save if the last jump came earlier. + jump_counts = counts[end] == counts[end - 1] ? counts[1:(end - 1)] : counts + @test jump_counts == 0:(length(jump_counts) - 1) + @test independent_snapshots(sol.u) + end + + # Forced saves from a callback and the integrator's interpolation call. + jprob = nested_prob((false, false)) + forced = DiscreteCallback((u, t, integrator) -> t == 5.0, + integrator -> savevalues!(integrator, true)) + integrator = init(jprob, SSAStepper(); callback = forced, tstops = [5.0], seed = 7) + solve!(integrator) + @test 5.0 in integrator.sol.t + @test independent_snapshots(integrator.sol.u) + snapshot = integrator(integrator.t) + @test snapshot == integrator.u + @test snapshot[1] !== integrator.u[1] + out = [[-1], [-1]] + integrator(out, integrator.t) + @test out == integrator.u + @test out[1] !== integrator.u[1] +end + +@testset "In-place state evaluation returns the destination" begin + # Scalar states are broadcast into the destination. + scalar_jump = ConstantRateJump((u, p, t) -> 1.0, + integrator -> (integrator.u += 1; nothing)) + scalar_jprob = JumpProblem(DiscreteProblem(10, (0.0, 1.0)), Direct(), scalar_jump) + integrator = init(scalar_jprob, SSAStepper(); seed = 1) + out = [0] + @test integrator(out, integrator.t) === out + @test out == [10] + @test integrator(integrator.t) == 10 + + # Flat states are broadcast, converting the element type as before. + integrator = init(bd_jprob(), SSAStepper(); seed = 1) + out = zeros(1) + @test integrator(out, integrator.t) === out + @test out == [10.0] + + # Arrays of static arrays are copied element-wise. + no_jump = ConstantRateJump((u, p, t) -> 0.0, integrator -> nothing) + static_jprob = JumpProblem(DiscreteProblem([SVector(1, 2)], (0.0, 1.0)), Direct(), + no_jump) + integrator = init(static_jprob, SSAStepper(); seed = 1) + out = [SVector(0, 0)] + @test integrator(out, integrator.t) === out + @test out == [SVector(1, 2)] +end + +# Record which jump aggregator each trajectory uses. The affect and the record live at +# module scope, so copying a problem does not copy them. +const USED_CALLBACKS = Set{UInt}() +const USED_CALLBACKS_LOCK = ReentrantLock() +function recording_affect!(integrator) + integrator.u[1] += 1 + lock(() -> push!(USED_CALLBACKS, objectid(integrator.cb.condition)), USED_CALLBACKS_LOCK) + nothing +end + +function recording_jprob() + jump = ConstantRateJump((u, p, t) -> 10.0, recording_affect!) + JumpProblem(DiscreteProblem([0], (0.0, 1.0)), Direct(), jump) +end + +function used_callbacks(f) + lock(() -> empty!(USED_CALLBACKS), USED_CALLBACKS_LOCK) + sol = f() + @test all(SciMLBase.successful_retcode, sol.u) + lock(() -> copy(USED_CALLBACKS), USED_CALLBACKS_LOCK) +end + +@testset "Ensemble copies of jump state" begin + trajectories = 64 + remake_func(prob, ctx) = remake(prob; u0 = [ctx.sim_id]) + for prob_func in (nothing, remake_func) + jprob = recording_jprob() + original = objectid(jprob.discrete_jump_aggregation) + ensemble(safetycopy) = prob_func === nothing ? + EnsembleProblem(jprob; safetycopy) : + EnsembleProblem(jprob; prob_func, safetycopy) + + # EnsembleSerial without safety copies reuses the original problem. + used = used_callbacks(() -> solve(ensemble(false), SSAStepper(), + EnsembleSerial(); trajectories, seed = 1)) + @test used == Set([original]) + + # Safety copies give every trajectory its own copy, made before `prob_func`. + used = used_callbacks(() -> solve(ensemble(true), SSAStepper(), + EnsembleSerial(); trajectories, seed = 1)) + @test length(used) == trajectories + @test original ∉ used + + # EnsembleThreads copies once per spawned task, and JumpProcesses adds no + # copies; with one thread it falls back to the serial path. + used = used_callbacks(() -> solve(ensemble(false), SSAStepper(), + EnsembleThreads(); trajectories, seed = 1)) + if Threads.nthreads() == 1 + @test used == Set([original]) + else + @test 1 <= length(used) <= Threads.nthreads() + @test original ∉ used + end + + used = used_callbacks(() -> solve(ensemble(true), SSAStepper(), + EnsembleThreads(); trajectories, seed = 1)) + @test length(used) == trajectories + @test original ∉ used + end +end + +@testset "SortingDirect ensembles replay from identical learned state" begin + paths(sol) = [(trajectory.t, trajectory.u) for trajectory in sol.u] + template = bd_jprob(; aggregator = SortingDirect(), u0 = [50], p = [5.0, 0.1]) + for ensemblealg in (EnsembleSerial(), EnsembleThreads()) + # Each replay starts from an independent copy of the same never-solved problem, + # with the same seed, batching, and thread count. + run(seed) = solve(EnsembleProblem(deepcopy(template)), SSAStepper(), + ensemblealg; trajectories = 32, batch_size = 16, seed) + @test paths(run(11)) == paths(run(11)) + @test paths(run(11)) != paths(run(12)) + end +end diff --git a/test/kernel_contract_tests.jl b/test/kernel_contract_tests.jl index 170106187..aeb135010 100644 --- a/test/kernel_contract_tests.jl +++ b/test/kernel_contract_tests.jl @@ -56,6 +56,50 @@ function run_kernel_rng_tests(backend, alg, jump_prob; kwargs...) end end +function run_kernel_alias_tests(backend, alg, jump_prob; kwargs...) + ensemble_prob = EnsembleProblem(jump_prob) + paths(sol) = [(traj.t, traj.u) for traj in sol.u] + + @testset "Kernel alias inputs: $(typeof(alg))" begin + # Kernels always build device-owned state, so the removed `alias_jump` + # keyword and `alias` choices are rejected rather than silently ignored. + # These paths bypass SciMLBase's keyword validation, so the never-supported + # plural `alias_jumps` is rejected explicitly as well. + for options in ((; alias_jump = true), (; alias_jump = false), + (; alias_jump = nothing), (; alias_jumps = false), (; alias = true), + (; alias = false), (; alias = SciMLBase.ODEAliasSpecifier(alias_u0 = true))) + @test_throws ArgumentError solve(ensemble_prob, alg, + EnsembleGPUKernel(backend); trajectories = 16, kwargs..., options...) + end + @test_throws "`alias_jump` keyword argument was removed" solve(ensemble_prob, + alg, EnsembleGPUKernel(backend); trajectories = 16, kwargs..., + alias_jump = true) + + if backend isa CPU + sol = solve(ensemble_prob, alg, EnsembleGPUKernel(backend); + trajectories = 16, kwargs..., alias = nothing) + @test all(SciMLBase.successful_retcode, sol.u) + end + + # A single trajectory delegates to the serial solver, whose rules apply: + # SSAStepper raises the migration error and supports `alias`, while the + # tau-leaping solvers' strict keyword signatures reject both keywords. + single(; options...) = solve(ensemble_prob, alg, EnsembleGPUKernel(backend); + trajectories = 1, seed = 123, kwargs..., options...) + if alg isa SSAStepper + @test_throws ArgumentError single(; alias_jump = true) + # `alias = false` keeps the problem's `u0` unmodified between the solves. + kernel = single(; alias = false) + serial = solve(ensemble_prob, alg, EnsembleSerial(); trajectories = 1, + seed = 123, kwargs..., alias = false) + @test paths(kernel) == paths(serial) + else + @test_throws MethodError single(; alias_jump = true) + @test_throws MethodError single(; alias = true) + end + end +end + function run_kernel_massaction_rate_tests(backend, alg, aggregator) ext = Base.get_extension(JumpProcesses, :JumpProcessesKernelAbstractionsExt) @@ -124,4 +168,5 @@ function run_regular_kernel_rng_tests(backend) rj = RegularJump(birth_rate!, birth_c!, 1) jp = JumpProblem(DiscreteProblem([0.0], (0.0, 2.0), (20.0,)), PureLeaping(), rj) run_kernel_rng_tests(backend, SimpleTauLeaping(), jp; dt = 0.25) + run_kernel_alias_tests(backend, SimpleTauLeaping(), jp; dt = 0.25) end diff --git a/test/massaction_state.jl b/test/massaction_state.jl index ea4a7e489..7333367ff 100644 --- a/test/massaction_state.jl +++ b/test/massaction_state.jl @@ -10,7 +10,7 @@ using JumpProcesses, Test jprob = JumpProblem(prob, algorithm, maj) # A parameter change after construction must be seen by initialization. prob.p[1] = 12.0 - integrator = init(jprob, SSAStepper(); seed = 12345, alias_jump = true) + integrator = init(jprob, SSAStepper(); seed = 12345) @test jprob.discrete_jump_aggregation.maj_rates == [2.0] integrator.p[1] = 18.0 reset_aggregated_jumps!(integrator) @@ -18,7 +18,7 @@ using JumpProcesses, Test @test jprob.massaction_jump.scaled_rates === nothing remade = remake(jprob; p = [24.0]) - init(remade, SSAStepper(); seed = 12345, alias_jump = true) + init(remade, SSAStepper(); seed = 12345) @test remade.discrete_jump_aggregation.maj_rates == [4.0] @test jprob.prob.p == [18.0] @test remade.massaction_jump === jprob.massaction_jump @@ -32,7 +32,7 @@ end jprob = JumpProblem(prob, algorithm, maj; spatial_system = CartesianGrid((2,)), hopping_constants = zeros(2)) prob.p[1] = 12.0 - integrator = init(jprob, SSAStepper(); seed = 12345, alias_jump = true) + integrator = init(jprob, SSAStepper(); seed = 12345) rx_rates = jprob.discrete_jump_aggregation.rx_rates @test rx_rates.maj_rates == [2.0] @test rx_rates.rates == [240.0 1008.0] @@ -43,7 +43,7 @@ end @test jprob.massaction_jump.scaled_rates === nothing remade = remake(jprob; p = [24.0]) - init(remade, SSAStepper(); seed = 12345, alias_jump = true) + init(remade, SSAStepper(); seed = 12345) @test remade.discrete_jump_aggregation.rx_rates.maj_rates == [4.0] @test remade.discrete_jump_aggregation.rx_rates.rates == [480.0 2016.0] @test jprob.prob.p == [18.0] @@ -62,7 +62,7 @@ end spatial_system = CartesianGrid((2,)), hopping_constants = zeros(2, 2)) @test jprob.massaction_jump.scaled_rates == [0.0, 0.0, 0.0, 0.0, 2.0, 2.0] @test jprob.prob.u0 == [6, 0, 9, 0] - init(jprob, SSAStepper(); seed = 12345, alias_jump = true) + init(jprob, SSAStepper(); seed = 12345) @test jprob.discrete_jump_aggregation.maj_rates == jprob.massaction_jump.scaled_rates if parameterized diff --git a/test/runtests.jl b/test/runtests.jl index 18b26382f..27cbdb4e5 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -46,6 +46,7 @@ end @time @safetestset "Save_positions test" begin include("save_positions.jl") end @time @safetestset "RNG kwarg tests" begin include("rng_kwarg_tests.jl") end @time @safetestset "Stochastic dispatch tests" begin include("stochastic_dispatch.jl") end + @time @safetestset "Jump-state ownership and alias controls" begin include("jump_state_ownership.jl") end @time @safetestset "Ensemble Uniqueness test" begin include("ensemble_uniqueness.jl") end @time @safetestset "Ensemble Problem Tests" begin include("ensemble_problems.jl") end @time @safetestset "A + B <--> C" begin include("reversible_binding.jl") end @@ -67,6 +68,7 @@ end if GROUP == "All" || GROUP == "ThreadSafety" @time @safetestset "Thread Safety test (threaded)" begin include("thread_safety.jl") end + @time @safetestset "Jump-state ownership (threaded)" begin include("jump_state_ownership.jl") end end if GROUP == "CUDA" diff --git a/test/saveat_regression.jl b/test/saveat_regression.jl index e6f544a7a..486a09a95 100644 --- a/test/saveat_regression.jl +++ b/test/saveat_regression.jl @@ -75,10 +75,10 @@ end prob = DiscreteProblem([0], (0.0, 1.0)) jump = ConstantRateJump(rate, affect!) jprob = JumpProblem(prob, Direct(), jump; save_positions) - # Each solve owns a fresh problem and RNG; keep access to its counters - # even when this test is invoked on a worker thread. - sol = solve( - jprob, SSAStepper(); saveat, save_start, save_end, rng, alias_jump = true) + # Each solve owns a fresh problem and RNG. SSAStepper reuses the problem's jump + # state rather than copying it, so the captured event log and RNG counters + # observe every draw, even when this test runs on a worker thread. + sol = solve(jprob, SSAStepper(); saveat, save_start, save_end, rng) sol, events, (rng.exponential_draws, rng.uniform_draws) end diff --git a/test/sortingdirect_reuse.jl b/test/sortingdirect_reuse.jl index a34527a34..ff3d6551d 100644 --- a/test/sortingdirect_reuse.jl +++ b/test/sortingdirect_reuse.jl @@ -26,21 +26,21 @@ end @testset "SortingDirect preserves learned search order" begin @testset "$representation" for representation in (:parameterized_massaction, :fixed_massaction, :constant) jprob = sortingdirect_reuse_problem(representation) - sol = solve(jprob, SSAStepper(); seed = 12345, alias_jump = true) + sol = solve(jprob, SSAStepper(); seed = 12345) @test SciMLBase.successful_retcode(sol) @test sum(last(sol.u)) > 0 # Seeded replay requires identical initial search order, as well as RNG # state. Each copy begins with the order learned by the completed solve. - replay = solve(deepcopy(jprob), SSAStepper(); seed = 12345, alias_jump = true) - repeated = solve(deepcopy(jprob), SSAStepper(); seed = 12345, alias_jump = true) + replay = solve(deepcopy(jprob), SSAStepper(); seed = 12345) + repeated = solve(deepcopy(jprob), SSAStepper(); seed = 12345) @test replay.t == repeated.t @test replay.u == repeated.u # Initialization and rate resets retain the ordering optimization. aggregation = jprob.discrete_jump_aggregation aggregation.jump_search_order .= [3, 1, 2] - integrator = init(jprob, SSAStepper(); seed = 12345, alias_jump = true) + integrator = init(jprob, SSAStepper(); seed = 12345) @test aggregation.jump_search_order == [3, 1, 2] if representation !== :fixed_massaction integrator.p .= [3.0, 4.0, 5.0] diff --git a/test/ssa_kernel_tests.jl b/test/ssa_kernel_tests.jl index f778b1fe0..c0a46d3f5 100644 --- a/test/ssa_kernel_tests.jl +++ b/test/ssa_kernel_tests.jl @@ -206,6 +206,7 @@ function run_ssa_kernel_tests(backend, nsims) run_kernel_massaction_rate_tests(backend, SSAStepper(), Direct()) rng_prob = majump_prob([5.0], [Pair{Int, Int}[]], [[1 => 1]], [0], (0.0, 4.0)) run_kernel_rng_tests(backend, SSAStepper(), rng_prob; saveat = 1.0) + run_kernel_alias_tests(backend, SSAStepper(), rng_prob; saveat = 1.0) run_regular_kernel_rng_tests(backend) return nothing end From b7ad542e8e2d5cf68d8295ba08f174e98d437fc3 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Fri, 9 Oct 2026 18:16:57 -0400 Subject: [PATCH 38/52] Validate SSAIntegrator evaluation times and add SSAStepper save_uprev Evaluating an SSAStepper integrator, as integrator(t) or integrator(out, t), ignored t and returned the current state. That was silently wrong for callbacks that evaluate earlier times, such as a SavingCallback or FunctionCallingCallback with save times, and the in-place form errored because get_tmp_cache was missing. Evaluation now returns the current state at integrator.t and raises an informative error at other times, unless the problem is solved with the new SSAStepper(; save_uprev = true). That option keeps the state at the start of each step in integrator.uprev, without allocating per step, so the integrator can be evaluated anywhere in [tprev, t], and it provides get_tmp_cache. SSAStepper is now parametric; SSAStepper() is unchanged. solve! counts its final advance to the end time as a step, and isdenseplot accepts both parameter values. Also: - save pending saveat times before solve!'s final callback pass, so an end-time callback no longer alters them or puts the saved times out of order; - reject alias_jump, alias_jumps, and alias stored on the wrapped problem, as well as save_uprev = true, in multi-trajectory EnsembleGPUKernel solves; - document the supported SSAStepper state types, the evaluation rules, and the tstops alternative in the docstrings, the solver page, a new tutorial section, the FAQ, the migration guide, and HISTORY. Validation on Julia 1.12.7 with the pinned StochasticDiffEqCore: InterfaceI (1,751 assertions), InterfaceII (3,337), four-thread ThreadSafety (333), strict QA (20, plus 5 pin checks), formatting, spelling and whitespace checks, the documentation blocks, and a Documenter build of the changed pages pass. Default-path benchmark allocations are unchanged. Co-Authored-By: Claude Opus 5.5 (1M context) --- HISTORY.md | 32 +- docs/Project.toml | 2 + docs/src/api.md | 3 +- docs/src/faq.md | 43 +- docs/src/jump_solve.md | 29 ++ docs/src/migration.md | 25 ++ .../tutorials/discrete_stochastic_example.md | 95 ++++ ext/JumpProcessesKernelAbstractionsExt.jl | 28 +- ext/explicit_tau_leaping.jl | 2 +- ext/ssa_stepper.jl | 6 +- src/SSA_stepper.jl | 202 ++++++++- src/aggregators/aggregators.jl | 11 +- src/jumps.jl | 6 + src/simple_regular_solve.jl | 7 +- test/kernel_contract_tests.jl | 27 ++ test/runtests.jl | 1 + test/saveat_regression.jl | 29 ++ test/ssa_integrator_evaluation.jl | 424 ++++++++++++++++++ 18 files changed, 928 insertions(+), 44 deletions(-) create mode 100644 test/ssa_integrator_evaluation.jl diff --git a/HISTORY.md b/HISTORY.md index d033fafe4..a53264b07 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -50,7 +50,8 @@ changes are not intended for a 9.x release. ensembles and problem reuse for patterns that balance performance and safety. - **Breaking**: The `alias_jump` keyword argument has been removed. Passing it to `solve`, `init`, `JumpProblem`, or a multi-trajectory `EnsembleGPUKernel` - solve raises an error; remove it, and solve a copy of the problem when + solve, or storing it on the problem wrapped by a `JumpProblem`, raises an + error; remove it, and solve a copy of the problem when independent jump state is needed. StochasticDiffEq SDE/RODE solvers ignore the keyword and keep their own policy: pass `alias = SciMLBase.SDEAliasSpecifier(; alias_jumps = false)` (or @@ -67,9 +68,34 @@ changes are not intended for a 9.x release. independent snapshots. `tstops` given as a number, a tuple, or an `AbstractVector` other than a `Vector` of the time type (such as a range or view) is copied into a `Vector` at `init`, so `add_tstop!` works with it. + - **Breaking**: Evaluating an `SSAStepper` integrator, as `integrator(t)` or + `integrator(out, t)`, at a time other than `integrator.t` now raises an error, + unless the problem is solved with `SSAStepper(; save_uprev = true)`, which also + allows times in the current step, `[integrator.tprev, integrator.t]`. + Previously the current state was returned for any `t`, which was silently wrong + for times before the most recent event, for example in a `SavingCallback` or + `FunctionCallingCallback` with save times. Alternatively, pass such times as + `tstops`, so that the integrator stops exactly at them. DiffEqCallbacks' + integrating callbacks are not yet supported with `SSAStepper`. + - New `SSAStepper(; save_uprev = true)` option: the integrator keeps the state at + the start of each step in `integrator.uprev`, copying the state on every step, + and provides `get_tmp_cache` for callbacks that evaluate it in place. The + default, `SSAStepper()`, keeps no copy. `SSAStepper` is now a parametric type; + `SSAStepper()` is unchanged and has type `SSAStepper{false}`. After `solve!`, + `integrator.tprev` is the time of the last event or tstop. + - Fixed `SSAStepper` saving pending `saveat` times before the end time after its + final callback pass. A callback that changed the state at the end time altered the + values saved at those earlier times, and the saved times came out of order. They + are now saved before the final callback pass. + - The state types `SSAStepper` supports are now documented: vectors and + `SVector`s of integers or floats, species × sites matrices for spatial + problems, and scalars for models without mass-action jumps. See "Supported + state types" in the solver documentation. - **Breaking**: Multi-trajectory `EnsembleGPUKernel` solves reject `alias` values - other than `nothing`, and the `alias_jumps` keyword, since kernels always build - device-owned state. + other than `nothing` and the `alias_jumps` keyword, whether passed to `solve` or + stored on the problem, since kernels always build device-owned state. They + also reject `SSAStepper(; save_uprev = true)`, since kernels never expose an + integrator. - The minimum supported SciMLBase version is now 3.34. - `SSAIntegrator` now supports the `SciMLBase` RNG interface (`has_rng`, `get_rng`, `set_rng!`). diff --git a/docs/Project.toml b/docs/Project.toml index c29a7a37d..f41c639bb 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,6 +1,7 @@ [deps] Catalyst = "479239e8-5488-4da2-87a7-35f2df7eef83" DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" +DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" @@ -18,6 +19,7 @@ StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" [compat] Catalyst = "16" DensityInterface = "0.4" +DiffEqCallbacks = "4.17.0" DifferentialEquations = "7.11, 8.0" Distributions = "0.25" Documenter = "1.1" diff --git a/docs/src/api.md b/docs/src/api.md index 2ac1a58e8..be77b9c7e 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -196,7 +196,8 @@ ignores the removed `alias_jump` keyword. Multi-trajectory [`EnsembleGPUKernel`](@ref) solves always build device-owned state. They reject the removed `alias_jump` keyword and any `alias` value other -than `nothing`. With `trajectories = 1`, the serial solver's rules apply. +than `nothing`, whether passed to `solve` or stored on the problem wrapped by the +`JumpProblem`. With `trajectories = 1`, the serial solver's rules apply. ## Random Number Generator Control diff --git a/docs/src/faq.md b/docs/src/faq.md index 1c1c3b7ff..00c7008e4 100644 --- a/docs/src/faq.md +++ b/docs/src/faq.md @@ -190,10 +190,34 @@ of `u[1]`, giving When using an ODE or SDE time-stepper that conforms to the [integrator interface](https://docs.sciml.ai/DiffEqDocs/stable/basics/integrator/), one -can simply use `integrator.uprev`. For efficiency reasons, the pure jump -[`SSAStepper`](@ref) integrator does not have such a field. If one needs -solution components at earlier times, one can save them within the callback -condition by making a functor: +can simply use `integrator.uprev`. The pure jump [`SSAStepper`](@ref) provides +the same field when the problem is solved with `SSAStepper(; save_uprev = true)`: +`integrator.uprev` then holds the state at the start of the current step, at +time `integrator.tprev`, and the integrator can be evaluated anywhere in +`[integrator.tprev, integrator.t]`. For example, to remove one unit of `u[4]` +whenever `u[2]` decreases: + +```julia +condition(u, t, integrator) = u[2] < integrator.uprev[2] +function affect!(integrator) + integrator.u[4] -= 1 + reset_aggregated_jumps!(integrator) + nothing +end +sol = solve(jprob, SSAStepper(; save_uprev = true); + callback = DiscreteCallback(condition, affect!)) +``` + +The option also enables `get_tmp_cache`, which callbacks that evaluate the +integrator in place need, such as a `SavingCallback` or `FunctionCallingCallback` +with save times for a vector state. Keeping `uprev` copies the state on every +step, so `SSAStepper()` does not; see [Saving with callbacks and evaluating the +integrator](@ref ssa_integrator_evaluation) for the trade-offs, and for the +alternative of passing the save times as `tstops`, which needs neither. + +If only a few components are needed, storing them in the callback avoids copying +the whole state on every step. For example, a functor can record the previous +value of `u[2]`: ```julia # stores the previous value of u[2] and represents the callback functions @@ -216,9 +240,20 @@ end # affect! function (upc::UprevCondition)(integrator) integrator.u[4] -= 1 + reset_aggregated_jumps!(integrator) nothing end upc = UprevCondition(u0[2]) cb = DiscreteCallback(upc, upc) ``` + +## Which state types can I use with `SSAStepper`? + +`SSAStepper` fully supports a `Vector` or an `SVector` of integers or floats, and, +for spatial problems, a species × sites `Matrix{Int}`, which `NSM` and +`DirectCRDirect` keep and other aggregators flatten into a `Vector`. Models without `MassActionJump`s can also use an integer or float +scalar state, with every non-spatial aggregator except `RSSA` and `RSSACR`. +Bounded `VariableRateJump`s require `Coevolve`. Other state types, such as nested +arrays, are not supported, even when a model built from user-written jumps appears +to work. See [Supported state types](@ref ssa_state_types) for details. diff --git a/docs/src/jump_solve.md b/docs/src/jump_solve.md index 2fa1a18b9..c9d51aceb 100644 --- a/docs/src/jump_solve.md +++ b/docs/src/jump_solve.md @@ -65,6 +65,35 @@ algorithms are optimized for pure jump problems. `SSAStepper` supports the common `alias` keyword argument for `u0`, `p`, and `tstops`, and reuses the problem's jump state rather than copying it; see [Jump state ownership and problem reuse](@ref jump_state_ownership). + Callbacks can only evaluate an `SSAStepper` integrator at its current time, + unless the problem is solved with `SSAStepper(; save_uprev = true)`; see + [Saving with callbacks and evaluating the integrator](@ref ssa_integrator_evaluation). + +### [Supported state types](@id ssa_state_types) + +`SSAStepper` supports the following types for the state `u`: + +| State `u` | Jump types and aggregators | Status | +|:-------------------------------------------------------------------------- |:--------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |:----------------------------------- | +| `Vector` of integers or floats | `MassActionJump`s and `ConstantRateJump`s with every non-spatial aggregator; bounded `VariableRateJump`s with `Coevolve` | Fully supported | +| `SVector` (from StaticArrays.jl) of integers or floats | As for `Vector` | Fully supported | +| `Matrix{Int}` of species × sites | Spatial problems: `MassActionJump`s plus hopping, from a `JumpProblem` with `hopping_constants` and a `spatial_system`. `NSM` and `DirectCRDirect` keep the matrix; any other aggregator flattens the problem into a `Vector` state, which requires per-species or per-species-and-site hopping constants | Fully supported | +| Integer or float scalar | `ConstantRateJump`s with every non-spatial aggregator except `RSSA` and `RSSACR`; bounded `VariableRateJump`s with `Coevolve`. `MassActionJump`s are not supported. | Supported without mass-action jumps | +| Anything else, for example nested arrays, `MVector`s, or `ArrayPartition`s | None | Not supported | + +Notes: + + - Bounded `VariableRateJump`s require the `Coevolve` aggregator for every state type. + - Spatial problems are described in the [spatial tutorial](tutorials/spatial.md). + - An `SVector` or a scalar cannot be modified in place, so `ConstantRateJump` and + `VariableRateJump` `affect!` functions must assign a new value to `integrator.u`, + for example `integrator.u = setindex(integrator.u, integrator.u[1] + 1, 1)` or + `integrator.u += 1`. Their rate functions receive the scalar or `SVector` state. + - Some aggregators need dependency graphs when a model has `ConstantRateJump`s or + `VariableRateJump`s; see the aggregator docstrings. + - Other state types are not supported, even when a model built from user-written jumps + appears to work with them. + - Coupled ODE and SDE problems follow the state-type rules of their solvers. ## RegularJump Compatible Methods diff --git a/docs/src/migration.md b/docs/src/migration.md index addcadfa2..ec7d6e2ea 100644 --- a/docs/src/migration.md +++ b/docs/src/migration.md @@ -194,3 +194,28 @@ mutable RNG from a custom ensemble `rng_func`. For identical seeded trajectories, start from identical learned state as well as identical model inputs; reusing a problem with a different learned order can change the trajectory. + +## Evaluating an `SSAStepper` integrator at earlier times + +`integrator(t)` and `integrator(out, t)` previously returned the current state of +an `SSAStepper` integrator for any `t`. That was silently wrong for times before +the most recent event, which callbacks such as DiffEqCallbacks' `SavingCallback` +and `FunctionCallingCallback` request at their save times. Evaluating at a time +other than `integrator.t` now raises an error, and there are two remedies: + +```julia +using DiffEqCallbacks +saved = SavedValues(Float64, Int) +cb = SavingCallback((u, t, integrator) -> sum(u), saved; saveat = 0.0:1.0:10.0) + +# Keep the state at the start of each step, at the cost of a copy per step: +sol = solve(jprob, SSAStepper(; save_uprev = true); callback = cb) + +# Or stop exactly at the save times, avoiding the per-step copy: +sol = solve(jprob, SSAStepper(); callback = cb, tstops = 0.0:1.0:10.0) +``` + +`SSAStepper` is now a parametric type, but `SSAStepper()` is unchanged. See +[Saving with callbacks and evaluating the integrator](@ref ssa_integrator_evaluation) +for when to choose each remedy, and [Supported state types](@ref ssa_state_types) +for the state types `SSAStepper` supports. diff --git a/docs/src/tutorials/discrete_stochastic_example.md b/docs/src/tutorials/discrete_stochastic_example.md index 03ef1bc17..5ffe7b105 100644 --- a/docs/src/tutorials/discrete_stochastic_example.md +++ b/docs/src/tutorials/discrete_stochastic_example.md @@ -689,6 +689,101 @@ the jump whenever `save_positions = (true, true)`. All other saving behavior is delegated to the integrator and controlled via keyword arguments to `solve` such as `saveat`. +## [Saving with callbacks and evaluating the integrator](@id ssa_integrator_evaluation) + +`saveat` saves the full state exactly at the requested times. To save something +else at a set of times, for example a function of the state, a `SavingCallback` +from DiffEqCallbacks.jl can be used. Such callbacks run at the end of each step +and evaluate the integrator, `integrator(t)`, at the save times that fall within +the step. With `SSAStepper` the sampled path is piecewise constant, and jumps and +callbacks only change the state at the end of a step, so the state at a save time +before the current time is the state at the start of the step. `SSAStepper` keeps +that state only when asked to. + +We use the SIR model from above, defined with `ConstantRateJump`s, and turn off +saving at every jump since only the callback saves: + +```@example ssa_eval +using JumpProcesses, DiffEqCallbacks +p = (0.1 / 1000.0, 0.01) # (β, ν) +rate1(u, p, t) = p[1] * u[1] * u[2] # β*S*I +function affect1!(integrator) + integrator.u[1] -= 1 + integrator.u[2] += 1 + nothing +end +rate2(u, p, t) = p[2] * u[2] # ν*I +function affect2!(integrator) + integrator.u[2] -= 1 + integrator.u[3] += 1 + nothing +end +prob = DiscreteProblem([999, 10, 0], (0.0, 250.0), p) +jump_prob = JumpProblem(prob, Direct(), ConstantRateJump(rate1, affect1!), + ConstantRateJump(rate2, affect2!); save_positions = (false, false)) +nothing # hide +``` + +We save the number of people who have ever been infected, `I(t) + R(t)`, every +25 days. Solving with `SSAStepper(; save_uprev = true)` makes the integrator keep +the state at the start of each step in `integrator.uprev`, so it can be evaluated +anywhere in the current step, `[integrator.tprev, integrator.t]`: + +```@example ssa_eval +save_times = 0.0:25.0:250.0 +ever_infected(u, t, integrator) = u[2] + u[3] +saved = SavedValues(Float64, Int) +callback = SavingCallback(ever_infected, saved; saveat = save_times) +solve(jump_prob, SSAStepper(; save_uprev = true); callback, seed = 1234) +saved.saveval +``` + +The values agree exactly with those computed from `saveat` for the same seed: + +```@example ssa_eval +full = solve(jump_prob, SSAStepper(); saveat = save_times, seed = 1234) +saved.saveval == [u[2] + u[3] for u in full.u] +``` + +Keeping `uprev` copies the state on every step, a cost proportional to the size +of the state that can dominate the per-step cost of the efficient aggregators on +large or spatial systems. Alternatively, pass the save times as `tstops`, so that +the integrator stops exactly at them and never needs an earlier state: + +```@example ssa_eval +saved_at_stops = SavedValues(Float64, Int) +callback = SavingCallback(ever_infected, saved_at_stops; saveat = save_times) +solve(jump_prob, SSAStepper(); callback, tstops = save_times, seed = 1234) +saved_at_stops.saveval == saved.saveval +``` + +This avoids the per-step copy, but adds a stop and a pass through the callbacks at +each save time. When every jump is saved, each stop also adds a saved state, so +pair it with `save_positions = (false, false)`. + +Without either, `SSAStepper()` can only be evaluated at the current time, so the +`SavingCallback` raises an `ArgumentError` at the first save time that falls +before the end of a step: + +```julia +solve(jump_prob, SSAStepper(); callback) # ArgumentError +``` + +In summary: + + - To save the full state at given times, use `saveat`, which is exact and cheapest. + - To save functions of the state at given times, pass those times as `tstops` + as well, or use `SSAStepper(; save_uprev = true)`. + - For callbacks that need the state at the start of the step, `integrator.uprev`, + or that evaluate the integrator at times within a step, use + `SSAStepper(; save_uprev = true)`. It also provides `get_tmp_cache`, which + callbacks that evaluate the integrator in place use. + +`save_uprev` works with every [supported state type](@ref ssa_state_types), +including scalar states. DiffEqCallbacks' integrating callbacks +(`IntegratingCallback`, `IntegratingSumCallback`) are not yet supported with +`SSAStepper`. + ## Defining the Jumps Directly: Mixing `ConstantRateJump`/`VariableRateJump` and `MassActionJump` Suppose we now want to add in to the original SIR model another jump that diff --git a/ext/JumpProcessesKernelAbstractionsExt.jl b/ext/JumpProcessesKernelAbstractionsExt.jl index 5e212792d..948c8740c 100644 --- a/ext/JumpProcessesKernelAbstractionsExt.jl +++ b/ext/JumpProcessesKernelAbstractionsExt.jl @@ -25,18 +25,28 @@ function seed_kernel_backend!(backend, seed, rng, kwargs) return nothing end +# Keywords stored on the problem wrapped by a `JumpProblem`, e.g. `DiscreteProblem(...; +# alias_jump = true)`. +wrapped_problem_kwargs(jump_prob::JumpProblem) = wrapped_problem_kwargs(jump_prob.prob) +wrapped_problem_kwargs(prob) = hasproperty(prob, :kwargs) ? prob.kwargs : (;) + # Kernels always build device-owned state, so the removed `alias_jump` keyword and -# `alias` controls are rejected rather than silently ignored. These paths bypass -# SciMLBase's keyword validation, so the never-supported plural spelling is rejected here. -function reject_kernel_alias_kwargs(kwargs) - haskey(kwargs, :alias_jump) && +# `alias` controls are rejected rather than silently ignored, whether passed to `solve` +# or stored on the wrapped problem. These paths bypass SciMLBase's keyword validation, so +# the never-supported plural spelling is rejected here too, passed or stored. +function reject_kernel_alias_kwargs(jump_prob, kwargs) + stored = wrapped_problem_kwargs(jump_prob) + (haskey(kwargs, :alias_jump) || haskey(stored, :alias_jump)) && throw(ArgumentError(JumpProcesses.ALIAS_JUMP_REMOVED_MSG)) - haskey(kwargs, :alias_jumps) && throw(ArgumentError( + (haskey(kwargs, :alias_jumps) || haskey(stored, :alias_jumps)) && throw(ArgumentError( "`alias_jumps` is not a keyword argument; it is a field of some alias " * "specifiers. EnsembleGPUKernel always builds device-owned jump state.")) - get(kwargs, :alias, nothing) === nothing || throw(ArgumentError( - "EnsembleGPUKernel does not support the `alias` keyword with multiple " * - "trajectories, since kernels always build device-owned state.")) + (get(kwargs, :alias, nothing) === nothing && + get(stored, :alias, nothing) === nothing) || + throw(ArgumentError( + "EnsembleGPUKernel does not support the `alias` keyword, whether passed to " * + "`solve` or stored on the problem, with multiple trajectories, since kernels " * + "always build device-owned state.")) return nothing end @@ -58,7 +68,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, ensemblealg.backend === nothing ? backend = CPU() : backend = ensemblealg.backend - reject_kernel_alias_kwargs(kwargs) + reject_kernel_alias_kwargs(ensembleprob.prob, kwargs) seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob diff --git a/ext/explicit_tau_leaping.jl b/ext/explicit_tau_leaping.jl index 3cb70a797..908877d84 100644 --- a/ext/explicit_tau_leaping.jl +++ b/ext/explicit_tau_leaping.jl @@ -251,7 +251,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, end ensemblealg.backend === nothing ? backend = CPU() : backend = ensemblealg.backend - reject_kernel_alias_kwargs(kwargs) + reject_kernel_alias_kwargs(ensembleprob.prob, kwargs) seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob diff --git a/ext/ssa_stepper.jl b/ext/ssa_stepper.jl index 1d5750205..f40c49935 100644 --- a/ext/ssa_stepper.jl +++ b/ext/ssa_stepper.jl @@ -318,9 +318,13 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, callback === nothing || error("EnsembleGPUKernel with SSAStepper does not support callbacks, since they \ would have to run inside the GPU kernel.") + JumpProcesses.ssa_save_uprev(alg) && + throw(ArgumentError("EnsembleGPUKernel with SSAStepper does not support \ + `SSAStepper(; save_uprev = true)` with multiple trajectories, since kernels \ + never expose an integrator to evaluate. Use `SSAStepper()`.")) ensemblealg.backend === nothing ? backend = CPU() : backend = ensemblealg.backend - reject_kernel_alias_kwargs(kwargs) + reject_kernel_alias_kwargs(ensembleprob.prob, kwargs) seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index 35c35c542..550d896e7 100644 --- a/src/SSA_stepper.jl +++ b/src/SSA_stepper.jl @@ -4,6 +4,56 @@ $(TYPEDEF) Highly efficient integrator for pure jump problems that involve only `ConstantRateJump`s, `MassActionJump`s, and/or `VariableRateJump`s *with rate bounds*. +## Constructor + +```julia +SSAStepper(; save_uprev = false) +``` + + - `save_uprev`: whether the integrator keeps the state at the start of each step in + `integrator.uprev`, so that it can be evaluated anywhere in the current step (see + "Evaluating the integrator" below). This copies the state on every step, at a cost + proportional to the size of the state, which can dominate the per-step cost of + efficient aggregators on large or spatial systems. `SSAStepper()` keeps no copy. + +## Supported state types + +The state `u` can be a `Vector` or an `SVector` of integers or floats; a species × sites +`Matrix{Int}` for spatial problems, which `NSM` and `DirectCRDirect` keep and other +aggregators flatten into a `Vector`; or, for models without `MassActionJump`s, an integer +or float scalar (not with `RSSA` or `RSSACR`). Bounded +`VariableRateJump`s require `Coevolve`, and `affect!` functions for `SVector` and scalar +states must assign a new value to `integrator.u`. Other state types are not supported. See +the [supported state +types](https://docs.sciml.ai/JumpProcesses/stable/jump_solve/#ssa_state_types) for +details. + +## Evaluating the integrator + +Callbacks, such as DiffEqCallbacks' `SavingCallback` and `FunctionCallingCallback` with +save times, may evaluate the integrator at times before the current one with +`integrator(t)` or `integrator(out, t)`. The sampled path is piecewise constant, and +jumps and callbacks only change the state at the end of a step, so: + + - Evaluating at the current time, `integrator.t`, returns the current state. + - With `SSAStepper(; save_uprev = true)`, evaluating at any `t` in + `[integrator.tprev, integrator.t)` returns the state at the start of the step, + `integrator.uprev`. `get_tmp_cache(integrator)` then also provides scratch space for + callbacks that evaluate the integrator in place. + - Any other time raises an `ArgumentError`: a path cannot be extrapolated, and earlier + states are only available from the saved solution. With the default, + `SSAStepper()`, every time other than `integrator.t` raises this error, as does + `get_tmp_cache`. + +`integrator(t)` returns a copy for array states, and `integrator(out, t)` writes into +`out` and returns it. Instead of `save_uprev = true`, the times at which callbacks +evaluate the integrator can be passed as `tstops`, so that it stops exactly at them. This +avoids the per-step copy, but adds a stop and a pass through the callbacks at each time, +and, when every jump is saved, a saved point at each stop, so pair it with +`save_positions = (false, false)`. To save the full state, `saveat` is exact and needs +neither. DiffEqCallbacks' integrating callbacks (`IntegratingCallback`, +`IntegratingSumCallback`) are not yet supported with `SSAStepper`. + ## Notes - Only works with `JumpProblem`s defined from `DiscreteProblem`s. @@ -76,7 +126,18 @@ see the [tutorial](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) for details. """ -struct SSAStepper <: SciMLBase.AbstractDEAlgorithm end +struct SSAStepper{SaveUprev} <: SciMLBase.AbstractDEAlgorithm + function SSAStepper{S}() where {S} + S isa Bool || + throw(ArgumentError("The `SSAStepper` type parameter must be a `Bool`, got $S.")) + new{S}() + end +end +SSAStepper(; save_uprev::Bool = false) = SSAStepper{save_uprev}() + +# Whether the integrator keeps the state at the start of each step (`save_uprev`). +ssa_save_uprev(::SSAStepper{S}) where {S} = S + SciMLBase.allows_late_binding_tstops(::SSAStepper) = true SciMLBase.supports_solve_rng(::JumpProblem, ::SSAStepper) = true @@ -89,7 +150,7 @@ Integrator for pure jump problems solved via `SSAStepper`. $(FIELDS) """ -mutable struct SSAIntegrator{F, uType, tType, tdirType, P, S, CB, SA, OPT, TS, R} <: +mutable struct SSAIntegrator{F, uType, tType, tdirType, P, S, CB, SA, OPT, TS, R, UP, TC} <: AbstractSSAIntegrator{SSAStepper, Nothing, uType, tType} """ The underlying `prob.f` function. Not currently used. @@ -104,10 +165,16 @@ mutable struct SSAIntegrator{F, uType, tType, tdirType, P, S, CB, SA, OPT, TS, R """ t::tType """ - The previous time a jump occurred. + The time at the start of the current step. A step ends at the next jump or `tstop`, + and `solve!`'s final advance to the end time also counts as a step. """ tprev::tType """ + The state at the start of the current step, at time `tprev`, when solving with + `SSAStepper(; save_uprev = true)`; otherwise `nothing`. + """ + uprev::UP + """ The direction time is changing in (must be positive, indicating time is increasing) """ tdir::tdirType @@ -167,6 +234,11 @@ mutable struct SSAIntegrator{F, uType, tType, tdirType, P, S, CB, SA, OPT, TS, R The random number generator. """ rng::R + """ + Scratch space returned by `get_tmp_cache` for `save_uprev = true` and array states + other than `SVector`s; otherwise `nothing`. + """ + tmp_cache::TC end SciMLBase.has_rng(::SSAIntegrator) = true @@ -184,17 +256,93 @@ function SciMLBase.set_rng!(integrator::SSAIntegrator, rng) nothing end -(integrator::SSAIntegrator)(t) = recursivecopy(integrator.u) -# Copy the current state into `out` and return it. Nested arrays are copied recursively, -# like saved states; other states, including scalars, are broadcast into `out`. -function (integrator::SSAIntegrator)(out, t) +# Copy the state `src` into the existing array `dest`. Nested arrays are copied +# recursively, like saved states; other states, including scalars and `SVector`s, are +# broadcast into `dest`. +@inline function copy_state!(dest, src) + if src isa AbstractArray && eltype(src) <: AbstractArray + recursivecopy!(dest, src) + else + dest .= src + end + return dest +end + +# Store the state at the start of a step in `uprev` (only with `save_uprev = true`). +# Scalar and `SVector` states are replaced; other states are copied into the existing +# buffer, so no allocation occurs. +@inline store_uprev!(integrator::SSAIntegrator) = store_uprev!(integrator, integrator.uprev) +@inline store_uprev!(::SSAIntegrator, ::Nothing) = nothing +@inline function store_uprev!(integrator::SSAIntegrator, uprev) u = integrator.u - if u isa AbstractArray && eltype(u) <: AbstractArray - recursivecopy!(out, u) + if u isa Union{Number, SVector} + integrator.uprev = u + else + copy_state!(uprev, u) + end + return nothing +end + +# The `uprev` and `get_tmp_cache` storage for a new integrator: none unless +# `save_uprev = true`, and no scratch array for scalar and `SVector` states. +ssa_uprev_storage(::SSAStepper{false}, _) = (nothing, nothing) +function ssa_uprev_storage(::SSAStepper{true}, u0) + if u0 isa Union{Number, SVector} + return u0, nothing else - out .= u + return recursivecopy(u0), recursivecopy(u0) end - return out +end + +# The state at time `t`, which must be `integrator.t`, or, with `save_uprev = true`, lie +# in `[integrator.tprev, integrator.t)`, where the state equals `uprev`, since jumps and +# callbacks only change the state at the end of a step. +@inline function state_at(integrator::SSAIntegrator, t) + t == integrator.t && return integrator.u + uprev = integrator.uprev + uprev === nothing && throw_ssa_evaluation_needs_uprev(integrator, t) + (integrator.tprev <= t < integrator.t) || + throw_ssa_evaluation_out_of_step(integrator, t) + return uprev +end + +@noinline function throw_ssa_evaluation_needs_uprev(integrator, t) + throw(ArgumentError("An `SSAStepper` integrator can only be evaluated at its current \ + time `integrator.t = $(integrator.t)`, but was evaluated at `t = $t`. To evaluate \ + it anywhere in the current step, `[integrator.tprev, integrator.t]`, as \ + `SavingCallback`s and `FunctionCallingCallback`s with save times require, solve \ + with `SSAStepper(; save_uprev = true)`. Alternatively, pass those times as \ + `tstops`, so that the integrator stops exactly at them.")) +end + +@noinline function throw_ssa_evaluation_out_of_step(integrator, t) + throw(ArgumentError("An `SSAStepper` integrator can only be evaluated in the current \ + step, `[integrator.tprev, integrator.t] = [$(integrator.tprev), \ + $(integrator.t)]`, but was evaluated at `t = $t`. A jump process path cannot be \ + extrapolated, and states before the current step are only available from the \ + saved solution.")) +end + +@noinline function throw_ssa_tmp_cache_needs_uprev() + throw(ArgumentError("`get_tmp_cache` is only available for `SSAStepper` integrators \ + solved with `SSAStepper(; save_uprev = true)`. It provides scratch space for \ + callbacks, such as an in-place `SavingCallback`, that evaluate the integrator \ + at times before `integrator.t`, which requires `save_uprev = true`. \ + Alternatively, pass the callback's times as `tstops`, so that the integrator \ + stops exactly at them.")) +end + +# Evaluate the piecewise-constant path at time `t`; see `state_at` for the valid times. +# Returns a copy for array states, and the value itself for scalar and `SVector` states. +(integrator::SSAIntegrator)(t) = recursivecopy(state_at(integrator, t)) + +# Copy the state at time `t` into the caller's `out` and return `out`. +(integrator::SSAIntegrator)(out, t) = copy_state!(out, state_at(integrator, t)) + +function SciMLBase.get_tmp_cache(integrator::SSAIntegrator) + integrator.uprev === nothing && throw_ssa_tmp_cache_needs_uprev() + cache = integrator.tmp_cache + return cache === nothing ? nothing : (cache,) end # Save a snapshot of the current state, matching OrdinaryDiffEq's save semantics @@ -247,14 +395,17 @@ function DiffEqBase.solve!(integrator::SSAIntegrator) # if the user terminated the solve we shouldn't advance in time any more if integrator.sol.retcode !== ReturnCode.Terminated - integrator.t = end_time - - # check callbacks one last time - if !(integrator.opts.callback.discrete_callbacks isa Tuple{}) - DiffEqBase.apply_discrete_callback!(integrator, - integrator.opts.callback.discrete_callbacks...) + # The advance to the end time is a final step without a jump, so it starts a new + # step for evaluating the integrator (and `tprev` always marks a step's start). + if integrator.t < end_time + store_uprev!(integrator) + integrator.tprev = integrator.t end + integrator.t = end_time + # Save the pending `saveat` times before the end time first: the state there is + # the one from before the final callback pass, and saving them afterwards would + # also put them after any saves the callbacks make at the end time. if integrator.saveat !== nothing && !isempty(integrator.saveat) # Split to help prediction while integrator.cur_saveat <= length(integrator.saveat) && @@ -264,6 +415,12 @@ function DiffEqBase.solve!(integrator::SSAIntegrator) end end + # check callbacks one last time + if !(integrator.opts.callback.discrete_callbacks isa Tuple{}) + DiffEqBase.apply_discrete_callback!(integrator, + integrator.opts.callback.discrete_callbacks...) + end + if integrator.save_end && integrator.sol.t[end] != end_time save_current_state!(integrator, end_time) end @@ -432,10 +589,14 @@ function SciMLBase.__init(jump_prob::JumpProblem, _tstops, _alias_tstops, copied_tstops = tType[], true, true end - integrator = SSAIntegrator(prob.f, u0, prob.tspan[1], prob.tspan[1], tdir, + # With `save_uprev = true`, `uprev` and the scratch cache are distinct from `u0`, even + # when `u0` aliases the problem's initial condition. + uprev, tmp_cache = ssa_uprev_storage(alg, u0) + + integrator = SSAIntegrator(prob.f, u0, prob.tspan[1], prob.tspan[1], uprev, tdir, p, sol, 1, prob.tspan[1], cb, _saveat, save_everystep, save_end, cur_saveat, opts, _tstops, 1, false, true, _alias_tstops, - copied_tstops, _rng) + copied_tstops, _rng, tmp_cache) cb.initialize(cb, integrator.u, prob.tspan[1], integrator) DiffEqBase.initialize!(opts.callback, integrator.u, prob.tspan[1], integrator) if save_start @@ -494,6 +655,7 @@ end end function DiffEqBase.step!(integrator::SSAIntegrator) + store_uprev!(integrator) integrator.tprev = integrator.t next_jump_time = integrator.tstop > integrator.t ? integrator.tstop : typemax(integrator.tstop) @@ -588,6 +750,6 @@ end function SciMLBase.isdenseplot(sol::ODESolution{ T, N, uType, uType2, DType, tType, rateType, discType, P, - SSAStepper}) where {T, N, uType, uType2, DType, tType, rateType, discType, P} + <:SSAStepper}) where {T, N, uType, uType2, DType, tType, rateType, discType, P} sol.dense end diff --git a/src/aggregators/aggregators.jl b/src/aggregators/aggregators.jl index b7f2f32b3..6bd27d992 100644 --- a/src/aggregators/aggregators.jl +++ b/src/aggregators/aggregators.jl @@ -66,6 +66,8 @@ struct SortingDirect <: AbstractAggregatorAlgorithm end """ The Rejection SSA method. One of the best methods for systems with hundreds to many thousands of jumps (along with `RSSACR`) and sparse dependency graphs. +Requires an array state, such as a `Vector` or `SVector`; scalar states are not +supported. V. H. Thanh, C. Priami and R. Zunino, Efficient rejection-based simulation of biochemical reactions with stochastic noise and delays, Journal of Chemical @@ -79,7 +81,8 @@ struct RSSA <: AbstractAggregatorAlgorithm end """ The Rejection SSA Composition-Rejection method. Often the best performer for -systems with tens of thousands of jumps and sparse dependency graphs. +systems with tens of thousands of jumps and sparse dependency graphs. Requires an +array state, such as a `Vector` or `SVector`; scalar states are not supported. V. H. Thanh, R. Zunino, and C. Priami, Efficient constant-time complexity algorithm for stochastic simulation of large reaction networks, IEEE/ACM @@ -161,7 +164,8 @@ struct CCNRM <: AbstractAggregatorAlgorithm end The Next Subvolume Method for spatial jump process simulations. Usually slower than `DirectCRDirect`. Uses an indexed priority queue tree structure to determine where on the grid/graph the next jump occurs, and then the `Direct` -method to determine which jump at the given location occurs. +method to determine which jump at the given location occurs. The state is a +species × sites matrix, with `u[i, j]` the number of species `i` at site `j`. Elf, Johan and Ehrenberg, M, Spontaneous separation of bi-stable biochemical systems into spatial domains of opposite phases,Systems Biology, 1(2), 230-236 @@ -172,7 +176,8 @@ struct NSM <: AbstractAggregatorAlgorithm end """ The Direct Composition-Rejection Direct method. Uses the `DirectCR` method to determine where on the grid/graph a jump occurs, and the `Direct` method to -determine which jump occurs at the sampled location. +determine which jump occurs at the sampled location. The state is a species × sites +matrix, with `u[i, j]` the number of species `i` at site `j`. Kevin R. Sanft and Hans G. Othmer, Constant-complexity stochastic simulation algorithm with optimal binning, Journal of Chemical Physics 143, 074108 diff --git a/src/jumps.jl b/src/jumps.jl index ef1688dd9..060745bbf 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -566,6 +566,12 @@ jprob = JumpProblem(prob, Direct(), maj) - Also see the [main docs](https://docs.sciml.ai/JumpProcesses/stable/jump_types/#Defining-a-Mass-Action-Jump) for how to specify reactions with no products or no reactants. + - For non-spatial problems the state `u` must be a `Vector` or an `SVector` of integers + or floats; scalar and other states are not supported. Spatial problems, built by + passing `hopping_constants` and a `spatial_system` to `JumpProblem`, use a species × + sites `Matrix{Int}`. `NSM` and `DirectCRDirect` keep the matrix, while other + aggregators flatten it into a `Vector`. See the [supported state + types](https://docs.sciml.ai/JumpProcesses/stable/jump_solve/#ssa_state_types). """ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump """ diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index 4053a91b9..5faf26b2a 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -1160,8 +1160,11 @@ and accepts its usual `rng`, `seed`, and ensemble `rng_func` keywords. Kernels always build device-owned state. For multiple trajectories, the removed `alias_jump` keyword raises its migration error, and the `alias_jumps` keyword and any -`alias` value other than `nothing` raise an `ArgumentError`. With `trajectories = 1`, the serial solver's rules -apply: `SSAStepper` supports `alias`, while the tau-leaping solvers reject both keywords. +`alias` value other than `nothing` raise an `ArgumentError`, whether passed to `solve` or +stored on the problem wrapped by the `JumpProblem`. Kernels never expose an integrator, so +`SSAStepper(; save_uprev = true)` also raises an `ArgumentError`. With +`trajectories = 1`, the serial solver's rules apply: `SSAStepper` supports `alias` and +`save_uprev`, while the tau-leaping solvers reject both alias keywords. ## Returns diff --git a/test/kernel_contract_tests.jl b/test/kernel_contract_tests.jl index aeb135010..ed0fe35cf 100644 --- a/test/kernel_contract_tests.jl +++ b/test/kernel_contract_tests.jl @@ -75,12 +75,39 @@ function run_kernel_alias_tests(backend, alg, jump_prob; kwargs...) alg, EnsembleGPUKernel(backend); trajectories = 16, kwargs..., alias_jump = true) + # The same keywords stored on the wrapped problem are rejected too. + wrapped = jump_prob.prob + for stored in ((; alias_jump = true), (; alias_jumps = false), (; alias = true), + (; alias = SciMLBase.ODEAliasSpecifier(alias_u0 = true))) + stored_prob = DiscreteProblem(wrapped.f, wrapped.u0, wrapped.tspan, wrapped.p; + stored...) + stored_ensemble = EnsembleProblem(remake(jump_prob; prob = stored_prob)) + @test_throws ArgumentError solve(stored_ensemble, alg, + EnsembleGPUKernel(backend); trajectories = 16, kwargs...) + end + stored_prob = DiscreteProblem(wrapped.f, wrapped.u0, wrapped.tspan, wrapped.p; + alias_jump = true) + @test_throws "`alias_jump` keyword argument was removed" solve( + EnsembleProblem(remake(jump_prob; prob = stored_prob)), alg, + EnsembleGPUKernel(backend); trajectories = 16, kwargs...) + if backend isa CPU sol = solve(ensemble_prob, alg, EnsembleGPUKernel(backend); trajectories = 16, kwargs..., alias = nothing) @test all(SciMLBase.successful_retcode, sol.u) end + # Kernels never expose an integrator, so keeping the pre-step state is rejected + # with multiple trajectories; a single trajectory uses the serial solver. + if alg isa SSAStepper + @test_throws "save_uprev" solve(ensemble_prob, SSAStepper(; save_uprev = true), + EnsembleGPUKernel(backend); trajectories = 16, kwargs...) + single_sol = solve(ensemble_prob, SSAStepper(; save_uprev = true), + EnsembleGPUKernel(backend); trajectories = 1, seed = 123, kwargs..., + alias = false) + @test all(SciMLBase.successful_retcode, single_sol.u) + end + # A single trajectory delegates to the serial solver, whose rules apply: # SSAStepper raises the migration error and supports `alias`, while the # tau-leaping solvers' strict keyword signatures reject both keywords. diff --git a/test/runtests.jl b/test/runtests.jl index 27cbdb4e5..dea382bb0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -43,6 +43,7 @@ end end if GROUP == "All" || GROUP == "InterfaceII" @time @safetestset "Saveat Regression test" begin include("saveat_regression.jl") end + @time @safetestset "SSAIntegrator evaluation and save_uprev" begin include("ssa_integrator_evaluation.jl") end @time @safetestset "Save_positions test" begin include("save_positions.jl") end @time @safetestset "RNG kwarg tests" begin include("rng_kwarg_tests.jl") end @time @safetestset "Stochastic dispatch tests" begin include("stochastic_dispatch.jl") end diff --git a/test/saveat_regression.jl b/test/saveat_regression.jl index 486a09a95..483d08111 100644 --- a/test/saveat_regression.jl +++ b/test/saveat_regression.jl @@ -158,3 +158,32 @@ end @test draws == (3, 3) end end + +# A callback that changes the state at the end time must not affect the pending `saveat` +# times before it, and the saved times must stay sorted: `solve!` saves those times +# before its final callback pass. +@testset "End-time callback after pending saveat times" begin + no_jumps = ConstantRateJump((u, p, t) -> 0.0, + integrator -> (integrator.u[1] += 1; nothing)) + ts = [0.0, 0.5, 0.75, 1.0] + @testset "jump save_positions=$jump_saves, callback save_positions=$cb_saves" for jump_saves in ( + (false, false), (true, true)), + cb_saves in ((true, true), (false, false)) + + set_to_ten = DiscreteCallback((u, t, integrator) -> t == 1.0, + integrator -> (integrator.u[1] = 10; nothing); save_positions = cb_saves) + jprob = JumpProblem(DiscreteProblem([0], (0.0, 1.0)), Direct(), no_jumps; + save_positions = jump_saves) + sol = solve(jprob, SSAStepper(); saveat = ts, callback = set_to_ten, seed = 1) + @test issorted(sol.t) + # Every requested time is saved, with the state before the callback at the + # earlier times and after it at the end time. + @test [first(sol.u[findfirst(==(t), sol.t)]) for t in ts[1:3]] == [0, 0, 0] + @test sol.t[end] == 1.0 + @test sol(0.75) == [0] + # A callback that saves after its effect records the new state at the end time. + # One that does not save may leave the earlier end-time save as the last entry, + # as OrdinaryDiffEq does for callbacks with `save_positions = (false, false)`. + cb_saves[2] && @test sol.u[end] == [10] + end +end diff --git a/test/ssa_integrator_evaluation.jl b/test/ssa_integrator_evaluation.jl new file mode 100644 index 000000000..f9e2dc0ee --- /dev/null +++ b/test/ssa_integrator_evaluation.jl @@ -0,0 +1,424 @@ +using JumpProcesses, DiffEqCallbacks, SciMLBase, StaticArrays, Test + +# Evaluating an `SSAIntegrator` at a time: only `integrator.t` by default, and anywhere in +# the current step, `[integrator.tprev, integrator.t]`, with `SSAStepper(; save_uprev = +# true)`, which keeps the state at the start of each step in `integrator.uprev`. + +errmsg(f) = + try + f() + "" + catch err + sprint(showerror, err) + end + +mutable_state(u) = !(u isa Union{Number, SVector}) + +# Birth–death models built from constant-rate jumps with user-written effects, for vector +# and scalar states. +const VECTOR_BIRTH = ConstantRateJump((u, p, t) -> 2.0, + integrator -> (integrator.u[1] += 1; nothing)) +const VECTOR_DEATH = ConstantRateJump((u, p, t) -> 0.1 * u[1], + integrator -> (integrator.u[1] -= 1; nothing)) +const SCALAR_BIRTH = ConstantRateJump((u, p, t) -> 2.0, + integrator -> (integrator.u += 1; nothing)) +const SCALAR_DEATH = ConstantRateJump((u, p, t) -> 0.1 * u, + integrator -> (integrator.u -= 1; nothing)) + +function vector_jprob(; u0 = [0], tspan = (0.0, 20.0), kwargs...) + JumpProblem(DiscreteProblem(u0, tspan), Direct(), VECTOR_BIRTH, VECTOR_DEATH; kwargs...) +end + +@testset "SSAStepper type and traits" begin + @test SSAStepper() === SSAStepper{false}() + @test SSAStepper(; save_uprev = true) === SSAStepper{true}() + @test_throws ArgumentError SSAStepper{1}() + jprob = vector_jprob() + for alg in (SSAStepper(), SSAStepper(; save_uprev = true)) + @test alg isa SSAStepper + @test SciMLBase.allows_late_binding_tstops(alg) + @test SciMLBase.supports_solve_rng(jprob, alg) + end + # Solving without an algorithm still selects the default `SSAStepper()`. + @test solve(jprob; seed = 1).alg === SSAStepper() +end + +@testset "Default SSAStepper only evaluates at the current time" begin + integrator = init(vector_jprob(), SSAStepper(); seed = 1) + @test integrator.uprev === nothing + @test integrator(integrator.t) == integrator.u + step!(integrator) + step!(integrator) + current = integrator(integrator.t) + @test current == integrator.u + @test current !== integrator.u + out = zeros(1) + @test integrator(out, integrator.t) === out + @test out == integrator.u + + @test_throws ArgumentError integrator(integrator.tprev) + @test_throws ArgumentError integrator(out, integrator.tprev) + msg = errmsg(() -> integrator(integrator.tprev)) + @test occursin("save_uprev = true", msg) + @test occursin("tstops", msg) + @test_throws ArgumentError SciMLBase.get_tmp_cache(integrator) + @test occursin("save_uprev = true", errmsg(() -> SciMLBase.get_tmp_cache(integrator))) +end + +@testset "SavingCallback is exact with save_uprev or tstops" begin + ts = collect(0.0:0.5:20.0) + # A pure-death model goes extinct early, so many save times fall after its last + # event and are evaluated during `solve!`'s final advance to the end time. + death_only = ConstantRateJump((u, p, t) -> 1.0 * u[1], + integrator -> (integrator.u[1] -= 1; nothing)) + models = (("vector birth–death", [0], (VECTOR_BIRTH, VECTOR_DEATH)), + ("scalar birth–death", 0, (SCALAR_BIRTH, SCALAR_DEATH)), + ("vector pure death", [5], (death_only,))) + for (name, u0, jumps) in models, save_positions in ((true, true), (false, false)) + + @testset "$name, save_positions = $save_positions" begin + dprob = DiscreteProblem(u0, (0.0, 20.0)) + jprob = JumpProblem(dprob, Direct(), jumps...; save_positions) + # Native `saveat` is exact and runs before jumps; the seeded path is the same + # for every saving choice. + reference_prob = JumpProblem(dprob, Direct(), jumps...; + save_positions = (false, false)) + reference = [first(u) + for u in solve(reference_prob, SSAStepper(); saveat = ts, + seed = 7).u] + saved_values() = SavedValues(Float64, typeof(first(u0))) + saving_callback(values) = SavingCallback((u, t, integrator) -> first(u), + values; saveat = ts) + + values = saved_values() + solve(jprob, SSAStepper(; save_uprev = true); + callback = saving_callback(values), seed = 7) + @test values.t == ts + @test values.saveval == reference + + # Stopping at the save times avoids evaluating earlier times altogether. + values = saved_values() + solve(jprob, SSAStepper(); callback = saving_callback(values), tstops = ts, + seed = 7) + @test values.t == ts + @test values.saveval == reference + + @test_throws ArgumentError solve(jprob, SSAStepper(); + callback = saving_callback(saved_values()), seed = 7) + end + end +end + +@testset "uprev is the state after the previous step's callbacks" begin + # A callback that resets `u` without saving leaves the reset state out of the saved + # solution, but `uprev` still records it. + birth = ConstantRateJump((u, p, t) -> 1.0, + integrator -> (integrator.u[1] += 1; nothing)) + reset = DiscreteCallback((u, t, integrator) -> u[1] == 3, + integrator -> (integrator.u[1] = 0; nothing); save_positions = (false, false)) + jprob = JumpProblem(DiscreteProblem([0], (0.0, 100.0)), Direct(), birth) + integrator = init(jprob, SSAStepper(; save_uprev = true); callback = reset, seed = 3) + while integrator.u[1] != 0 || integrator.t == 0 + step!(integrator) + end + reset_time = integrator.t + @test integrator.sol.u[end] == [3] + step!(integrator) + @test integrator.tprev == reset_time + @test integrator(integrator.tprev) == [0] + @test integrator((integrator.tprev + integrator.t) / 2) == [0] +end + +@testset "Changing u between manual steps" begin + integrator = init(vector_jprob(), SSAStepper(; save_uprev = true); seed = 1) + step!(integrator) + step!(integrator) + previous = copy(integrator.uprev) + integrator.u[1] = 100 + reset_aggregated_jumps!(integrator) + @test integrator(integrator.t) == [100] + # The state on `[tprev, t)` is unchanged by the modification at `t`. + @test integrator(integrator.tprev) == previous + step!(integrator) + @test integrator(integrator.tprev) == [100] +end + +@testset "End of solve" begin + for save_uprev in (false, true) + integrator = init(vector_jprob(), SSAStepper(; save_uprev); seed = 5) + solve!(integrator) + @test integrator.t == 20.0 + # The final advance to the end time starts a step at the last event. + @test integrator.tprev == integrator.sol.t[end - 1] + @test integrator.tprev < 20.0 + if save_uprev + for t in (integrator.tprev, (integrator.tprev + 20.0) / 2, 20.0) + @test integrator(t) == integrator.u + end + else + @test_throws ArgumentError integrator(integrator.tprev) + end + end + + # A final callback that changes `u` at the end time: earlier times in the last step + # keep the state after the last event. + jprob = vector_jprob(; save_positions = (false, false)) + after_last_event = solve(jprob, SSAStepper(); seed = 5).u[end] + final_change = DiscreteCallback((u, t, integrator) -> t == 20.0, + integrator -> (integrator.u[1] = 1000; nothing)) + integrator = init(jprob, SSAStepper(; save_uprev = true); callback = final_change, + seed = 5) + solve!(integrator) + @test integrator(20.0) == [1000] + @test integrator.uprev == after_last_event + @test integrator(integrator.tprev) == after_last_event + @test integrator((integrator.tprev + 20.0) / 2) == after_last_event +end + +@testset "Terminated solves do not advance" begin + birth = ConstantRateJump((u, p, t) -> 2.0, + integrator -> (integrator.u[1] += 1; nothing)) + stop = DiscreteCallback((u, t, integrator) -> u[1] >= 3, terminate!) + jprob = JumpProblem(DiscreteProblem([0], (0.0, 20.0)), Direct(), birth) + integrator = init(jprob, SSAStepper(; save_uprev = true); callback = stop, seed = 2) + solve!(integrator) + @test integrator.sol.retcode == ReturnCode.Terminated + @test integrator.t < 20.0 + @test integrator(integrator.t) == integrator.u + @test integrator(integrator.tprev) == integrator.uprev + @test_throws ArgumentError integrator((integrator.t + 20.0) / 2) +end + +@testset "isdenseplot for both SSAStepper settings" begin + for alg in (SSAStepper(), SSAStepper(; save_uprev = true)) + # SciMLBase's generic method returns `true` for every `DiscreteProblem`. + sparse = solve(vector_jprob(; save_positions = (false, false)), alg; saveat = 1.0, + seed = 1) + @test !SciMLBase.isdenseplot(sparse) + @test SciMLBase.isdenseplot(solve(vector_jprob(), alg; seed = 1)) + end +end + +# Supported state types: vectors and `SVector`s of integers or floats (here with +# mass-action jumps) and species × sites matrices for the spatial solvers. +function birth_death_maj_jprob(u0; tspan = (0.0, 20.0), kwargs...) + maj = MassActionJump([[0 => 1], [1 => 1]], [[1 => 1], [1 => -1]]; + param_idxs = [1, 2]) + JumpProblem(DiscreteProblem(u0, tspan, [2.0, 0.1]), Direct(), maj; kwargs...) +end + +function spatial_jprob(aggregator; tspan = (0.0, 20.0), kwargs...) + u0 = zeros(Int, 1, 3) + u0[1, 2] = 10 + JumpProblem(DiscreteProblem(u0, tspan), aggregator, JumpSet(nothing); + hopping_constants = ones(1, 3), spatial_system = CartesianGridRej((3,)), kwargs...) +end + +const STATE_CASES = (("Vector{Int}", kw -> birth_death_maj_jprob([10]; kw...)), + ("Vector{Float64}", kw -> birth_death_maj_jprob([10.0]; kw...)), + ("SVector{Int}", kw -> birth_death_maj_jprob(SA[10]; kw...)), + ("SVector{Float64}", kw -> birth_death_maj_jprob(SA[10.0]; kw...)), + ("Matrix{Int} with NSM", kw -> spatial_jprob(NSM(); kw...)), + ("Matrix{Int} with DirectCRDirect", kw -> spatial_jprob(DirectCRDirect(); kw...))) + +function check_current_evaluation(integrator) + u = integrator.u + value = integrator(integrator.t) + @test value == u + @test value isa typeof(u) + if mutable_state(u) + @test value !== u + state = copy(u) + value .= -1 + @test integrator.u == state + end + out = zeros(Float64, size(u)) + @test integrator(out, integrator.t) === out + @test out == u +end + +function check_step_evaluation(integrator, save_uprev) + step!(integrator) + before = copy(integrator.u) + step!(integrator) + @test integrator.u != before + check_current_evaluation(integrator) + middle = (integrator.tprev + integrator.t) / 2 + if save_uprev + for t in (integrator.tprev, middle) + value = integrator(t) + @test value == before + @test value isa typeof(integrator.u) + if mutable_state(integrator.u) + @test value !== integrator.uprev + @test value !== integrator.u + @test value !== only(SciMLBase.get_tmp_cache(integrator)) + value .= -1 + @test integrator(t) == before + end + out = zeros(Float64, size(integrator.u)) + @test integrator(out, t) === out + @test out == before + end + cache = SciMLBase.get_tmp_cache(integrator) + if mutable_state(integrator.u) + @test cache isa Tuple{typeof(integrator.u)} + @test only(cache) === only(SciMLBase.get_tmp_cache(integrator)) + @test only(cache) !== integrator.uprev + @test only(cache) !== integrator.u + @test integrator.uprev !== integrator.u + else + @test cache === nothing + end + else + @test integrator.uprev === nothing + @test_throws ArgumentError integrator(middle) + @test_throws ArgumentError integrator(zeros(Float64, size(integrator.u)), middle) + @test_throws ArgumentError SciMLBase.get_tmp_cache(integrator) + end +end + +function check_alias(make_jprob, save_uprev) + for alias_u0 in (nothing, true, false) + jprob = make_jprob((;)) + u0 = copy(jprob.prob.u0) + alias = alias_u0 === nothing ? nothing : + SciMLBase.DiscreteAliasSpecifier(; alias_u0) + integrator = init(jprob, SSAStepper(; save_uprev); seed = 1, alias) + if mutable_state(u0) + @test (integrator.u === jprob.prob.u0) == (alias_u0 === true) + if save_uprev + @test integrator.uprev !== integrator.u + @test integrator.uprev !== jprob.prob.u0 + end + else + @test integrator.u == jprob.prob.u0 + end + step!(integrator) + step!(integrator) + alias_u0 === true || @test jprob.prob.u0 == u0 + end +end + +function check_saving(make_jprob, save_uprev) + jprob = make_jprob((; save_positions = (false, false))) + u0 = copy(jprob.prob.u0) + integrator = init(jprob, SSAStepper(; save_uprev); seed = 1, saveat = 1.0) + solve!(integrator) + sol = integrator.sol + @test length(sol.t) == 21 + @test eltype(sol.u) == typeof(integrator.u) + @test sol.u[1] == u0 + if mutable_state(integrator.u) + @test allunique(objectid.(sol.u)) + @test all(saved -> saved !== integrator.u, sol.u) + end +end + +@testset "Supported state type: $name" for (name, make_jprob) in STATE_CASES + for save_uprev in (false, true) + integrator = init(make_jprob((;)), SSAStepper(; save_uprev); seed = 1) + check_current_evaluation(integrator) + check_step_evaluation(integrator, save_uprev) + check_alias(make_jprob, save_uprev) + check_saving(make_jprob, save_uprev) + end +end + +@testset "Scalar states in models without mass-action jumps" begin + graphs = (; dep_graph = [[1, 2], [1, 2]], vartojumps_map = [[2]], + jumptovars_map = [[1], [1]]) + aggregators = filter(agg -> !(agg isa Union{RSSA, RSSACR}), + JumpProcesses.JUMP_AGGREGATORS) + bounded_birth = VariableRateJump((u, p, t) -> 1.0 + 0.1 * sin(t), + integrator -> (integrator.u += 1; nothing); urate = (u, p, t) -> 1.1, + rateinterval = (u, p, t) -> Inf) + for u0 in (0, 0.0) + for aggregator in aggregators + jprob = JumpProblem(DiscreteProblem(u0, (0.0, 20.0)), aggregator, SCALAR_BIRTH, + SCALAR_DEATH; save_positions = (false, false), graphs...) + sol = solve(jprob, SSAStepper(); saveat = 1.0, seed = 1) + @test length(sol.u) == 21 + @test eltype(sol.u) == typeof(u0) + end + jprob = JumpProblem(DiscreteProblem(u0, (0.0, 20.0)), Coevolve(), bounded_birth; + dep_graph = [[1]]) + sol = solve(jprob, SSAStepper(); seed = 1) + @test eltype(sol.u) == typeof(u0) + @test sol.u[end] > u0 + + jprob = JumpProblem(DiscreteProblem(u0, (0.0, 20.0)), Direct(), SCALAR_BIRTH, + SCALAR_DEATH) + for save_uprev in (false, true) + integrator = init(jprob, SSAStepper(; save_uprev); seed = 1) + step!(integrator) + before = integrator.u + step!(integrator) + @test integrator(integrator.t) == integrator.u + @test integrator(integrator.t) isa typeof(u0) + out = zeros(1) + @test integrator(out, integrator.t) === out + @test out == [integrator.u] + middle = (integrator.tprev + integrator.t) / 2 + if save_uprev + @test integrator(middle) == before + @test integrator(middle) isa typeof(u0) + @test integrator(out, middle) === out + @test out == [before] + @test SciMLBase.get_tmp_cache(integrator) === nothing + else + @test_throws ArgumentError integrator(middle) + end + for alias in (nothing, true, false) + @test SciMLBase.successful_retcode(solve(jprob, + SSAStepper(; save_uprev); seed = 1, alias)) + end + end + end +end + +function allocations_per_step(integrator) + step!(integrator) + step!(integrator) + return @allocated step!(integrator) +end + +@testset "Stepping with save_uprev does not allocate" begin + # Measured with solution saving off, since saved snapshots allocate by design. + for (name, make_jprob) in STATE_CASES + name in ("Vector{Int}", "SVector{Int}", "Matrix{Int} with NSM") || continue + jprob = make_jprob((; tspan = (0.0, 1.0e6), save_positions = (false, false))) + integrator = init(jprob, SSAStepper(; save_uprev = true); seed = 1, + save_start = false, save_end = false) + @test allocations_per_step(integrator) == 0 + end +end + +@testset "Evaluation and stepping infer" begin + for save_uprev in (false, true) + integrator = init(birth_death_maj_jprob([10]), SSAStepper(; save_uprev); seed = 1) + step!(integrator) + @inferred step!(integrator) + @inferred integrator(integrator.t) + @inferred integrator(zeros(1), integrator.t) + if save_uprev + @inferred integrator(integrator.tprev) + @inferred SciMLBase.get_tmp_cache(integrator) + end + end +end + +@testset "Integrators and ensembles own their uprev buffers" begin + jprob = birth_death_maj_jprob([10]) + first_integrator = init(jprob, SSAStepper(; save_uprev = true); seed = 1) + second_integrator = init(deepcopy(jprob), SSAStepper(; save_uprev = true); seed = 2) + @test first_integrator.uprev !== second_integrator.uprev + @test only(SciMLBase.get_tmp_cache(first_integrator)) !== + only(SciMLBase.get_tmp_cache(second_integrator)) + + paths(sol) = [(trajectory.t, trajectory.u) for trajectory in sol.u] + run(seed) = solve(EnsembleProblem(jprob), SSAStepper(; save_uprev = true), + EnsembleThreads(); trajectories = 16, seed) + @test all(SciMLBase.successful_retcode, run(4).u) + @test paths(run(4)) == paths(run(4)) +end From 951c81905c0f65a8991f50875817e0ee563f666f Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 10 Oct 2026 09:35:52 -0400 Subject: [PATCH 39/52] Forward stochastic init to StochasticDiffEqCore and fix init double merge Released StochasticDiffEqCore (through 2.2.3) defines an __init method for JumpProblems that is ambiguous with JumpProcesses' OrdinaryDiffEqCore extension, so init(jprob, sde_alg) could not reach Core's initializer. A new extension, JumpProcessesStochasticDiffEqCoreExt, adds one method, __init(::JumpProblem, ::Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}), that forwards to Core's exported _sde_init, so init followed by solve! matches solve. Core versions with per-algorithm JumpProblem methods (OrdinaryDiffEq.jl#4614) are more specific and take precedence, and TauLeaping keeps Core's jump-algorithm method. StochasticDiffEqCore becomes a weak dependency with compat 2.2.1, the first release whose __solve(::JumpProblem) calls _sde_init directly. The temporary [sources] pin is removed, and CI returns to the registered dependency matrix: the grouped tests (lts, 1, pre, and 32-bit), the downgrade job, ThreadSafety on lts, 1, and pre, and benchmarks on 1 and lts. Also: - fix init on ODE, DAE, and FunctionMap integrators merging the JumpProblem's stored keywords again after DiffEqBase's init_call, which ran stored callbacks twice per step and reintroduced them under merge_callbacks = false; - import SciMLBase explicitly in two tests that relied on newer OrdinaryDiffEq and StochasticDiffEq releases to provide it, which the downgrade job exposed; - decide stochastic method ownership in the dispatch tests from Core's method table, check that the forwarding method is never overwritten, and add a filtered ambiguity check covering __init, __solve, and supports_solve_rng; - document the dispatch pathways and the keyword-merge contract, and add HISTORY entries. Validation on registered dependencies, Julia 1.12.7 and 1.10.12: InterfaceI (1,789 assertions), InterfaceII (3,344), four-thread ThreadSafety (333), the stochastic suite on Core 2.2.1, 2.2.3, and OrdinaryDiffEq.jl#4614's revision (752 each), the downgrade All group on the action-derived minimum graph (5,468), strict QA (20, plus 5 registry checks), formatting, spelling and whitespace checks, and a Documenter build of the changed pages pass. Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/Downgrade.yml | 3 - .github/workflows/Tests.yml | 16 +--- .github/workflows/ThreadSafety.yml | 3 +- .github/workflows/benchmark.yml | 3 +- HISTORY.md | 10 +++ Project.toml | 8 +- docs/src/api.md | 28 +++++-- ext/JumpProcessesOrdinaryDiffEqCoreExt.jl | 9 ++- ext/JumpProcessesStochasticDiffEqCoreExt.jl | 25 ++++++ src/solve.jl | 8 +- test/callbacks.jl | 85 ++++++++++++++++++++- test/extended_jump_array_remake.jl | 2 +- test/rng_kwarg_tests.jl | 2 +- test/stochastic_dispatch.jl | 49 +++++++++++- 14 files changed, 207 insertions(+), 44 deletions(-) create mode 100644 ext/JumpProcessesStochasticDiffEqCoreExt.jl diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml index 8a630a96a..9ae531357 100644 --- a/.github/workflows/Downgrade.yml +++ b/.github/workflows/Downgrade.yml @@ -13,9 +13,6 @@ on: jobs: downgrade: name: "Downgrade" - # The downgrade action discards URL sources. Restore this job after the - # fixed StochasticDiffEqCore is released and the temporary pin is removed. - if: ${{ false }} uses: "SciML/.github/.github/workflows/downgrade.yml@v1" with: group: All diff --git a/.github/workflows/Tests.yml b/.github/workflows/Tests.yml index 6ce571b22..45dc20491 100644 --- a/.github/workflows/Tests.yml +++ b/.github/workflows/Tests.yml @@ -18,19 +18,5 @@ concurrency: jobs: tests: - # Restore grouped-tests.yml and test/test_groups.toml's full matrix after - # OrdinaryDiffEq.jl#4614 is released and the temporary Core source is removed. - strategy: - fail-fast: false - matrix: - group: ["QA", "InterfaceI", "InterfaceII"] - name: "${{ matrix.group }} (latest Julia)" - uses: "SciML/.github/.github/workflows/tests.yml@v1" - with: - project: "." - group: "${{ matrix.group }}" - julia-version: "1" - num-threads: "1" - coverage-directories: "src,ext" - timeout-minutes: 120 + uses: "SciML/.github/.github/workflows/grouped-tests.yml@v1" secrets: "inherit" diff --git a/.github/workflows/ThreadSafety.yml b/.github/workflows/ThreadSafety.yml index e79efb8f9..45939c730 100644 --- a/.github/workflows/ThreadSafety.yml +++ b/.github/workflows/ThreadSafety.yml @@ -22,9 +22,10 @@ jobs: strategy: fail-fast: false matrix: - # Restore lts/pre after the temporary Core source pin is removed. version: + - "lts" - "1" + - "pre" runs-on: ubuntu-latest steps: - uses: actions/checkout@v7 diff --git a/.github/workflows/benchmark.yml b/.github/workflows/benchmark.yml index 6e349481b..4f5b471ec 100644 --- a/.github/workflows/benchmark.yml +++ b/.github/workflows/benchmark.yml @@ -16,8 +16,7 @@ jobs: strategy: fail-fast: false matrix: - # Restore lts when the full version matrix is enabled again. - version: ["1"] + version: ["1", "lts"] env: # Force consistent Julia depot path for self-hosted runners JULIA_DEPOT_PATH: ~/.julia diff --git a/HISTORY.md b/HISTORY.md index a53264b07..81543940f 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -87,6 +87,16 @@ changes are not intended for a 9.x release. final callback pass. A callback that changed the state at the end time altered the values saved at those earlier times, and the saved times came out of order. They are now saved before the final callback pass. + - Fixed `init` on ODE, DAE, and `FunctionMap` integrators for `JumpProblem`s: + callbacks stored in the `JumpProblem` (`JumpProblem(...; callback)`) were merged + twice, so they ran twice per step, and were reintroduced even with + `merge_callbacks = false`. `init` now combines stored, wrapped-problem, and + call-level callbacks exactly as `solve` does. + - StochasticDiffEqCore is now a weak dependency: when it is installed, version + 2.2.1 or later is required. A new extension sends stochastic `init` of a + `JumpProblem` to StochasticDiffEqCore's own initializer on every supported version, + so `init` followed by `solve!` matches `solve`. StochasticDiffEqCore versions with + their own per-algorithm `JumpProblem` initializers take precedence automatically. - The state types `SSAStepper` supports are now documented: vectors and `SVector`s of integers or floats, species × sites matrices for spatial problems, and scalars for models without mass-action jumps. See "Supported diff --git a/Project.toml b/Project.toml index 627d1de7f..9da983bb8 100644 --- a/Project.toml +++ b/Project.toml @@ -27,11 +27,13 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8" +StochasticDiffEqCore = "19c5a474-6cd1-4a5f-be79-46dc34e54d7f" [extensions] JumpProcessesForwardDiffExt = "ForwardDiff" JumpProcessesKernelAbstractionsExt = ["Adapt", "KernelAbstractions"] JumpProcessesOrdinaryDiffEqCoreExt = "OrdinaryDiffEqCore" +JumpProcessesStochasticDiffEqCoreExt = ["OrdinaryDiffEqCore", "StochasticDiffEqCore"] [compat] ADTypes = "1.22.0" @@ -64,15 +66,11 @@ StableRNGs = "1" StaticArrays = "1.9.18" Statistics = "1" StochasticDiffEq = "7.0.0, 7" -StochasticDiffEqCore = "2.2.4" +StochasticDiffEqCore = "2.2.1" SymbolicIndexingInterface = "0.3.48" Test = "1" julia = "1.10" -[sources] -# Temporary until OrdinaryDiffEq.jl#4614 is merged and released. -StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "77c07bffe2a225a95436190570eb8795f69342d0", subdir = "lib/StochasticDiffEqCore"} - [extras] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" diff --git a/docs/src/api.md b/docs/src/api.md index be77b9c7e..da472ce9d 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -296,7 +296,7 @@ backends. |:----------------------------------------------------- |:-------------------------------------------------------------------- |:------------------------------------------------------------- |:------------------- | | `SSAStepper` | JumpProcesses (`solve.jl`) | JumpProcesses (`SSA_stepper.jl`) | No | | OrdinaryDiffEq ODE/DAE (e.g., `Tsit5`, `DFBDF`) | JumpProcesses (`solve.jl`) | JumpProcesses' OrdinaryDiffEqCore extension → OrdinaryDiffEq | Yes | -| StochasticDiffEq SDE/RODE (e.g., `SRIW1`, `RandomEM`) | StochasticDiffEqCore | StochasticDiffEqCore's JumpProblem initializer | No | +| StochasticDiffEq SDE/RODE (e.g., `SRIW1`, `RandomEM`) | StochasticDiffEqCore | StochasticDiffEqCore's JumpProblem initializer (see below) | No | | StochasticDiffEq jump algorithms (e.g., `TauLeaping`) | StochasticDiffEqCore | StochasticDiffEqCore's specialized jump-algorithm initializer | No | | All five CPU tau-leaping algorithms | JumpProcesses (`simple_regular_solve.jl`, custom `DiffEqBase.solve`) | N/A | No | @@ -307,11 +307,27 @@ For **OrdinaryDiffEq ODE/DAE solvers**, `rng` is resolved via `resolve_rng` in ` (`solve.jl`) and forwarded to OrdinaryDiffEq's `init`, which stores it on the `ODEIntegrator`. -For **StochasticDiffEq SDE/RODE solvers**, the backend's more specific `__solve` -and `__init` methods handle the original JumpProblem and unchanged RNG/seed -inputs. Both routes use the backend's JumpProblem initializer, including its RNG -policy, jump-state copying, and callback setup. DiffEqBase merges stored problem -keywords and user callbacks before dispatching `init`. +For **StochasticDiffEq SDE/RODE solvers**, both `solve` and `init` reach the +backend's JumpProblem initializer, StochasticDiffEqCore's `_sde_init`, with the +original JumpProblem and unchanged RNG/seed inputs, so they share its RNG policy, +jump-state copying, and callback setup. `solve` goes through the backend's more +specific `__solve`. For `init`, StochasticDiffEqCore versions with per-algorithm +`__init(::JumpProblem, ...)` methods own dispatch directly; on other versions, +JumpProcesses' `JumpProcessesStochasticDiffEqCoreExt` extension adds one +`__init(::JumpProblem, ::Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm})` +method that forwards to `_sde_init`, resolving an otherwise ambiguous dispatch. The +backend's per-algorithm methods are more specific and take precedence wherever they +exist, as does its method for jump algorithms such as `TauLeaping`. + +DiffEqBase's `init_call` merges the JumpProblem's stored keywords and user callbacks +once, before dispatching to `__init`; JumpProcesses' `__init` methods do not merge +them again, and discard `merge_callbacks`. The `solve` path merges once in `__solve`. +Keywords stored on the wrapped problem depend on the pathway. On the `__jump_init` +pathways (OrdinaryDiffEq ODE/DAE solvers and `FunctionMap`), the wrapped problem's +own `init` merges them, callbacks included, exactly once, whatever `merge_callbacks` +is. `SSAStepper` and StochasticDiffEqCore's `_sde_init` do not merge callbacks +stored on the wrapped problem; on these pathways, store callbacks on the +`JumpProblem` or pass them to `solve`/`init`. Jump-state copying on these pathways follows the backend's alias specifier; see [Jump state ownership and problem reuse](@ref jump_state_ownership). In contrast, diff --git a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl index df907d6b0..edb0029bb 100644 --- a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl +++ b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl @@ -1,13 +1,13 @@ module JumpProcessesOrdinaryDiffEqCoreExt using JumpProcesses -import DiffEqBase import SciMLBase import OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, DAEAlgorithm, StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm +# `init_call` has already merged the problem's stored keywords before `__init`; see the +# generic `__init` in src/solve.jl. `merge_callbacks` is discarded, not forwarded. function _jump_init(_jump_prob, alg; merge_callbacks = true, kwargs...) - kwargs = DiffEqBase.merge_problem_kwargs(_jump_prob; merge_callbacks, kwargs...) JumpProcesses.__jump_init(_jump_prob, alg; kwargs...) end @@ -16,8 +16,9 @@ function SciMLBase.__init( alg::Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm, StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; kwargs...) where {IIP, P} - # Cover the full intersection with OrdinaryDiffEqCore's initializer. - # StochasticDiffEqCore's narrower methods own SDE/RODE initialization. + # Cover the full intersection with OrdinaryDiffEqCore's initializer. SDE/RODE + # initialization belongs to StochasticDiffEqCore: its per-algorithm methods where it + # has them, otherwise the forwarding method in JumpProcessesStochasticDiffEqCoreExt. _jump_init(_jump_prob, alg; kwargs...) end diff --git a/ext/JumpProcessesStochasticDiffEqCoreExt.jl b/ext/JumpProcessesStochasticDiffEqCoreExt.jl new file mode 100644 index 000000000..599f78185 --- /dev/null +++ b/ext/JumpProcessesStochasticDiffEqCoreExt.jl @@ -0,0 +1,25 @@ +module JumpProcessesStochasticDiffEqCoreExt + +using JumpProcesses: JumpProblem +import SciMLBase +import OrdinaryDiffEqCore: StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm +import StochasticDiffEqCore + +# Forward stochastic `init` of a `JumpProblem` to StochasticDiffEqCore's own initializer, +# `_sde_init`, which is exported for packages that need to bypass `__init` dispatch. On +# Core versions whose `__init` covers `Union{AbstractRODEProblem, JumpProblem}`, that +# method and the `OrdinaryDiffEqCore` extension's four-family method are each more +# specific in one argument, so without this method `init(jprob, sde_alg)` is ambiguous. +# Core versions with per-algorithm `JumpProblem` methods (OrdinaryDiffEq.jl#4614) are more +# specific than this union method and take precedence, and Core's jump-algorithm method +# (e.g. `TauLeaping`) is also more specific. Keep the signature distinct from Core's: do +# not split it into separate SDE and RODE methods. +# +# `init_call` has already merged the `JumpProblem`'s stored keywords, and `_sde_init` +# documents that they must not be merged again. +function SciMLBase.__init(prob::JumpProblem, + alg::Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; kwargs...) + StochasticDiffEqCore._sde_init(prob, alg; kwargs...) +end + +end diff --git a/src/solve.jl b/src/solve.jl index 58a304654..b6a5553de 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -57,12 +57,14 @@ function SciMLBase.__solve(jump_prob::JumpProblem; kwargs...) error("Auto-solver selection is currently only implemented for JumpProblems defined over DiscreteProblems. Please explicitly specify a solver algorithm in calling solve.") end +# Public `init` reaches `__init` through DiffEqBase's `init_call`, which has already merged +# the problem's stored keywords (honoring `merge_callbacks`), so they are not merged again +# here. `merge_callbacks` is accepted and discarded so that it is not forwarded to the +# wrapped problem's `init`, whose own stored keywords are merged there exactly once, as on +# the `solve` path. function SciMLBase.__init(_jump_prob::JumpProblem{IIP, P}, alg::SciMLBase.AbstractDEAlgorithm; merge_callbacks = true, kwargs...) where { IIP, P} - # Merge jump_prob.kwargs with passed kwargs - kwargs = DiffEqBase.merge_problem_kwargs(_jump_prob; merge_callbacks, kwargs...) - __jump_init(_jump_prob, alg; kwargs...) end diff --git a/test/callbacks.jl b/test/callbacks.jl index b7ce5614e..b656e57bc 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -1,4 +1,4 @@ -using JumpProcesses, OrdinaryDiffEq, SciMLBase, Test +using JumpProcesses, OrdinaryDiffEq, SciMLBase, DiffEqBase, Test using StableRNGs rng = StableRNG(12345) @@ -515,3 +515,86 @@ end n_continuous = length(integrator_vr.opts.callback.continuous_callbacks) @test n_continuous == 1 end + +# Callbacks can be stored on the wrapped problem, stored on the `JumpProblem`, or passed +# at the call site. `init` must combine them exactly as `solve` does: DiffEqBase's +# `init_call` merges the `JumpProblem`'s stored keywords once (honoring +# `merge_callbacks`). On the `__jump_init` paths (ODE and `FunctionMap`), the wrapped +# problem's own stored keywords are merged once by its inner `init`; `SSAStepper` and +# the StochasticDiffEqCore backend do not merge callbacks stored on the wrapped problem. +@testset "Stored and call-level callbacks: init matches solve" begin + using OrdinaryDiffEqFunctionMap, StochasticDiffEq + log = Symbol[] + logging_callback(name) = DiscreteCallback((u, t, integrator) -> true, + integrator -> (push!(log, name); nothing); save_positions = (false, false)) + jump = ConstantRateJump((u, p, t) -> 1.0, integrator -> (integrator.u[1] += 1; nothing)) + noop!(du, u, p, t) = (du .= 0; nothing) + noop_rode!(du, u, p, t, W) = (du .= 0; nothing) + wrapped = logging_callback(:wrapped) + ode = ODEProblem(noop!, [0.0], (0.0, 0.8); callback = wrapped) + discrete = DiscreteProblem([0], (0.0, 8.0); callback = wrapped) + sde = SDEProblem(noop!, noop!, [0.0], (0.0, 0.8); callback = wrapped) + rode = RODEProblem(noop_rode!, [0.0], (0.0, 0.8); callback = wrapped) + cases = (("ODE", ode, Tsit5(), (; adaptive = false, dt = 0.1)), + ("FunctionMap", discrete, FunctionMap(), (;)), + ("SSAStepper", discrete, SSAStepper(), (;)), + ("SDE", sde, EM(), (; dt = 0.1)), + ("RODE", rode, RandomEM(), (; dt = 0.1))) + + function run_and_log(entry, jprob, alg; kwargs...) + empty!(log) + if entry === :solve + solve(jprob, alg; kwargs...) + else + solve!(init(jprob, alg; kwargs...)) + end + copy(log) + end + + @testset "$name, merge_callbacks = $merge_callbacks" for (name, prob, alg, options) in cases, + merge_callbacks in (true, false) + + jprob = JumpProblem(prob, Direct(), jump; callback = logging_callback(:jumpproblem)) + call = logging_callback(:call) + kwargs = (; options..., callback = call, merge_callbacks, seed = 1) + solved = run_and_log(:solve, jprob, alg; kwargs...) + initialized = run_and_log(:init, jprob, alg; kwargs...) + # Same callbacks, same number of runs, same order. + @test initialized == solved + steps = count(==(:call), solved) + @test steps > 0 + @test count(==(:jumpproblem), initialized) == (merge_callbacks ? steps : 0) + if name in ("ODE", "FunctionMap") + # The wrapped problem's callback is merged once by its inner `init`, whatever + # `merge_callbacks` says about the `JumpProblem`'s stored callback. + @test count(==(:wrapped), initialized) == steps + end + end + + @testset "Direct __init discards merge_callbacks" begin + prob = ODEProblem(noop!, [0.0], (0.0, 0.8); callback = logging_callback(:wrapped)) + jprob = JumpProblem(prob, Direct(), jump; callback = logging_callback(:jumpproblem)) + options = (; adaptive = false, dt = 0.1, seed = 1) + expected = run_and_log(:init, jprob, Tsit5(); options..., + callback = logging_callback(:call)) + # `__init` receives keywords that are already merged; a stray `merge_callbacks` + # must neither merge again nor reach the wrapped problem's `init`. + merged = DiffEqBase.merge_problem_kwargs(jprob; callback = logging_callback(:call)) + empty!(log) + solve!(SciMLBase.__init(jprob, Tsit5(); options..., merged..., + merge_callbacks = false)) + @test log == expected + end + + @testset "Stored tstops are honored once" begin + prob = ODEProblem(noop!, [0.0], (0.0, 1.0)) + jprob = JumpProblem(prob, Direct(), jump; tstops = [0.55]) + for sol in (solve(jprob, Tsit5(); seed = 1), + solve!(init(jprob, Tsit5(); seed = 1))) + @test count(==(0.55), sol.t) >= 1 + end + integrator = init(jprob, Tsit5(); seed = 1) + solve!(integrator) + @test integrator.sol.t == solve(jprob, Tsit5(); seed = 1).t + end +end diff --git a/test/extended_jump_array_remake.jl b/test/extended_jump_array_remake.jl index 7a32c64d7..5f6e4add3 100644 --- a/test/extended_jump_array_remake.jl +++ b/test/extended_jump_array_remake.jl @@ -1,7 +1,7 @@ # Tests for JumpProblem remake with VariableRateJumps (ExtendedJumpArray case) # This tests the fix for symbolic u0 with ExtendedJumpArray -using JumpProcesses, OrdinaryDiffEq, Test, SymbolicIndexingInterface +using JumpProcesses, OrdinaryDiffEq, SciMLBase, Test, SymbolicIndexingInterface using StableRNGs @testset "remake JumpProblem with VariableRateJumps (ExtendedJumpArray)" begin diff --git a/test/rng_kwarg_tests.jl b/test/rng_kwarg_tests.jl index 1728b03db..11f91b82f 100644 --- a/test/rng_kwarg_tests.jl +++ b/test/rng_kwarg_tests.jl @@ -1,4 +1,4 @@ -using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq, Test +using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq, SciMLBase, Test using StableRNGs, Random # ========================================================================== diff --git a/test/stochastic_dispatch.jl b/test/stochastic_dispatch.jl index 9f9243351..439de4600 100644 --- a/test/stochastic_dispatch.jl +++ b/test/stochastic_dispatch.jl @@ -23,12 +23,57 @@ function stochastic_rng_options(kind) (; rng = StableRNG(789), seed = 456) end +const SDECore = StochasticDiffEq.StochasticDiffEqCore +const SDE_BRIDGE = Base.get_extension(JumpProcesses, :JumpProcessesStochasticDiffEqCoreExt) +const BRIDGE_SIGNATURE = Tuple{typeof(SciMLBase.__init), JumpProblem, + Union{SDECore.StochasticDiffEqAlgorithm, SDECore.StochasticDiffEqRODEAlgorithm}} + +# Whether StochasticDiffEqCore defines its own `__init(::JumpProblem, ::A)` for the +# algorithm supertype `A`, as Core versions with OrdinaryDiffEq.jl#4614 do. Decided from +# the method table rather than Core's version number. +function core_owns_jump_init(A) + signature = Tuple{typeof(SciMLBase.__init), JumpProblem, A} + any(m -> m.module === SDECore && m.sig == signature, methods(SciMLBase.__init)) +end + +# Whether a method is the forwarding method, or a `__init`, `__solve`, or +# `supports_solve_rng` method for `JumpProblem`s. +function involves_jump_init(m) + m.module === SDE_BRIDGE && return true + m.name in (:__init, :__solve, :supports_solve_rng) || return false + any(T -> T isa Type && T <: JumpProblem, Base.unwrap_unionall(m.sig).parameters) +end + @testset "Stochastic initialization dispatches to the backend" begin - for (kind, alg) in ((:SDE, EM()), (:RODE, RandomEM())) + @test SDE_BRIDGE !== nothing + for (kind, alg, A) in ((:SDE, EM(), SDECore.StochasticDiffEqAlgorithm), + (:RODE, RandomEM(), SDECore.StochasticDiffEqRODEAlgorithm)) jprob = stochastic_jump_problem(kind) method = which(SciMLBase.__init, (typeof(jprob), typeof(alg))) - @test method.module === StochasticDiffEq.StochasticDiffEqCore + if core_owns_jump_init(A) + # Core's per-algorithm method is more specific than the forwarding method. + @test method.module === SDECore + @test method.sig == Tuple{typeof(SciMLBase.__init), JumpProblem, A} + else + # Otherwise JumpProcesses' extension forwards to Core's `_sde_init`. + @test method.module === SDE_BRIDGE + @test method.sig == BRIDGE_SIGNATURE + end end + + # The forwarding method exists once, owned by the extension: Core never overwrites it. + bridge_methods = filter(m -> m.sig == BRIDGE_SIGNATURE, collect(methods(SciMLBase.__init))) + @test length(bridge_methods) == 1 + @test only(bridge_methods).module === SDE_BRIDGE + + # Core's jump-algorithm method (e.g. `TauLeaping`) is more specific than the bridge. + jprob = stochastic_jump_problem(:SDE) + @test which(SciMLBase.__init, (typeof(jprob), typeof(TauLeaping()))).module === SDECore + + # No ambiguity involves the forwarding method or a `JumpProblem` method of `__init`, + # `__solve`, or `supports_solve_rng`. + ambiguities = Test.detect_ambiguities(JumpProcesses, SDECore, SDE_BRIDGE) + @test !any(pair -> any(involves_jump_init, pair), ambiguities) end @testset "Stochastic solve/init use the same RNG policy" begin From d73ff5fa48cba9069d7c32be9b2832800e4c49b9 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 10 Oct 2026 09:36:03 -0400 Subject: [PATCH 40/52] Reject callbacks stored on the problem wrapped by a JumpProblem Callbacks stored on the wrapped problem, for example ODEProblem(f, u0, tspan, p; callback), ran on OrdinaryDiffEq solvers, including FunctionMap, through the wrapped problem's own init, but SSAStepper and StochasticDiffEq solvers silently ignored them. Rather than behave differently by solver, the JumpProblem constructors, and remake with a new wrapped problem, now raise an ArgumentError when the wrapped problem stores a callback other than nothing or an empty CallbackSet. The message points to JumpProblem(...; callback), solve/init, and remake(prob; callback = nothing). The check runs before spatial flattening, which rebuilds the problem without its keywords. ModelingToolkitBase and Catalyst attach their events to the outer JumpProblem, so the problems they build are unaffected. The JumpProblem and remake docstrings, the JumpProblem documentation, the FAQ, the migration guide, and HISTORY state the rule. The init/solve callback tests no longer store callbacks on the wrapped problem. New tests cover the rejection on each constructor path, the accepted empty forms, the error message, and the remake remedy. Validation on registered dependencies, Julia 1.12.7 and 1.10.12: InterfaceI (1,861 assertions), InterfaceII (3,344), four-thread ThreadSafety (333), the stochastic suite on Core 2.2.1 and OrdinaryDiffEq.jl#4614's revision (824 each), the downgrade All group on the action-derived minimum graph (5,540), strict QA (20, plus 5 registry checks), the package-only minimum graph, formatting, spelling and whitespace checks, and a Documenter build of the changed pages pass. Co-Authored-By: Claude Opus 5.5 (1M context) --- HISTORY.md | 20 ++++++- docs/src/api.md | 11 ++-- docs/src/faq.md | 10 ++++ docs/src/jump_types.md | 15 +++++ docs/src/migration.md | 22 ++++++++ src/problem.jl | 45 +++++++++++++++ test/callbacks.jl | 123 +++++++++++++++++++++++++++++++++++------ 7 files changed, 220 insertions(+), 26 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 81543940f..06049ba17 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -58,6 +58,22 @@ changes are not intended for a 9.x release. `SciMLBase.RODEAliasSpecifier` for RODEs) to copy jump state, or `alias_jumps = true` to reuse it. When unspecified, those solvers reuse jump state on thread 1 and copy it on other threads. + - **Breaking**: Callbacks stored on the problem wrapped by a `JumpProblem` are not + supported. The `JumpProblem` constructor, and `remake` with a new wrapped problem, + now raise an `ArgumentError` if that problem stores a `callback` (other than + `nothing` or an empty `CallbackSet`). Previously, OrdinaryDiffEq solvers + (including `FunctionMap`) ran such callbacks, while `SSAStepper` and + StochasticDiffEq SDE/RODE solvers silently ignored them. Pass callbacks to the + `JumpProblem` constructor or to `solve`/`init` instead: + ```julia + # Before (no longer works): + oprob = ODEProblem(f, u0, tspan, p; callback = cb) + jprob = JumpProblem(oprob, Direct(), jump) + # After: + oprob = ODEProblem(f, u0, tspan, p) + jprob = JumpProblem(oprob, Direct(), jump; callback = cb) + ``` + `remake(prob; callback = nothing)` removes callbacks from an existing problem. - `SSAStepper` now supports the common `alias` keyword argument for its `u0`, `p`, and `tstops` inputs, given as a `Bool`, a `SciMLBase.DiscreteAliasSpecifier`, or an `ODEAliasSpecifier` (which also controls `tstops`). The defaults match @@ -90,8 +106,8 @@ changes are not intended for a 9.x release. - Fixed `init` on ODE, DAE, and `FunctionMap` integrators for `JumpProblem`s: callbacks stored in the `JumpProblem` (`JumpProblem(...; callback)`) were merged twice, so they ran twice per step, and were reintroduced even with - `merge_callbacks = false`. `init` now combines stored, wrapped-problem, and - call-level callbacks exactly as `solve` does. + `merge_callbacks = false`. `init` now combines stored and call-level callbacks + exactly as `solve` does. - StochasticDiffEqCore is now a weak dependency: when it is installed, version 2.2.1 or later is required. A new extension sends stochastic `init` of a `JumpProblem` to StochasticDiffEqCore's own initializer on every supported version, diff --git a/docs/src/api.md b/docs/src/api.md index da472ce9d..026babd1e 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -322,12 +322,11 @@ exist, as does its method for jump algorithms such as `TauLeaping`. DiffEqBase's `init_call` merges the JumpProblem's stored keywords and user callbacks once, before dispatching to `__init`; JumpProcesses' `__init` methods do not merge them again, and discard `merge_callbacks`. The `solve` path merges once in `__solve`. -Keywords stored on the wrapped problem depend on the pathway. On the `__jump_init` -pathways (OrdinaryDiffEq ODE/DAE solvers and `FunctionMap`), the wrapped problem's -own `init` merges them, callbacks included, exactly once, whatever `merge_callbacks` -is. `SSAStepper` and StochasticDiffEqCore's `_sde_init` do not merge callbacks -stored on the wrapped problem; on these pathways, store callbacks on the -`JumpProblem` or pass them to `solve`/`init`. +Only the `__jump_init` pathways (OrdinaryDiffEq ODE/DAE solvers and `FunctionMap`) +read keywords stored on the wrapped problem, through its own `init`; `SSAStepper` and +StochasticDiffEqCore's `_sde_init` never do. So that every solver runs the same +callbacks, the `JumpProblem` constructor, and `remake` with a new wrapped problem, +reject a wrapped problem that stores a `callback`. Jump-state copying on these pathways follows the backend's alias specifier; see [Jump state ownership and problem reuse](@ref jump_state_ownership). In contrast, diff --git a/docs/src/faq.md b/docs/src/faq.md index 00c7008e4..e2508c49a 100644 --- a/docs/src/faq.md +++ b/docs/src/faq.md @@ -147,6 +147,16 @@ Callbacks can be used with `ConstantRateJump`s, `MassActionJump`s, and discrete callbacks can be used (otherwise a different time stepper is needed). When using an ODE or SDE time stepper, any callback should work. +Pass callbacks to `solve` or `init`, as in the example below, or to the +`JumpProblem` constructor with its `callback` keyword. + +!!! warning "Callbacks on the wrapped problem are not supported" + + Do not store callbacks on the problem a `JumpProblem` wraps, such as the + `DiscreteProblem` or `ODEProblem` passed to `JumpProblem`. The `JumpProblem` + constructor raises an `ArgumentError` if that problem stores a `callback`; see + [Defining a `JumpProblem`](@ref defining_jump_problem). + *Note, when modifying `u` or `p` within a callback, you must call [`reset_aggregated_jumps!`](@ref) after making updates.* This ensures that the underlying jump simulation algorithms know to reinitialize their internal data diff --git a/docs/src/jump_types.md b/docs/src/jump_types.md index 65b42a4c1..b018421bd 100644 --- a/docs/src/jump_types.md +++ b/docs/src/jump_types.md @@ -278,6 +278,21 @@ and the internals will automatically build the `JumpSet`. `save_positions` determines whether to save the state of the system just before and/or after jumps occur. +Callbacks to run in every solve can be passed to the constructor with the +`callback` keyword, `JumpProblem(prob, aggregator, jumps...; callback)`, and +further callbacks can be passed to `solve` or `init`. + +!!! warning "Callbacks on the wrapped problem are not supported" + + Do not store callbacks on `prob` itself, for example with + `ODEProblem(f, u0, tspan, p; callback)` or `DiscreteProblem(u0, tspan, p; callback)`. + The `JumpProblem` constructor raises an `ArgumentError` if `prob` stores a + `callback`. Before JumpProcesses 10, OrdinaryDiffEq solvers (including + `FunctionMap`) ran such callbacks, while `SSAStepper` and StochasticDiffEq solvers + silently ignored them. Pass the callbacks to the `JumpProblem` constructor or to + `solve`/`init` instead, and use `remake(prob; callback = nothing)` to remove them + from an existing problem. + Note that a `JumpProblem`/`JumpSet` can only have 1 `RegularJump` (since a `RegularJump` itself describes multiple processes together). Similarly, it can only have one `MassActionJump` (since it also describes multiple processes diff --git a/docs/src/migration.md b/docs/src/migration.md index ec7d6e2ea..a0d84997a 100644 --- a/docs/src/migration.md +++ b/docs/src/migration.md @@ -219,3 +219,25 @@ sol = solve(jprob, SSAStepper(); callback = cb, tstops = 0.0:1.0:10.0) [Saving with callbacks and evaluating the integrator](@ref ssa_integrator_evaluation) for when to choose each remedy, and [Supported state types](@ref ssa_state_types) for the state types `SSAStepper` supports. + +## Callbacks on the wrapped problem + +Callbacks stored on the problem a `JumpProblem` wraps were run by OrdinaryDiffEq +solvers (including `FunctionMap`), but silently ignored by `SSAStepper` and +StochasticDiffEq solvers. They are no longer supported: the `JumpProblem` +constructor, and `remake` with a new wrapped problem, raise an `ArgumentError` if +that problem stores a `callback`. Pass the callbacks to the `JumpProblem` +constructor, or to `solve`/`init`, instead: + +```julia +# Before (no longer works): +oprob = ODEProblem(f, u0, tspan, p; callback = cb) +jprob = JumpProblem(oprob, Direct(), jump) + +# After: +oprob = ODEProblem(f, u0, tspan, p) +jprob = JumpProblem(oprob, Direct(), jump; callback = cb) +# or pass `callback = cb` to `solve`/`init`. +``` + +For an existing problem, `remake(oprob; callback = nothing)` removes its callbacks. diff --git a/src/problem.jl b/src/problem.jl index 66468b81d..27a24253f 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -65,6 +65,18 @@ $(FIELDS) default is `VR_FRM`. - `tstops`, time stops to pass through to the solver. Can be an `AbstractVector` of times or a callable `(p, tspan) -> times`. + - `callback`, callbacks to run in every solve of the problem. They are combined with any + `callback` passed to `solve` or `init`, which replaces them instead when + `merge_callbacks = false` is also passed. Do not store callbacks on the wrapped problem + `prob`; see the warning below. + +!!! warning "Callbacks on the wrapped problem are not supported" + + Pass callbacks to the `JumpProblem` constructor, with the `callback` keyword, or to + `solve`/`init`. Do not store them on the wrapped problem, for example with + `ODEProblem(f, u0, tspan, p; callback)`: the constructor raises an `ArgumentError` if + `prob` stores a `callback`, since only some solvers would run it. Use + `remake(prob; callback = nothing)` to remove callbacks from an existing problem. Please see the [tutorial page](https://docs.sciml.ai/JumpProcesses/stable/tutorials/discrete_stochastic_example/) in @@ -141,8 +153,35 @@ mutable struct JumpProblem{iip, P, A, C, J <: Union{Nothing, AbstractJumpAggrega """ kwargs::K end + +# Callbacks stored on the wrapped problem would run only on solvers that `init` the wrapped +# problem (OrdinaryDiffEq, including `FunctionMap`); `SSAStepper` and StochasticDiffEq never +# read them. Reject them, so that every solver runs the same callbacks. +function check_no_wrapped_callbacks(prob) + if hasproperty(prob, :kwargs) && haskey(prob.kwargs, :callback) && + !is_empty_callback(prob.kwargs[:callback]) + throw(ArgumentError(""" + Callbacks stored on the problem wrapped by a `JumpProblem` are not supported, \ + but the wrapped `$(nameof(typeof(prob)))` stores a `callback` keyword argument. \ + OrdinaryDiffEq solvers would run these callbacks, while `SSAStepper` and \ + StochasticDiffEq solvers would ignore them. Pass the callbacks to the \ + `JumpProblem` constructor instead, `JumpProblem(prob, aggregator, jumps...; \ + callback)`, or to `solve`/`init`. To remove callbacks from an existing problem, \ + use `remake(prob; callback = nothing)`.""")) + end + nothing +end + +is_empty_callback(::Nothing) = true +function is_empty_callback(cb::CallbackSet) + isempty(cb.continuous_callbacks) && isempty(cb.discrete_callbacks) +end +is_empty_callback(_) = false + +# Also used by `remake`, so a new wrapped problem is checked too. function JumpProblem(p::P, a::A, dj::J, jc::C, cj::J1, vj::J2, rj::J3, mj::J4, kwargs::K) where {P, A, J, C, J1, J2, J3, J4, K} + check_no_wrapped_callbacks(p) iip = isinplace_jump(p, rj) JumpProblem{iip, P, A, C, J, J1, J2, J3, J4, K}(p, a, dj, jc, cj, vj, rj, mj, kwargs) @@ -166,6 +205,9 @@ Base.@pure remaker_of(prob::T) where {T <: JumpProblem} = SciMLBase.parameterles Construct a `JumpProblem` with a new initial condition `u0`, parameters `p`, time span `tspan`, or wrapped problem `prob`. +A new wrapped problem `prob` must not store callbacks, as for the `JumpProblem` +constructor; an `ArgumentError` is raised if it does. + The new problem shares the original's jump aggregator and jump callbacks rather than copying them. Solvers owned by JumpProcesses re-initialize this state from the problem being solved at every `init`, so the original and remade problems can be solved one after @@ -301,6 +343,8 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS throw(ArgumentError("`useiszero` is no longer a keyword argument for `JumpProblem`. Set `useiszero` on the `MassActionJump` directly instead.")) end haskey(kwargs, :alias_jump) && throw(ArgumentError(ALIAS_JUMP_REMOVED_MSG)) + # Check before spatial flattening, which rebuilds the problem without its keywords. + check_no_wrapped_callbacks(prob) # keep the original MAJ (including {Nothing} parameterized MAJs); # fill_scaled_rates! in each aggregator's initialize! handles rate setup @@ -381,6 +425,7 @@ function JumpProblem(prob, aggregator::PureLeaping, jumps::JumpSet; throw(ArgumentError("`useiszero` is no longer a keyword argument for `JumpProblem`. Set `useiszero` on the `MassActionJump` directly instead.")) end haskey(kwargs, :alias_jump) && throw(ArgumentError(ALIAS_JUMP_REMOVED_MSG)) + check_no_wrapped_callbacks(prob) # Validate no spatial systems (not currently supported) (spatial_system !== nothing || hopping_constants !== nothing) && diff --git a/test/callbacks.jl b/test/callbacks.jl index b656e57bc..af5c20864 100644 --- a/test/callbacks.jl +++ b/test/callbacks.jl @@ -516,12 +516,10 @@ end @test n_continuous == 1 end -# Callbacks can be stored on the wrapped problem, stored on the `JumpProblem`, or passed -# at the call site. `init` must combine them exactly as `solve` does: DiffEqBase's -# `init_call` merges the `JumpProblem`'s stored keywords once (honoring -# `merge_callbacks`). On the `__jump_init` paths (ODE and `FunctionMap`), the wrapped -# problem's own stored keywords are merged once by its inner `init`; `SSAStepper` and -# the StochasticDiffEqCore backend do not merge callbacks stored on the wrapped problem. +# Callbacks can be stored on the `JumpProblem` or passed at the call site (callbacks +# stored on the wrapped problem are rejected; see the next testset). `init` must combine +# them exactly as `solve` does: DiffEqBase's `init_call` merges the `JumpProblem`'s stored +# keywords once, honoring `merge_callbacks`. @testset "Stored and call-level callbacks: init matches solve" begin using OrdinaryDiffEqFunctionMap, StochasticDiffEq log = Symbol[] @@ -530,11 +528,10 @@ end jump = ConstantRateJump((u, p, t) -> 1.0, integrator -> (integrator.u[1] += 1; nothing)) noop!(du, u, p, t) = (du .= 0; nothing) noop_rode!(du, u, p, t, W) = (du .= 0; nothing) - wrapped = logging_callback(:wrapped) - ode = ODEProblem(noop!, [0.0], (0.0, 0.8); callback = wrapped) - discrete = DiscreteProblem([0], (0.0, 8.0); callback = wrapped) - sde = SDEProblem(noop!, noop!, [0.0], (0.0, 0.8); callback = wrapped) - rode = RODEProblem(noop_rode!, [0.0], (0.0, 0.8); callback = wrapped) + ode = ODEProblem(noop!, [0.0], (0.0, 0.8)) + discrete = DiscreteProblem([0], (0.0, 8.0)) + sde = SDEProblem(noop!, noop!, [0.0], (0.0, 0.8)) + rode = RODEProblem(noop_rode!, [0.0], (0.0, 0.8)) cases = (("ODE", ode, Tsit5(), (; adaptive = false, dt = 0.1)), ("FunctionMap", discrete, FunctionMap(), (;)), ("SSAStepper", discrete, SSAStepper(), (;)), @@ -564,21 +561,16 @@ end steps = count(==(:call), solved) @test steps > 0 @test count(==(:jumpproblem), initialized) == (merge_callbacks ? steps : 0) - if name in ("ODE", "FunctionMap") - # The wrapped problem's callback is merged once by its inner `init`, whatever - # `merge_callbacks` says about the `JumpProblem`'s stored callback. - @test count(==(:wrapped), initialized) == steps - end end @testset "Direct __init discards merge_callbacks" begin - prob = ODEProblem(noop!, [0.0], (0.0, 0.8); callback = logging_callback(:wrapped)) + prob = ODEProblem(noop!, [0.0], (0.0, 0.8)) jprob = JumpProblem(prob, Direct(), jump; callback = logging_callback(:jumpproblem)) options = (; adaptive = false, dt = 0.1, seed = 1) expected = run_and_log(:init, jprob, Tsit5(); options..., callback = logging_callback(:call)) # `__init` receives keywords that are already merged; a stray `merge_callbacks` - # must neither merge again nor reach the wrapped problem's `init`. + # must not merge them again. merged = DiffEqBase.merge_problem_kwargs(jprob; callback = logging_callback(:call)) empty!(log) solve!(SciMLBase.__init(jprob, Tsit5(); options..., merged..., @@ -598,3 +590,98 @@ end @test integrator.sol.t == solve(jprob, Tsit5(); seed = 1).t end end + +# Callbacks stored on the wrapped problem would run only on solvers that `init` the +# wrapped problem (OrdinaryDiffEq), and be ignored by `SSAStepper` and StochasticDiffEq, so +# the `JumpProblem` constructors and `remake` reject them. +@testset "Callbacks stored on the wrapped problem are rejected" begin + using StochasticDiffEq + function rejects_wrapped_callbacks(f) + err = try + f() + nothing + catch e + e + end + err isa ArgumentError && occursin("are not supported", err.msg) + end + discrete_cb = DiscreteCallback((u, t, integrator) -> false, integrator -> nothing) + continuous_cb = ContinuousCallback((u, t, integrator) -> t - 0.5, integrator -> nothing) + stored_callbacks = (discrete_cb, continuous_cb, CallbackSet(discrete_cb)) + crj = ConstantRateJump((u, p, t) -> 1.0, integrator -> (integrator.u[1] += 1; nothing)) + vrj = VariableRateJump((u, p, t) -> 1.0, integrator -> (integrator.u[1] += 1; nothing)) + maj = MassActionJump([1.0], [[1 => 1]], [[1 => -1]]) + rj = RegularJump((out, u, p, t) -> (out[1] = 1.0), + (du, u, p, t, counts, mark) -> (du[1] = counts[1]), 1) + noop!(du, u, p, t) = (du .= 0; nothing) + noop_rode!(du, u, p, t, W) = (du .= 0; nothing) + tspan = (0.0, 1.0) + discrete(kw) = DiscreteProblem([10], tspan; kw...) + leaping(kw) = DiscreteProblem([10.0], tspan; kw...) + spatial(kw) = DiscreteProblem([10 0], tspan; kw...) # one species on two sites + ode(kw) = ODEProblem(noop!, [0.0], tspan; kw...) + sde(kw) = SDEProblem(noop!, noop!, [0.0], tspan; kw...) + rode(kw) = RODEProblem(noop_rode!, [0.0], tspan; kw...) + hopping_constants = fill(1.0, 1, 2) + spatial_options = (; hopping_constants, spatial_system = CartesianGrid((2,))) + + # Each case wraps a problem through a different constructor path: its name, the + # wrapped-problem builder, and the remaining `JumpProblem` arguments and keywords. + cases = (("DiscreteProblem, Direct", discrete, (Direct(), crj), (;)), + ("DiscreteProblem, selected aggregator", discrete, (maj,), (;)), + ("DiscreteProblem, PureLeaping", leaping, (PureLeaping(), rj), (;)), + # Flattening rebuilds the problem without its keywords, so the check must precede it. + ("spatial DiscreteProblem, flattened", spatial, (Direct(), maj), spatial_options), + ("spatial DiscreteProblem, NSM", spatial, (NSM(), maj), spatial_options), + ("ODEProblem, Direct", ode, (Direct(), crj), (;)), + ("ODEProblem, VR_FRM", ode, (vrj,), (; vr_aggregator = VR_FRM())), + ("ODEProblem, VR_Direct", ode, (vrj,), (; vr_aggregator = VR_Direct())), + ("SDEProblem", sde, (Direct(), crj), (;)), + ("RODEProblem", rode, (Direct(), crj), (;))) + + @testset "$name" for (name, problem, args, options) in cases + build(kw) = JumpProblem(problem(kw), args...; options...) + for callback in stored_callbacks + @test rejects_wrapped_callbacks(() -> build((; callback))) + end + # No stored callback, `nothing`, or an empty `CallbackSet` is accepted. + for kw in ((;), (; callback = nothing), (; callback = CallbackSet())) + @test build(kw) isa JumpProblem + end + end + + @testset "The error names the wrapped problem type" begin + err = try + JumpProblem(ode((; callback = discrete_cb)), Direct(), crj) + catch e + e + end + @test err isa ArgumentError + @test occursin("wrapped `ODEProblem`", err.msg) + @test occursin("remake(prob; callback = nothing)", err.msg) + end + + @testset "Accepted problems solve" begin + jprob = JumpProblem(discrete((; callback = CallbackSet())), Direct(), crj) + @test SciMLBase.successful_retcode(solve(jprob, SSAStepper(); seed = 1)) + jprob = JumpProblem(ode((; callback = nothing)), Direct(), crj) + @test SciMLBase.successful_retcode(solve(jprob, Tsit5(); seed = 1)) + end + + @testset "remake(prob; callback = nothing) removes stored callbacks" begin + for problem in (discrete, ode, sde, rode) + prob = problem((; callback = discrete_cb)) + @test rejects_wrapped_callbacks(() -> JumpProblem(prob, Direct(), crj)) + stripped = remake(prob; callback = nothing) + @test JumpProblem(stripped, Direct(), crj) isa JumpProblem + end + end + + @testset "remake checks a new wrapped problem" begin + jprob = JumpProblem(discrete((;)), Direct(), crj) + newprob = DiscreteProblem([5], tspan; callback = discrete_cb) + @test rejects_wrapped_callbacks(() -> remake(jprob; prob = newprob)) + @test remake(jprob; prob = DiscreteProblem([5], tspan)).prob.u0 == [5] + @test remake(jprob; u0 = [5]).prob.u0 == [5] + end +end From 6309b434b970e2243b0cd1fd12b51733e5d9b93d Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 10 Oct 2026 10:45:51 -0400 Subject: [PATCH 41/52] Run benchmarks on both APIs and mark Float32 GPU explicit tau broken The benchmark job runs benchmark/benchmarks.jl against both the PR and master. The script used v10's solve-time rng, which master's SSAStepper rejects, so both benchmark jobs failed while benchmarking master. The script now detects which API it runs against, since JumpProblem(...; rng) raises an ArgumentError under v10. On master it reseeds the problem's StableRNG with seed in every sample, so both branches simulate the same trajectory and their timings are comparable. On CUDA, the explicit tau-leaping kernel fails to compile for Float32 rates. pois_rand uses PoissonRandom's procf for larger rates, and procf is not type-stable for Float32: two of its results are Union{Float32, Float64}. Master's kernel makes the same call, but v10's new kernel rate test is the first to run it with Float32 on a GPU, and the error also stopped that testset before its Float64 pass. The Float32 GPU solve is now @test_broken, so the Float64 pass runs and an upstream fix shows up as an unexpected pass. Co-Authored-By: Claude Opus 5.5 (1M context) --- benchmark/benchmarks.jl | 50 ++++++++++++++++++++++++++--------- test/kernel_contract_tests.jl | 25 ++++++++++++------ 2 files changed, 54 insertions(+), 21 deletions(-) diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl index 71e731c72..292ba7b8e 100644 --- a/benchmark/benchmarks.jl +++ b/benchmark/benchmarks.jl @@ -41,17 +41,43 @@ SUITE["construct"]["jumpproblem_constantrate"] = @benchmarkable JumpProblem( $dprob, Direct(), $jump1, $jump2 ) +# The benchmark job runs this script against both this branch and the base branch. +# JumpProcesses 10 takes the RNG at `solve`, and `JumpProblem(...; rng)` raises an +# `ArgumentError`; earlier versions store it in the `JumpProblem` and reject it in `solve`. +# Either way, every sample simulates the same trajectory, from a `StableRNG` seeded with +# 12345, so the two branches' timings are comparable. +const RNG_AT_SOLVE = try + JumpProblem(dprob, Direct(), maj; rng = StableRNG(1)) + false +catch err + err isa ArgumentError || rethrow() + true +end + +if RNG_AT_SOLVE + jump_problem(args...) = JumpProblem(args...) + function benchmark_solve(jprob) + @benchmarkable solve($jprob, SSAStepper(); rng) setup=(rng=StableRNG(12345)) evals=1 + end +else + # `seed` reseeds the problem's RNG. + jump_problem(args...) = JumpProblem(args...; rng = StableRNG(12345)) + function benchmark_solve(jprob) + @benchmarkable solve($jprob, SSAStepper(); seed = 12345) evals=1 + end +end + # ============================================================================= -# Solves (SSAStepper, fresh seeded RNG for each sample) +# Solves (SSAStepper, seeded RNG) # ============================================================================= SUITE["solve"] = BenchmarkGroup() -jprob_maj = JumpProblem(dprob, Direct(), maj) -jprob_cr = JumpProblem(dprob, Direct(), jump1, jump2) +jprob_maj = jump_problem(dprob, Direct(), maj) +jprob_cr = jump_problem(dprob, Direct(), jump1, jump2) -SUITE["solve"]["massaction"] = @benchmarkable solve($jprob_maj, SSAStepper(); rng) setup=(rng=StableRNG(12345)) evals=1 -SUITE["solve"]["constantrate"] = @benchmarkable solve($jprob_cr, SSAStepper(); rng) setup=(rng=StableRNG(12345)) evals=1 +SUITE["solve"]["massaction"] = benchmark_solve(jprob_maj) +SUITE["solve"]["constantrate"] = benchmark_solve(jprob_cr) # ============================================================================= # Aggregators @@ -59,12 +85,10 @@ SUITE["solve"]["constantrate"] = @benchmarkable solve($jprob_cr, SSAStepper(); r SUITE["aggregators"] = BenchmarkGroup() -jprob_rdirect = JumpProblem(dprob, RDirect(), maj) -jprob_sorting = JumpProblem(dprob, SortingDirect(), maj) -jprob_nrm = JumpProblem(dprob, NRM(), maj) +jprob_rdirect = jump_problem(dprob, RDirect(), maj) +jprob_sorting = jump_problem(dprob, SortingDirect(), maj) +jprob_nrm = jump_problem(dprob, NRM(), maj) -SUITE["aggregators"]["RDirect"] = @benchmarkable solve($jprob_rdirect, SSAStepper(); rng) setup=(rng=StableRNG(12345)) evals=1 -SUITE["aggregators"]["SortingDirect"] = @benchmarkable solve( - $jprob_sorting, SSAStepper(); rng -) setup=(rng=StableRNG(12345)) evals=1 -SUITE["aggregators"]["NRM"] = @benchmarkable solve($jprob_nrm, SSAStepper(); rng) setup=(rng=StableRNG(12345)) evals=1 +SUITE["aggregators"]["RDirect"] = benchmark_solve(jprob_rdirect) +SUITE["aggregators"]["SortingDirect"] = benchmark_solve(jprob_sorting) +SUITE["aggregators"]["NRM"] = benchmark_solve(jprob_nrm) diff --git a/test/kernel_contract_tests.jl b/test/kernel_contract_tests.jl index ed0fe35cf..eb437eab6 100644 --- a/test/kernel_contract_tests.jl +++ b/test/kernel_contract_tests.jl @@ -169,14 +169,23 @@ function run_kernel_massaction_rate_tests(backend, alg, aggregator) jp = JumpProblem(DiscreteProblem(u0, (zero(RT), one(RT)), RT[0]), aggregator, maj; save_positions = (false, false)) active = remake(jp; p = RT[0.06]) - inactive_sol = solve(EnsembleProblem(jp), alg, EnsembleGPUKernel(backend); - trajectories = 16, saveat = one(RT)) - active_sol = solve(EnsembleProblem(active), alg, EnsembleGPUKernel(backend); - trajectories = 16, saveat = one(RT)) - @test all(traj -> all(u -> u == u0, traj.u), inactive_sol.u) - @test any(traj -> traj.u[end][2] > 0, active_sol.u) - @test all(traj -> all(u -> u[1] + 3u[2] == 30, traj.u), active_sol.u) - @test all(SciMLBase.successful_retcode, active_sol.u) + kernel_solve(prob) = solve(EnsembleProblem(prob), alg, + EnsembleGPUKernel(backend); trajectories = 16, saveat = one(RT)) + if alg isa SimpleExplicitTauLeaping && RT === Float32 && + !(backend isa KernelAbstractions.CPU) + # PoissonRandom's sampler for larger rates is not type-stable for Float32 + # rates (its `procf` mixes Float32 and Float64 results), so this kernel + # fails to compile for Float32 on GPU backends; the CPU backend tolerates + # the dynamic dispatch. An unexpected pass means the sampler was fixed. + @test_broken (kernel_solve(active); true) + else + inactive_sol = kernel_solve(jp) + active_sol = kernel_solve(active) + @test all(traj -> all(u -> u == u0, traj.u), inactive_sol.u) + @test any(traj -> traj.u[end][2] > 0, active_sol.u) + @test all(traj -> all(u -> u[1] + 3u[2] == 30, traj.u), active_sol.u) + @test all(SciMLBase.successful_retcode, active_sol.u) + end @test jp.prob.p == RT[0] @test jp.massaction_jump.scaled_rates === nothing @test active.massaction_jump.scaled_rates === nothing From f161b3618c965a224775fa96adf59b0788152e75 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 10 Oct 2026 11:40:10 -0400 Subject: [PATCH 42/52] Load specific solver packages in the docs instead of DifferentialEquations The docs depended on the DifferentialEquations, OrdinaryDiffEq, and StochasticDiffEq umbrella packages, but the examples use only Tsit5, Rodas4P, FunctionMap, and SRIW1. Depend on OrdinaryDiffEqTsit5, OrdinaryDiffEqRosenbrock, OrdinaryDiffEqFunctionMap, and StochasticDiffEqHighOrder instead, and load them in the tutorials. JumpProcesses re-exports the problem types and solve, and the solver packages bring in OrdinaryDiffEqCore and StochasticDiffEqCore, which activate JumpProcesses' extensions. The README and docs home page now say to install JumpProcesses plus a solver package. They no longer describe JumpProcesses as part of DifferentialEquations.jl, since DifferentialEquations 8 no longer includes it. The advanced point-process helper now defaults to Rodas4P(), because the solver packages alone provide no default ODE algorithm. The changed pages build locally. The only failing examples are the Catalyst tutorial blocks that need the pending ModelingToolkitBase migration. Co-Authored-By: Claude Opus 5.5 (1M context) --- README.md | 28 +++++++------------ docs/Project.toml | 10 +++---- .../applications/advanced_point_process.md | 3 +- docs/src/assets/Project.toml | 16 +++++++---- docs/src/index.md | 26 ++++++----------- .../tutorials/discrete_stochastic_example.md | 20 +++++++------ .../tutorials/ensembles_and_problem_reuse.md | 2 +- docs/src/tutorials/jump_diffusion.md | 6 ++-- .../src/tutorials/point_process_simulation.md | 4 +-- docs/src/tutorials/simple_poisson_process.md | 10 +++---- 10 files changed, 56 insertions(+), 69 deletions(-) diff --git a/README.md b/README.md index 578131a2d..cb0e09963 100644 --- a/README.md +++ b/README.md @@ -19,9 +19,7 @@ Kinetic Monte Carlo methods across different fields of science. It also enables incorporation of jump processes into hybrid jump-ODE and jump-SDE models, including piecewise deterministic Markov processes (PDMPs) and jump diffusions. -JumpProcesses is a component package in the [SciML](https://sciml.ai/) ecosystem, -and one of the core solver libraries included in -[DifferentialEquations.jl](https://github.com/JuliaDiffEq/DifferentialEquations.jl). +JumpProcesses is a component package in the [SciML](https://sciml.ai/) ecosystem. For information on using the package, [see the stable documentation](https://docs.sciml.ai/JumpProcesses/stable/). Use the @@ -45,26 +43,20 @@ including documentation editing/writing. See also the [contribution section](#co ## Installation -There are two ways to install `JumpProcesses.jl`. First, users may install the meta -`DifferentialEquations.jl` package, which installs and wraps `OrdinaryDiffEq.jl` -for solving ODEs, `StochasticDiffEq.jl` for solving SDEs, and `JumpProcesses.jl`, -along with a number of other useful packages for solving models involving ODEs, -SDEs and/or jump process. This single install will provide the user with all of -the facilities for developing and solving Jump problems. - -To install the `DifferentialEquations.jl` package, refer to the following link -for complete [installation -details](https://docs.sciml.ai/DiffEqDocs/stable/). - -If the user wishes to separately install the `JumpProcesses.jl` library, which is a -lighter dependency than `DifferentialEquations.jl`, then the following code will -install `JumpProcesses.jl` using the Julia package manager: +Install `JumpProcesses.jl` with the Julia package manager: ```julia using Pkg Pkg.add("JumpProcesses") ``` +`JumpProcesses.jl` includes `SSAStepper` for pure jump processes. To couple jumps to ODEs +or SDEs, also install a solver package, for example `OrdinaryDiffEqTsit5` for the `Tsit5` +ODE solver or `StochasticDiffEqHighOrder` for the `SRIW1` SDE solver. See the +[OrdinaryDiffEq.jl](https://docs.sciml.ai/OrdinaryDiffEq/stable/) and +[StochasticDiffEq.jl](https://docs.sciml.ai/DiffEqDocs/stable/solvers/sde_solve/) +documentation for the available solvers and the packages that provide them. + ## Examples ### Stochastic Chemical Kinetics SIR Model @@ -132,7 +124,7 @@ Let's solve an ODE for exponential growth, but coupled to a constant rate jump (Poisson) process that halves the solution each time it fires ```julia -using DifferentialEquations, Plots +using JumpProcesses, OrdinaryDiffEqTsit5, Plots # du/dt = u is the ODE part function f(du, u, p, t) diff --git a/docs/Project.toml b/docs/Project.toml index f41c639bb..09fd1af98 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -2,34 +2,32 @@ Catalyst = "479239e8-5488-4da2-87a7-35f2df7eef83" DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" -DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" -OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" OrdinaryDiffEqFunctionMap = "d3585ca7-f5d3-4ba6-8057-292ed1abd90f" OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce" +OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" PointProcesses = "af0b7596-9bb0-472a-a012-63904f2b4c55" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" -StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" +StochasticDiffEqHighOrder = "0520c28c-50fd-4d16-9c96-902fc80b3bab" [compat] Catalyst = "16" DensityInterface = "0.4" DiffEqCallbacks = "4.17.0" -DifferentialEquations = "7.11, 8.0" Distributions = "0.25" Documenter = "1.1" Graphs = "1.9" JumpProcesses = "9.13.2" SciMLBase = "3" -OrdinaryDiffEq = "6.59, 7" OrdinaryDiffEqFunctionMap = "1, 2" OrdinaryDiffEqRosenbrock = "1, 2" +OrdinaryDiffEqTsit5 = "1, 2" Plots = "1.39" PointProcesses = "0.5" StableRNGs = "1.0" -StochasticDiffEq = "6.63, 7" +StochasticDiffEqHighOrder = "2" diff --git a/docs/src/applications/advanced_point_process.md b/docs/src/applications/advanced_point_process.md index cdc99de53..9b6ee2987 100644 --- a/docs/src/applications/advanced_point_process.md +++ b/docs/src/applications/advanced_point_process.md @@ -362,7 +362,6 @@ not need to evolve time incrementally by small deltas. The `SSAStepper` allow us to step through time one candidate at a time. ```@example tpp-advanced -using OrdinaryDiffEq using OrdinaryDiffEqFunctionMap: FunctionMap using OrdinaryDiffEqRosenbrock: Rodas4P using Random @@ -577,7 +576,7 @@ conditional intensity using an ODEProblem. function integrated_intensity(pp::SciMLPointProcess, t, h; - alg = nothing, + alg = Rodas4P(), saveat = [], save_positions = (true, true)) p = params(pp) diff --git a/docs/src/assets/Project.toml b/docs/src/assets/Project.toml index 546453a95..09fd1af98 100644 --- a/docs/src/assets/Project.toml +++ b/docs/src/assets/Project.toml @@ -1,29 +1,33 @@ [deps] Catalyst = "479239e8-5488-4da2-87a7-35f2df7eef83" DensityInterface = "b429d917-457f-4dbc-8f4c-0cc954292b1d" -DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa" +DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" -OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed" +OrdinaryDiffEqFunctionMap = "d3585ca7-f5d3-4ba6-8057-292ed1abd90f" +OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce" +OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" PointProcesses = "af0b7596-9bb0-472a-a012-63904f2b4c55" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" -StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" +StochasticDiffEqHighOrder = "0520c28c-50fd-4d16-9c96-902fc80b3bab" [compat] Catalyst = "16" DensityInterface = "0.4" -DifferentialEquations = "7.11" +DiffEqCallbacks = "4.17.0" Distributions = "0.25" Documenter = "1.1" Graphs = "1.9" JumpProcesses = "9.13.2" SciMLBase = "3" -OrdinaryDiffEq = "6.59, 7" +OrdinaryDiffEqFunctionMap = "1, 2" +OrdinaryDiffEqRosenbrock = "1, 2" +OrdinaryDiffEqTsit5 = "1, 2" Plots = "1.39" PointProcesses = "0.5" StableRNGs = "1.0" -StochasticDiffEq = "6.63, 7" +StochasticDiffEqHighOrder = "2" diff --git a/docs/src/index.md b/docs/src/index.md index 7d0e03205..e8a28b74c 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -7,9 +7,7 @@ Carlo's methods, thinning method, and Ogata's method. It also enables the incorporation of jump processes into hybrid jump-ODE models, including piecewise deterministic Markov processes, and into hybrid jump-SDE models, including jump diffusions. It is a component package in the -[SciML](https://sciml.ai/) ecosystem, and one of the core solver libraries -included in -[DifferentialEquations.jl](https://docs.sciml.ai/DiffEqDocs/stable/). +[SciML](https://sciml.ai/) ecosystem. Historically, jump processes have been developed in the context of dynamical systems to describe dynamics with discontinuous changes — the jumps — in a system's @@ -119,26 +117,20 @@ In addition to the tutorials linked above, the documentation contains ## Installation -There are two ways to install `JumpProcesses.jl`. First, users may install the meta -`DifferentialEquations.jl` package, which installs and wraps `OrdinaryDiffEq.jl` -for solving ODEs, `StochasticDiffEq.jl` for solving SDEs, and `JumpProcesses.jl`, -along with a number of other useful packages for solving models involving ODEs, -SDEs and/or jump process. This single install will provide the user with all -the facilities for developing and solving Jump problems. - -To install the `DifferentialEquations.jl` package, refer to the following link -for complete [installation -details](https://docs.sciml.ai/DiffEqDocs/stable). - -If the user wishes to install the `JumpProcesses.jl` library separately, which is a -lighter dependency than `DifferentialEquations.jl`, then the following code will -install `JumpProcesses.jl` using the Julia package manager: +Install `JumpProcesses.jl` with the Julia package manager: ```julia using Pkg Pkg.add("JumpProcesses") ``` +`JumpProcesses.jl` includes `SSAStepper` for pure jump processes. To couple jumps to ODEs +or SDEs, also install a solver package, for example `OrdinaryDiffEqTsit5` for the `Tsit5` +ODE solver or `StochasticDiffEqHighOrder` for the `SRIW1` SDE solver. See the +[OrdinaryDiffEq.jl](https://docs.sciml.ai/OrdinaryDiffEq/stable/) and +[StochasticDiffEq.jl](https://docs.sciml.ai/DiffEqDocs/stable/solvers/sde_solve/) +documentation for the available solvers and the packages that provide them. + ## Contributing and Getting Help - Please refer to the diff --git a/docs/src/tutorials/discrete_stochastic_example.md b/docs/src/tutorials/discrete_stochastic_example.md index 5ffe7b105..8fa9a0e3d 100644 --- a/docs/src/tutorials/discrete_stochastic_example.md +++ b/docs/src/tutorials/discrete_stochastic_example.md @@ -43,14 +43,18 @@ reaction, and necessary data structures for the simulation algorithms, such as dependency graphs, are automatically calculated. We'll make use of [JumpProcesses.jl](https://docs.sciml.ai/JumpProcesses/stable/) -together with [OrdinaryDiffEq.jl](https://docs.sciml.ai/OrdinaryDiffEq/stable/) -(for ODE solvers), Plots.jl, and (optionally) Catalyst.jl in this tutorial. If -not already installed, they can be added as follows: +together with OrdinaryDiffEqTsit5 (for the `Tsit5` ODE solver from +[OrdinaryDiffEq.jl](https://docs.sciml.ai/OrdinaryDiffEq/stable/)), +StochasticDiffEqHighOrder (for the `SRIW1` SDE solver from +[StochasticDiffEq.jl](https://docs.sciml.ai/DiffEqDocs/stable/solvers/sde_solve/)), +Plots.jl, and (optionally) Catalyst.jl in this tutorial. If not already +installed, they can be added as follows: ```julia using Pkg Pkg.add("JumpProcesses") -Pkg.add("OrdinaryDiffEq") +Pkg.add("OrdinaryDiffEqTsit5") +Pkg.add("StochasticDiffEqHighOrder") Pkg.add("Plots") Pkg.add("Catalyst") # optional ``` @@ -58,12 +62,12 @@ Pkg.add("Catalyst") # optional Let's now load the required packages and set some default plot settings ```julia -using JumpProcesses, OrdinaryDiffEq, Plots, LinearAlgebra +using JumpProcesses, OrdinaryDiffEqTsit5, Plots, LinearAlgebra default(; lw = 2) ``` ```@setup tut2 -using JumpProcesses, OrdinaryDiffEq, Plots, LinearAlgebra +using JumpProcesses, OrdinaryDiffEqTsit5, Plots, LinearAlgebra default(; lw = 2) ``` @@ -820,7 +824,7 @@ differential equation. Let's define an ODE problem, where the continuous part only acts on some new 4th component: ```@example tut2 -using OrdinaryDiffEq +using OrdinaryDiffEqTsit5 function f(du, u, p, t) du[4] = u[2] * u[3] / 1e5 - u[1] * u[4] / 1e5 nothing @@ -954,7 +958,7 @@ problem except with multiplicative noise on `u[4]` by using an `SDEProblem` instead: ```@example tut2 -using StochasticDiffEq +using StochasticDiffEqHighOrder function g(du, u, p, t) du[4] = 0.1u[4] end diff --git a/docs/src/tutorials/ensembles_and_problem_reuse.md b/docs/src/tutorials/ensembles_and_problem_reuse.md index 5c0ff634a..1f5b37285 100644 --- a/docs/src/tutorials/ensembles_and_problem_reuse.md +++ b/docs/src/tutorials/ensembles_and_problem_reuse.md @@ -227,7 +227,7 @@ the `alias_jumps` field of their alias specifier. When it is unspecified they re jump state on thread 1 and copy it on other threads. To request a copy explicitly: ```julia -using StochasticDiffEq +using StochasticDiffEqHighOrder sol = solve(sde_jprob, SRIW1(); alias = SciMLBase.SDEAliasSpecifier(alias_jumps = false)) ``` diff --git a/docs/src/tutorials/jump_diffusion.md b/docs/src/tutorials/jump_diffusion.md index 1caacba4d..6a10f5d66 100644 --- a/docs/src/tutorials/jump_diffusion.md +++ b/docs/src/tutorials/jump_diffusion.md @@ -31,15 +31,15 @@ not already installed ```julia using Pkg Pkg.add("JumpProcesses") -Pkg.add("OrdinaryDiffEq") -Pkg.add("StochasticDiffEq") +Pkg.add("OrdinaryDiffEqTsit5") +Pkg.add("StochasticDiffEqHighOrder") Pkg.add("Plots") ``` We then load these packages, and set some plotting defaults, as ```@example tut3 -using JumpProcesses, StochasticDiffEq, OrdinaryDiffEq, Plots +using JumpProcesses, StochasticDiffEqHighOrder, OrdinaryDiffEqTsit5, Plots default(; lw = 2) ``` diff --git a/docs/src/tutorials/point_process_simulation.md b/docs/src/tutorials/point_process_simulation.md index c3d65dad7..968f4fd60 100644 --- a/docs/src/tutorials/point_process_simulation.md +++ b/docs/src/tutorials/point_process_simulation.md @@ -140,8 +140,8 @@ By breaking the problem formulation and solver selection into specifying a and generating a realization via `solve`, JumpProcesses' has the flexibility to specify and simulate a broad variety of problem types. The base problem allow us to combine TPPs with other types of dynamics such as ODEs or SDEs, by replacing -`DiscreteProblem` with `ODEProblem` or `SDEProblem` from -[DifferentialEquations.jl](https://docs.sciml.ai/DiffEqDocs/stable/). For +`DiscreteProblem` with `ODEProblem` or `SDEProblem` (see the +[SciML differential equations documentation](https://docs.sciml.ai/DiffEqDocs/stable/)). For instance, we can declare a conditional intensity function that follows an ODE using `ODEProblem`. The `JumpProblem` allows us to combine multiple TPPs together as we will see in the next section. The simulation algorithm (i.e. diff --git a/docs/src/tutorials/simple_poisson_process.md b/docs/src/tutorials/simple_poisson_process.md index 20cd2e97e..e5f27caf3 100644 --- a/docs/src/tutorials/simple_poisson_process.md +++ b/docs/src/tutorials/simple_poisson_process.md @@ -318,20 +318,18 @@ are essentially defining a combined ODE-jump process, i.e., a [piecewise deterministic Markov process](https://en.wikipedia.org/wiki/Piecewise-deterministic_Markov_process), but one where the ODE is trivial and does not change the state. To use this -problem type and the ODE solvers, we first load `OrdinaryDiffEq.jl` or -`DifferentialEquations.jl`. If neither is installed, we first +problem type we need an ODE solver; here we use `Tsit5` from the +`OrdinaryDiffEqTsit5` package. If it is not installed, we first ```julia using Pkg -Pkg.add("OrdinaryDiffEq") -# or Pkg.add("DifferentialEquations") +Pkg.add("OrdinaryDiffEqTsit5") ``` and then load it via ```@example tut1 -using OrdinaryDiffEq -# or using DifferentialEquations +using OrdinaryDiffEqTsit5 ``` We can then construct our ODE problem with a trivial ODE derivative component. From a870daa6406e9e8f6d8a761e63a873ec0d7aa8e3 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 10 Oct 2026 13:04:27 -0400 Subject: [PATCH 43/52] Avoid a per-merge allocation in mass-action jump merging massaction_jump_combine checked both jumps for custom parameter mappers by looping over a tuple of two differently typed jumps, which allocated 128 bytes on every merge. It now checks each jump explicitly. Merging 10,000 scalar MassActionJumps allocates 0.73 MB again, as in released 9.33.1, rather than 2.0 MB. Merging a jump with fixed rates and one with a parameter mapping, in either order, now raises an informative error instead of a MethodError. New tests merge scalar-form MassActionJumps in any order with scalar or collection-form jumps, through JumpSet and JumpProblem. They check the merged stoichiometry, parameter indices, and rates, that the problem solves, and that the user's jumps are unchanged. They also cover the mixed-merge error. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/jumps.jl | 18 +++++++--- test/massaction_state.jl | 72 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 85 insertions(+), 5 deletions(-) diff --git a/src/jumps.jl b/src/jumps.jl index 060745bbf..a1f7e645d 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -1061,12 +1061,20 @@ function massaction_jump_combine(maj1::Nothing, maj2::MassActionJump) setup_majump_to_merge(maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, maj2.param_mapper, maj2.rescale_rates_on_update) end +function uses_custom_mapper(maj) + using_params(maj) && !(maj.param_mapper isa MassActionJumpParamMapper) +end + function massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) - for m in (maj1, maj2) - (using_params(m) && !(m.param_mapper isa MassActionJumpParamMapper)) && - error("Cannot merge MassActionJumps backed by custom parameter mappers. " * - "Construct a single MassActionJump with all reactions instead.") - end + # Check the two jumps separately: looping over a tuple of differently typed jumps + # allocates on every merge. + (uses_custom_mapper(maj1) || uses_custom_mapper(maj2)) && + error("Cannot merge MassActionJumps backed by custom parameter mappers. " * + "Construct a single MassActionJump with all reactions instead.") + (using_params(maj1) == using_params(maj2)) || + error("Cannot merge a MassActionJump that has fixed rates with one that has a " * + "parameter mapping (`param_idxs` or `param_mapper`). Give every jump " * + "fixed rates, or every jump a parameter mapping.") (maj1.rescale_rates_on_update == maj2.rescale_rates_on_update) || error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") majump_merge!(maj1, maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, diff --git a/test/massaction_state.jl b/test/massaction_state.jl index 7333367ff..61cb3da8f 100644 --- a/test/massaction_state.jl +++ b/test/massaction_state.jl @@ -106,6 +106,78 @@ end end end +# Scalar-form jumps hold one reaction without enclosing vectors. +@testset "Scalar mass-action jumps merge in any order" begin + p = [2.0, 3.0, 5.0] + fixed_scalar1() = MassActionJump(2.0, [1 => 1], [1 => -1, 2 => 1]) + fixed_scalar2() = MassActionJump(3.0, [2 => 2], [2 => -2, 3 => 1]) # scaled by 1/2! + fixed_vector() = MassActionJump([5.0], [[3 => 1]], [[3 => -1]]) + param_scalar1() = MassActionJump([1 => 1], [1 => -1, 2 => 1]; param_idxs = 1) + param_scalar2() = MassActionJump([2 => 2], [2 => -2, 3 => 1]; param_idxs = 2) + param_vector() = MassActionJump([[3 => 1]], [[3 => -1]]; param_idxs = [3]) + cases = (("fixed scalar, fixed scalar", (fixed_scalar1, fixed_scalar2), [2.0, 1.5]), + ("fixed scalar, fixed vector", (fixed_scalar1, fixed_vector), [2.0, 5.0]), + ("fixed vector, fixed scalar", (fixed_vector, fixed_scalar1), [5.0, 2.0]), + ("parameterized scalar, parameterized scalar", (param_scalar1, param_scalar2), + [2.0, 1.5]), + ("parameterized scalar, parameterized vector", (param_scalar1, param_vector), + [2.0, 5.0]), + ("parameterized vector, parameterized scalar", (param_vector, param_scalar1), + [5.0, 2.0]), + ("single fixed scalar", (fixed_scalar1,), [2.0]), + ("single parameterized scalar", (param_scalar1,), [2.0])) + dprob = DiscreteProblem([10, 10, 10], (0.0, 1.0), p) + + # A jump's reactions as a vector, whether the jump is in scalar or collection form. + reactions(stoich) = eltype(stoich) <: Pair ? [stoich] : stoich + + @testset "$name" for (name, builders, rates) in cases + jumps = map(build -> build(), builders) + fresh_jumps = map(build -> build(), builders) + # The merged jump lists each jump's reactions and parameter indices in order. + expected_rs = reduce(vcat, map(maj -> reactions(maj.reactant_stoch), fresh_jumps)) + expected_ns = reduce(vcat, map(maj -> reactions(maj.net_stoch), fresh_jumps)) + parameterized = first(fresh_jumps).param_mapper !== nothing + function check_merged(maj) + @test maj.reactant_stoch == expected_rs + @test maj.net_stoch == expected_ns + if parameterized + @test maj.param_mapper.param_idxs == + reduce(vcat, map(m -> m.param_mapper.param_idxs, fresh_jumps)) + end + dest = zeros(JumpProcesses.get_num_majumps(maj)) + JumpProcesses.fill_scaled_rates!(dest, maj, p) + @test dest == rates + end + for jset in (JumpSet(jumps...), JumpSet(; massaction_jumps = collect(jumps))) + check_merged(jset.massaction_jump) + end + jprob = JumpProblem(dprob, Direct(), jumps...) + check_merged(jprob.massaction_jump) + @test successful_retcode(solve(jprob, SSAStepper(); seed = 1)) + + # Merging leaves the user's jumps unchanged. + for (maj, fresh) in zip(jumps, fresh_jumps) + @test maj.reactant_stoch == fresh.reactant_stoch + @test maj.net_stoch == fresh.net_stoch + @test maj.scaled_rates == fresh.scaled_rates + if fresh.param_mapper !== nothing + @test maj.param_mapper.param_idxs == fresh.param_mapper.param_idxs + end + end + end + + @testset "Mixing fixed-rate and parameterized jumps is diagnosed" begin + for builders in ((fixed_scalar1, param_scalar1), (param_scalar1, fixed_scalar1), + (fixed_vector, param_vector), (param_vector, fixed_vector)) + jumps = map(build -> build(), builders) + @test_throws r"fixed rates with one that has a parameter mapping" JumpSet(jumps...) + @test_throws r"fixed rates with one that has a parameter mapping" JumpSet(; + massaction_jumps = collect(jumps)) + end + end +end + @testset "Removed reset rate-refresh keyword is diagnosed" begin maj = MassActionJump([[1 => 1]], [[1 => -1]]; param_idxs = 1) jprob = JumpProblem(DiscreteProblem([5], (0.0, 1.0), [1.0]), Direct(), maj) From 7eac8f717159cdf55f9a9b6ecaad931ef2021c11 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 10 Oct 2026 15:44:40 -0400 Subject: [PATCH 44/52] Remove the StochasticDiffEqCore extension and test against its v10 branch OrdinaryDiffEq.jl#4707 gave StochasticDiffEqCore per-algorithm __init(::JumpProblem, ...) methods, and OrdinaryDiffEq.jl#4836 narrows Core's remaining union __init to AbstractRODEProblem. With both, stochastic init of a JumpProblem dispatches to Core without ambiguity, so JumpProcessesStochasticDiffEqCoreExt and the StochasticDiffEqCore weak dependency are removed. The OrdinaryDiffEqCore extension keeps its four-family method, which covers the intersection with OrdinaryDiffEqCore's initializer. Until #4836 is released, tests pin StochasticDiffEqCore to that branch through [sources]. Julia 1.10 ignores [sources], and the downgrade action discards URL sources; registered StochasticDiffEqCore's init is ambiguous without the extension. So InterfaceI and InterfaceII drop lts, and the downgrade job is off, until the pin is removed. QA, ThreadSafety, and benchmarks keep lts: they do not call stochastic init. stochastic_dispatch.jl now checks that SDE, RODE, and TauLeaping init dispatch to Core's methods, and runs the filtered ambiguity check over JumpProcesses, Core, and the OrdinaryDiffEqCore extension. HISTORY and the solver-dispatch docs are updated. Validation, Julia 1.12.7 with the pinned branch: stochastic_dispatch.jl (287), rng_kwarg_tests.jl (36), callbacks.jl (162), variable_rate.jl (76), and QA (20) pass. Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/Downgrade.yml | 4 ++ HISTORY.md | 8 ++-- Project.toml | 11 ++++-- docs/src/api.md | 10 ++--- ext/JumpProcessesOrdinaryDiffEqCoreExt.jl | 4 +- ext/JumpProcessesStochasticDiffEqCoreExt.jl | 25 ------------ test/stochastic_dispatch.jl | 43 +++++---------------- test/test_groups.toml | 10 ++++- 8 files changed, 38 insertions(+), 77 deletions(-) delete mode 100644 ext/JumpProcessesStochasticDiffEqCoreExt.jl diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml index 9ae531357..d5b91b32f 100644 --- a/.github/workflows/Downgrade.yml +++ b/.github/workflows/Downgrade.yml @@ -13,6 +13,10 @@ on: jobs: downgrade: name: "Downgrade" + # The downgrade action discards URL sources. Restore this job after the + # StochasticDiffEqCore branch for JumpProcesses 10 is released and the + # temporary pin is removed. + if: ${{ false }} uses: "SciML/.github/.github/workflows/downgrade.yml@v1" with: group: All diff --git a/HISTORY.md b/HISTORY.md index 06049ba17..8e11a629c 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -108,11 +108,9 @@ changes are not intended for a 9.x release. twice, so they ran twice per step, and were reintroduced even with `merge_callbacks = false`. `init` now combines stored and call-level callbacks exactly as `solve` does. - - StochasticDiffEqCore is now a weak dependency: when it is installed, version - 2.2.1 or later is required. A new extension sends stochastic `init` of a - `JumpProblem` to StochasticDiffEqCore's own initializer on every supported version, - so `init` followed by `solve!` matches `solve`. StochasticDiffEqCore versions with - their own per-algorithm `JumpProblem` initializers take precedence automatically. + - Stochastic `init` of a `JumpProblem` now reaches StochasticDiffEqCore's own + per-algorithm `JumpProblem` initializers, so `init` followed by `solve!` matches + `solve`. This requires a StochasticDiffEqCore release that supports JumpProcesses 10. - The state types `SSAStepper` supports are now documented: vectors and `SVector`s of integers or floats, species × sites matrices for spatial problems, and scalars for models without mass-action jumps. See "Supported diff --git a/Project.toml b/Project.toml index 9da983bb8..f0a72ce7e 100644 --- a/Project.toml +++ b/Project.toml @@ -27,13 +27,11 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8" -StochasticDiffEqCore = "19c5a474-6cd1-4a5f-be79-46dc34e54d7f" [extensions] JumpProcessesForwardDiffExt = "ForwardDiff" JumpProcessesKernelAbstractionsExt = ["Adapt", "KernelAbstractions"] JumpProcessesOrdinaryDiffEqCoreExt = "OrdinaryDiffEqCore" -JumpProcessesStochasticDiffEqCoreExt = ["OrdinaryDiffEqCore", "StochasticDiffEqCore"] [compat] ADTypes = "1.22.0" @@ -66,11 +64,18 @@ StableRNGs = "1" StaticArrays = "1.9.18" Statistics = "1" StochasticDiffEq = "7.0.0, 7" -StochasticDiffEqCore = "2.2.1" +StochasticDiffEqCore = "2.2.3" SymbolicIndexingInterface = "0.3.48" Test = "1" julia = "1.10" +[sources] +# Temporary: tests use the StochasticDiffEqCore branch for JumpProcesses 10 +# (OrdinaryDiffEq.jl#4836), which removes the need for a JumpProcesses StochasticDiffEqCore +# extension. Remove when it is released. Julia 1.10 ignores [sources]; see +# test/test_groups.toml. +StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "7a0294b8a8368f7b1fb17bf5ec4e56d80db4e8cf", subdir = "lib/StochasticDiffEqCore"} + [extras] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" diff --git a/docs/src/api.md b/docs/src/api.md index 026babd1e..454de5b4d 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -311,13 +311,9 @@ For **StochasticDiffEq SDE/RODE solvers**, both `solve` and `init` reach the backend's JumpProblem initializer, StochasticDiffEqCore's `_sde_init`, with the original JumpProblem and unchanged RNG/seed inputs, so they share its RNG policy, jump-state copying, and callback setup. `solve` goes through the backend's more -specific `__solve`. For `init`, StochasticDiffEqCore versions with per-algorithm -`__init(::JumpProblem, ...)` methods own dispatch directly; on other versions, -JumpProcesses' `JumpProcessesStochasticDiffEqCoreExt` extension adds one -`__init(::JumpProblem, ::Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm})` -method that forwards to `_sde_init`, resolving an otherwise ambiguous dispatch. The -backend's per-algorithm methods are more specific and take precedence wherever they -exist, as does its method for jump algorithms such as `TauLeaping`. +specific `__solve`. For `init`, the backend's per-algorithm `__init(::JumpProblem, ...)` +methods, and its method for jump algorithms such as `TauLeaping`, are more specific than +JumpProcesses' OrdinaryDiffEqCore extension method and own dispatch directly. DiffEqBase's `init_call` merges the JumpProblem's stored keywords and user callbacks once, before dispatching to `__init`; JumpProcesses' `__init` methods do not merge diff --git a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl index edb0029bb..ceeceb0c9 100644 --- a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl +++ b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl @@ -17,8 +17,8 @@ function SciMLBase.__init( StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; kwargs...) where {IIP, P} # Cover the full intersection with OrdinaryDiffEqCore's initializer. SDE/RODE - # initialization belongs to StochasticDiffEqCore: its per-algorithm methods where it - # has them, otherwise the forwarding method in JumpProcessesStochasticDiffEqCoreExt. + # initialization belongs to StochasticDiffEqCore, whose per-algorithm methods are + # more specific. _jump_init(_jump_prob, alg; kwargs...) end diff --git a/ext/JumpProcessesStochasticDiffEqCoreExt.jl b/ext/JumpProcessesStochasticDiffEqCoreExt.jl deleted file mode 100644 index 599f78185..000000000 --- a/ext/JumpProcessesStochasticDiffEqCoreExt.jl +++ /dev/null @@ -1,25 +0,0 @@ -module JumpProcessesStochasticDiffEqCoreExt - -using JumpProcesses: JumpProblem -import SciMLBase -import OrdinaryDiffEqCore: StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm -import StochasticDiffEqCore - -# Forward stochastic `init` of a `JumpProblem` to StochasticDiffEqCore's own initializer, -# `_sde_init`, which is exported for packages that need to bypass `__init` dispatch. On -# Core versions whose `__init` covers `Union{AbstractRODEProblem, JumpProblem}`, that -# method and the `OrdinaryDiffEqCore` extension's four-family method are each more -# specific in one argument, so without this method `init(jprob, sde_alg)` is ambiguous. -# Core versions with per-algorithm `JumpProblem` methods (OrdinaryDiffEq.jl#4614) are more -# specific than this union method and take precedence, and Core's jump-algorithm method -# (e.g. `TauLeaping`) is also more specific. Keep the signature distinct from Core's: do -# not split it into separate SDE and RODE methods. -# -# `init_call` has already merged the `JumpProblem`'s stored keywords, and `_sde_init` -# documents that they must not be merged again. -function SciMLBase.__init(prob::JumpProblem, - alg::Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; kwargs...) - StochasticDiffEqCore._sde_init(prob, alg; kwargs...) -end - -end diff --git a/test/stochastic_dispatch.jl b/test/stochastic_dispatch.jl index 439de4600..e7bc1d5dd 100644 --- a/test/stochastic_dispatch.jl +++ b/test/stochastic_dispatch.jl @@ -24,55 +24,32 @@ function stochastic_rng_options(kind) end const SDECore = StochasticDiffEq.StochasticDiffEqCore -const SDE_BRIDGE = Base.get_extension(JumpProcesses, :JumpProcessesStochasticDiffEqCoreExt) -const BRIDGE_SIGNATURE = Tuple{typeof(SciMLBase.__init), JumpProblem, - Union{SDECore.StochasticDiffEqAlgorithm, SDECore.StochasticDiffEqRODEAlgorithm}} +const ODECORE_EXT = Base.get_extension(JumpProcesses, :JumpProcessesOrdinaryDiffEqCoreExt) -# Whether StochasticDiffEqCore defines its own `__init(::JumpProblem, ::A)` for the -# algorithm supertype `A`, as Core versions with OrdinaryDiffEq.jl#4614 do. Decided from -# the method table rather than Core's version number. -function core_owns_jump_init(A) - signature = Tuple{typeof(SciMLBase.__init), JumpProblem, A} - any(m -> m.module === SDECore && m.sig == signature, methods(SciMLBase.__init)) -end - -# Whether a method is the forwarding method, or a `__init`, `__solve`, or -# `supports_solve_rng` method for `JumpProblem`s. +# Whether a method is a `__init`, `__solve`, or `supports_solve_rng` method for +# `JumpProblem`s. function involves_jump_init(m) - m.module === SDE_BRIDGE && return true m.name in (:__init, :__solve, :supports_solve_rng) || return false any(T -> T isa Type && T <: JumpProblem, Base.unwrap_unionall(m.sig).parameters) end @testset "Stochastic initialization dispatches to the backend" begin - @test SDE_BRIDGE !== nothing + # Core's per-algorithm methods are more specific than JumpProcesses' extension method. for (kind, alg, A) in ((:SDE, EM(), SDECore.StochasticDiffEqAlgorithm), (:RODE, RandomEM(), SDECore.StochasticDiffEqRODEAlgorithm)) jprob = stochastic_jump_problem(kind) method = which(SciMLBase.__init, (typeof(jprob), typeof(alg))) - if core_owns_jump_init(A) - # Core's per-algorithm method is more specific than the forwarding method. - @test method.module === SDECore - @test method.sig == Tuple{typeof(SciMLBase.__init), JumpProblem, A} - else - # Otherwise JumpProcesses' extension forwards to Core's `_sde_init`. - @test method.module === SDE_BRIDGE - @test method.sig == BRIDGE_SIGNATURE - end + @test method.module === SDECore + @test method.sig == Tuple{typeof(SciMLBase.__init), JumpProblem, A} end - # The forwarding method exists once, owned by the extension: Core never overwrites it. - bridge_methods = filter(m -> m.sig == BRIDGE_SIGNATURE, collect(methods(SciMLBase.__init))) - @test length(bridge_methods) == 1 - @test only(bridge_methods).module === SDE_BRIDGE - - # Core's jump-algorithm method (e.g. `TauLeaping`) is more specific than the bridge. + # So is Core's jump-algorithm method (e.g. `TauLeaping`). jprob = stochastic_jump_problem(:SDE) @test which(SciMLBase.__init, (typeof(jprob), typeof(TauLeaping()))).module === SDECore - # No ambiguity involves the forwarding method or a `JumpProblem` method of `__init`, - # `__solve`, or `supports_solve_rng`. - ambiguities = Test.detect_ambiguities(JumpProcesses, SDECore, SDE_BRIDGE) + # No ambiguity involves a `JumpProblem` method of `__init`, `__solve`, or + # `supports_solve_rng`. + ambiguities = Test.detect_ambiguities(JumpProcesses, SDECore, ODECORE_EXT) @test !any(pair -> any(involves_jump_init, pair), ambiguities) end diff --git a/test/test_groups.toml b/test/test_groups.toml index bfc13e0c8..c953a3782 100644 --- a/test/test_groups.toml +++ b/test/test_groups.toml @@ -1,5 +1,8 @@ [InterfaceI] -versions = ["lts", "1", "pre"] +# Restore "lts" when the temporary StochasticDiffEqCore pin is removed: Julia 1.10 +# ignores [sources], and registered StochasticDiffEqCore lacks the stochastic `init` +# dispatch these tests check. +versions = ["1", "pre"] ["InterfaceI 32-bit"] group = "InterfaceI" @@ -8,7 +11,10 @@ os = ["ubuntu-latest"] arch = "x86" [InterfaceII] -versions = ["lts", "1", "pre"] +# Restore "lts" when the temporary StochasticDiffEqCore pin is removed: Julia 1.10 +# ignores [sources], and registered StochasticDiffEqCore lacks the stochastic `init` +# dispatch these tests check. +versions = ["1", "pre"] [QA] versions = ["lts", "1"] From cb6aa2fca3e144c98a4105d1eb32fc5459a72d83 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 10 Oct 2026 18:40:24 -0400 Subject: [PATCH 45/52] Fix SSAStepper erroring when saving only the final state solve! checked integrator.sol.t[end] before saving the end state, which threw a BoundsError when nothing had been saved yet: save_start = false, save_positions = (false, false), no saveat, and save_end = true. The check now treats an empty solution as needing the end save, so the solution holds just the final time and state. The same line is on master. saveat_regression.jl gains a "Final-only saving" testset: no jumps, jumps (the final state equals the number of jumps that fired), and an end time already saved by a callback in the final pass, which must not be saved twice. Co-Authored-By: Claude Opus 5.5 (1M context) --- HISTORY.md | 3 +++ src/SSA_stepper.jl | 3 ++- test/saveat_regression.jl | 42 +++++++++++++++++++++++++++++++++++++++ 3 files changed, 47 insertions(+), 1 deletion(-) diff --git a/HISTORY.md b/HISTORY.md index 8e11a629c..9e44c846a 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -103,6 +103,9 @@ changes are not intended for a 9.x release. final callback pass. A callback that changed the state at the end time altered the values saved at those earlier times, and the saved times came out of order. They are now saved before the final callback pass. + - Fixed `SSAStepper` erroring with a `BoundsError` when saving only the final state: + `save_end = true` with nothing saved earlier, for example with `save_start = false`, + `save_positions = (false, false)`, and no `saveat`. The final state is now saved. - Fixed `init` on ODE, DAE, and `FunctionMap` integrators for `JumpProblem`s: callbacks stored in the `JumpProblem` (`JumpProblem(...; callback)`) were merged twice, so they ran twice per step, and were reintroduced even with diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index 550d896e7..8af965090 100644 --- a/src/SSA_stepper.jl +++ b/src/SSA_stepper.jl @@ -421,7 +421,8 @@ function DiffEqBase.solve!(integrator::SSAIntegrator) integrator.opts.callback.discrete_callbacks...) end - if integrator.save_end && integrator.sol.t[end] != end_time + if integrator.save_end && + (isempty(integrator.sol.t) || integrator.sol.t[end] != end_time) save_current_state!(integrator, end_time) end end diff --git a/test/saveat_regression.jl b/test/saveat_regression.jl index 483d08111..a31674606 100644 --- a/test/saveat_regression.jl +++ b/test/saveat_regression.jl @@ -187,3 +187,45 @@ end cb_saves[2] && @test sol.u[end] == [10] end end + +# Saving only the final state must work when nothing was saved earlier: no start save, no +# event saves, and no `saveat`. +@testset "Final-only saving" begin + function final_only(rate_value; callback = nothing) + fired = Ref(0) + jump = ConstantRateJump((u, p, t) -> rate_value, + integrator -> (integrator.u[1] += 1; fired[] += 1; nothing)) + jprob = JumpProblem(DiscreteProblem([0], (0.0, 1.0)), Direct(), jump; + save_positions = (false, false)) + sol = solve(jprob, SSAStepper(); save_start = false, save_end = true, callback, + seed = 1) + sol, fired[] + end + + @testset "no jumps" begin + sol, fired = final_only(0.0) + @test SciMLBase.successful_retcode(sol) + @test fired == 0 + @test sol.t == [1.0] + @test sol.u == [[0]] + end + + @testset "with jumps" begin + sol, fired = final_only(20.0) + @test SciMLBase.successful_retcode(sol) + @test fired > 0 + @test sol.t == [1.0] + @test sol.u == [[fired]] + end + + # An end time already saved, here by a callback in the final callback pass, is not + # saved again. + @testset "end time already saved" begin + saves_at_end = DiscreteCallback((u, t, integrator) -> t == 1.0, + integrator -> nothing; save_positions = (false, true)) + sol, fired = final_only(20.0; callback = saves_at_end) + @test SciMLBase.successful_retcode(sol) + @test sol.t == [1.0] + @test sol.u == [[fired]] + end +end From 66d4f298fd98b1afe51d1c6cf0ba2c56a58e5eae Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sat, 10 Oct 2026 18:40:44 -0400 Subject: [PATCH 46/52] Track the StochasticDiffEqCore v10 branch instead of a fixed commit The temporary [sources] pin now follows the sdecore/jumpprocesses-v10 branch of OrdinaryDiffEq.jl#4836, so this PR's CI tests its latest state without a pin bump for every change there. Exact commits will be pinned for the 10.0 release validation. Co-Authored-By: Claude Opus 5.5 (1M context) --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index f0a72ce7e..4ebce5804 100644 --- a/Project.toml +++ b/Project.toml @@ -74,7 +74,7 @@ julia = "1.10" # (OrdinaryDiffEq.jl#4836), which removes the need for a JumpProcesses StochasticDiffEqCore # extension. Remove when it is released. Julia 1.10 ignores [sources]; see # test/test_groups.toml. -StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "7a0294b8a8368f7b1fb17bf5ec4e56d80db4e8cf", subdir = "lib/StochasticDiffEqCore"} +StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "sdecore/jumpprocesses-v10", subdir = "lib/StochasticDiffEqCore"} [extras] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" From ab206563c63d4de143d0cd3a766c1124929c0eb7 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sun, 11 Oct 2026 07:34:01 -0400 Subject: [PATCH 47/52] Show unseeded and single-RNG serial solves in the reuse tutorial The ensembles and problem reuse tutorial seeded every example, which suggested that repeated serial solves need a seed. The serial-solves section now shows two examples: solves drawing from Julia's default RNG, and one generator passed to every solve for reproducible samples, with a warning against creating a new same-seed generator per solve, which makes every solve follow the same path. Seeds that only made the rendered output reproducible are now hidden Random.seed! lines, or dropped where the output is deterministic. The SortingDirect replay section keeps its seeds, since reproducibility is its subject. docs/make.jl ignores the stable-docs URL of this tutorial in linkcheck: the JumpProblem docstring links it, and the page exists only once the 10.0 docs are deployed. Remove the entry then. Co-Authored-By: Claude Opus 5.5 (1M context) --- docs/make.jl | 2 + .../tutorials/ensembles_and_problem_reuse.md | 37 +++++++++++++++---- 2 files changed, 31 insertions(+), 8 deletions(-) diff --git a/docs/make.jl b/docs/make.jl index 92675e33e..ba37a4a6d 100644 --- a/docs/make.jl +++ b/docs/make.jl @@ -24,6 +24,8 @@ mathengine = MathJax3(Dict(:loader => Dict("load" => ["[tex]/require", "[tex]/ma makedocs( sitename = "JumpProcesses.jl", authors = "Chris Rackauckas", modules = [JumpProcesses], clean = true, linkcheck = true, checkdocs = :exports, + # The stable docs have this tutorial only once the 10.0 docs are deployed; remove then. + linkcheck_ignore = [r"^https://docs\.sciml\.ai/JumpProcesses/stable/tutorials/ensembles_and_problem_reuse/?$"], format = Documenter.HTML(; assets = ["assets/favicon.ico"], canonical = "https://docs.sciml.ai/JumpProcesses/", prettyurls = (get(ENV, "CI", nothing) == "true"), diff --git a/docs/src/tutorials/ensembles_and_problem_reuse.md b/docs/src/tutorials/ensembles_and_problem_reuse.md index 1f5b37285..72e75a2f0 100644 --- a/docs/src/tutorials/ensembles_and_problem_reuse.md +++ b/docs/src/tutorials/ensembles_and_problem_reuse.md @@ -34,12 +34,30 @@ nothing # hide Solving the same problem repeatedly in a loop is the fastest way to generate many samples: each solve re-initializes the problem's jump state from the problem's current initial condition and parameters, so nothing needs to be copied or rebuilt. +No seed is needed either. Each solve draws from Julia's default random number +generator, whose state advances from one solve to the next: ```@example reuse -final_counts = [solve(jprob, SSAStepper(); seed).u[end][1] for seed in 1:1000] +using Random # hide +Random.seed!(1) # hide +final_counts = [solve(jprob, SSAStepper()).u[end][1] for _ in 1:1000] sum(final_counts) / length(final_counts) ``` +To make the samples reproducible, create one generator and pass it to every solve. +Each solve continues the generator's stream, so the samples still differ: + +```@example reuse +using Random +rng = Xoshiro(1234) +final_counts = [solve(jprob, SSAStepper(); rng).u[end][1] for _ in 1:1000] +sum(final_counts) / length(final_counts) +``` + +Do not create a new generator with the same seed for each solve, for example +`solve(jprob, SSAStepper(); rng = Xoshiro(1234))` inside the loop: every solve would +then follow the same path. + ## Parameter sweeps with `remake` `remake` creates a problem with a new initial condition, parameters, or time span. @@ -55,9 +73,10 @@ Since every solve re-initializes the shared state from the problem being solved, problems can be solved one after another: ```@example reuse +Random.seed!(2) # hide mean_counts = map([0.5, 1.0, 2.0]) do birth_rate prob = remake(jprob; p = [birth_rate, 0.1]) - sum(solve(prob, SSAStepper(); seed).u[end][1] for seed in 1:500) / 500 + sum(solve(prob, SSAStepper()).u[end][1] for _ in 1:500) / 500 end ``` @@ -81,7 +100,7 @@ handles the copying needed to run trajectories safely: ```@example reuse ensemble = EnsembleProblem(jprob) -esol = solve(ensemble, SSAStepper(), EnsembleThreads(); trajectories = 1000, seed = 1) +esol = solve(ensemble, SSAStepper(), EnsembleThreads(); trajectories = 1000) length(esol.u) # the trajectories ``` @@ -154,10 +173,10 @@ per task and reused for all of that task's solves: ```@example reuse nsims = 1000 chunks = Iterators.partition(1:nsims, cld(nsims, Threads.nthreads())) -tasks = map(chunks) do seeds +tasks = map(chunks) do chunk Threads.@spawn begin local_prob = deepcopy(jprob) # one independent copy per task - [solve(local_prob, SSAStepper(); seed).u[end][1] for seed in seeds] + [solve(local_prob, SSAStepper()).u[end][1] for _ in chunk] end end final_counts = reduce(vcat, fetch.(tasks)) @@ -169,8 +188,9 @@ at the same time. For the same reason, integrators that are alive at the same ti independent problems: ```@example reuse -integrator1 = init(deepcopy(jprob), SSAStepper(); seed = 1) -integrator2 = init(deepcopy(jprob), SSAStepper(); seed = 2) +Random.seed!(3) # hide +integrator1 = init(deepcopy(jprob), SSAStepper()) +integrator2 = init(deepcopy(jprob), SSAStepper()) step!(integrator1) step!(integrator2) integrator1.t, integrator2.t @@ -212,8 +232,9 @@ Permitting `u0` to be aliased avoids a copy, but the solve then mutates the prob `u0`, so the next solve starts from the final state of the previous one: ```@example reuse +Random.seed!(4) # hide prob_u0 = deepcopy(jprob) -solve(prob_u0, SSAStepper(); seed = 1, alias = SciMLBase.DiscreteAliasSpecifier(alias_u0 = true)) +solve(prob_u0, SSAStepper(); alias = SciMLBase.DiscreteAliasSpecifier(alias_u0 = true)) prob_u0.prob.u0 ``` From fa5c52028e4559065e5068b7470dfa0e46c2bbf0 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sun, 11 Oct 2026 07:35:58 -0400 Subject: [PATCH 48/52] Build the docs against the downstream JumpProcesses 10 branches The Catalyst tutorial needs ModelingToolkitBase and Catalyst versions that support the JumpProcesses 10 mapper and solve-time RNG APIs. Their branches now exist, so docs/Project.toml pins Catalyst, ModelingToolkitBase, and StochasticDiffEqCore to them through [sources], with ModelingToolkitBase and StochasticDiffEqCore as direct docs dependencies so the pins apply. The pins track branches; exact commits will be pinned for the 10.0 release build, and the pins removed once those releases are registered. The asset copy of the docs project is kept in sync. The Documentation workflow runs on pushes and pull requests again, as on master. A local build against these branches ran every doctest and example, including the Catalyst tutorial, and rendered the full site. Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/Documentation.yml | 8 +++++--- docs/Project.toml | 12 ++++++++++++ docs/src/assets/Project.toml | 12 ++++++++++++ 3 files changed, 29 insertions(+), 3 deletions(-) diff --git a/.github/workflows/Documentation.yml b/.github/workflows/Documentation.yml index 6dec07596..fb51f43a6 100644 --- a/.github/workflows/Documentation.yml +++ b/.github/workflows/Documentation.yml @@ -1,9 +1,11 @@ name: Documentation on: - # Restore push (master and tags) and pull_request after ModelingToolkitBase - # supports the v10 JumpProblem and mapper APIs used by the Catalyst tutorial. - workflow_dispatch: + push: + branches: + - master + tags: '*' + pull_request: jobs: build: diff --git a/docs/Project.toml b/docs/Project.toml index 09fd1af98..747ee7252 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -6,6 +6,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5" +ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" OrdinaryDiffEqFunctionMap = "d3585ca7-f5d3-4ba6-8057-292ed1abd90f" OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce" @@ -13,6 +14,7 @@ OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" PointProcesses = "af0b7596-9bb0-472a-a012-63904f2b4c55" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +StochasticDiffEqCore = "19c5a474-6cd1-4a5f-be79-46dc34e54d7f" StochasticDiffEqHighOrder = "0520c28c-50fd-4d16-9c96-902fc80b3bab" [compat] @@ -23,6 +25,7 @@ Distributions = "0.25" Documenter = "1.1" Graphs = "1.9" JumpProcesses = "9.13.2" +ModelingToolkitBase = "1.77" SciMLBase = "3" OrdinaryDiffEqFunctionMap = "1, 2" OrdinaryDiffEqRosenbrock = "1, 2" @@ -30,4 +33,13 @@ OrdinaryDiffEqTsit5 = "1, 2" Plots = "1.39" PointProcesses = "0.5" StableRNGs = "1.0" +StochasticDiffEqCore = "2.2.3" StochasticDiffEqHighOrder = "2" + +[sources] +# Temporary: build the docs against the JumpProcesses 10 branches of Catalyst, +# ModelingToolkitBase, and StochasticDiffEqCore (SciML/OrdinaryDiffEq.jl#4836). Pin exact +# commits for the 10.0 release build, and remove once those releases are registered. +Catalyst = {url = "https://github.com/isaacsas/Catalyst.jl.git", rev = "catalyst/jumpprocesses-v10"} +ModelingToolkitBase = {url = "https://github.com/isaacsas/ModelingToolkit.jl.git", rev = "mtkbase/jumpprocesses-v10", subdir = "lib/ModelingToolkitBase"} +StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "sdecore/jumpprocesses-v10", subdir = "lib/StochasticDiffEqCore"} diff --git a/docs/src/assets/Project.toml b/docs/src/assets/Project.toml index 09fd1af98..747ee7252 100644 --- a/docs/src/assets/Project.toml +++ b/docs/src/assets/Project.toml @@ -6,6 +6,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6" JumpProcesses = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5" +ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" OrdinaryDiffEqFunctionMap = "d3585ca7-f5d3-4ba6-8057-292ed1abd90f" OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce" @@ -13,6 +14,7 @@ OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" Plots = "91a5bcdd-55d7-5caf-9e0b-520d859cae80" PointProcesses = "af0b7596-9bb0-472a-a012-63904f2b4c55" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" +StochasticDiffEqCore = "19c5a474-6cd1-4a5f-be79-46dc34e54d7f" StochasticDiffEqHighOrder = "0520c28c-50fd-4d16-9c96-902fc80b3bab" [compat] @@ -23,6 +25,7 @@ Distributions = "0.25" Documenter = "1.1" Graphs = "1.9" JumpProcesses = "9.13.2" +ModelingToolkitBase = "1.77" SciMLBase = "3" OrdinaryDiffEqFunctionMap = "1, 2" OrdinaryDiffEqRosenbrock = "1, 2" @@ -30,4 +33,13 @@ OrdinaryDiffEqTsit5 = "1, 2" Plots = "1.39" PointProcesses = "0.5" StableRNGs = "1.0" +StochasticDiffEqCore = "2.2.3" StochasticDiffEqHighOrder = "2" + +[sources] +# Temporary: build the docs against the JumpProcesses 10 branches of Catalyst, +# ModelingToolkitBase, and StochasticDiffEqCore (SciML/OrdinaryDiffEq.jl#4836). Pin exact +# commits for the 10.0 release build, and remove once those releases are registered. +Catalyst = {url = "https://github.com/isaacsas/Catalyst.jl.git", rev = "catalyst/jumpprocesses-v10"} +ModelingToolkitBase = {url = "https://github.com/isaacsas/ModelingToolkit.jl.git", rev = "mtkbase/jumpprocesses-v10", subdir = "lib/ModelingToolkitBase"} +StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "sdecore/jumpprocesses-v10", subdir = "lib/StochasticDiffEqCore"} From 146338434b97c21201d1deb31ef7f151210a1ee9 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sun, 11 Oct 2026 09:05:55 -0400 Subject: [PATCH 49/52] Test that VR_FRM thresholds are redrawn at every solve after remake remake now zeroes jump_u rather than resampling it, relying on the VR_FRM callback initializers to draw new thresholds from the integrator's RNG at every init. A new testset checks this for ODE (Tsit5) and SDE (SOSRI, EM) problems when solving the remade problem, solving the original problem (whose aliased jump_u remake also zeroes), and re-solving one integrator with reinit!. For a single constant-rate jump, every solve must succeed, the threshold saved at t = 0 after initialization must be nonzero and give the first jump time, and the first jump times must be distinct and pass a Kolmogorov-Smirnov test against the exponential distribution. Also corrects the remake comment that still described jump_u as resampled. Co-Authored-By: Claude Opus 5.5 (1M context) --- src/problem.jl | 3 +- test/extended_jump_array_remake.jl | 77 +++++++++++++++++++++++++++++- 2 files changed, 78 insertions(+), 2 deletions(-) diff --git a/src/problem.jl b/src/problem.jl index 27a24253f..ce4e053eb 100644 --- a/src/problem.jl +++ b/src/problem.jl @@ -238,7 +238,8 @@ function DiffEqBase.remake(jprob::JumpProblem; u0 = missing, p = missing, resolved_u0, _ = SciMLBase.updated_u0_p( prob, u0, missing; interpret_symbolicmap, use_defaults) - # Extract state values and wrap with fresh jump_u (resampled) + # Extract state values and wrap with a zeroed jump_u; the `VR_FRM` callbacks + # draw new thresholds from the integrator's RNG at every `init`. # This handles both: # - symbolic maps (resolved_u0 is ExtendedJumpArray with updated u) # - numeric arrays (resolved_u0 is Vector{Float64}) diff --git a/test/extended_jump_array_remake.jl b/test/extended_jump_array_remake.jl index 5f6e4add3..c64325ebe 100644 --- a/test/extended_jump_array_remake.jl +++ b/test/extended_jump_array_remake.jl @@ -1,7 +1,8 @@ # Tests for JumpProblem remake with VariableRateJumps (ExtendedJumpArray case) # This tests the fix for symbolic u0 with ExtendedJumpArray -using JumpProcesses, OrdinaryDiffEq, SciMLBase, Test, SymbolicIndexingInterface +using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq, SciMLBase, Test, + SymbolicIndexingInterface using StableRNGs @testset "remake JumpProblem with VariableRateJumps (ExtendedJumpArray)" begin @@ -119,3 +120,77 @@ using StableRNGs @test sol1.u[2].jump_u != sol3.u[2].jump_u end end + +# `remake` zeroes `jump_u`, so every solve must draw new `VR_FRM` thresholds. A single jump +# with constant rate `λ` first fires at an Exp(λ) time `τ`, and the threshold saved at +# `t = 0` after callback initialization must be `jump_u = -λτ`. The first saved point at +# `t = 0` precedes callback initialization, so its `jump_u` is not a threshold. +@testset "VR_FRM draws new thresholds at every solve after remake" begin + λ = 2.0 + N = 1000 + f!(du, u, p, t) = (du .= 0; nothing) + g!(du, u, p, t) = (du[1] = 0.1; du[2] = 0; nothing) + rate(u, p, t) = λ + affect!(integrator) = (integrator.u[2] += 1; nothing) + vrj = VariableRateJump(rate, affect!) + + # Returns the first jump time, and whether the solve succeeded and the threshold saved + # at `t = 0` after initialization is nonzero and gives that time. + function check_sol(sol) + i0 = findlast(==(0.0), sol.t) + threshold = -sol.u[i0].jump_u[1] + i1 = findfirst(u -> u.u[2] > 0, sol.u) + i1 === nothing && return (Inf, false) + τ = sol.t[i1] + τ, SciMLBase.successful_retcode(sol) && (sol.u[i1].u[2] == 1) && + (threshold > 0) && isapprox(τ, threshold / λ; rtol = 1e-6) + end + + # Kolmogorov-Smirnov distance between the sampled times and Exp(λ) + function ks_distance(τs) + n = length(τs) + F = 1 .- exp.(-λ .* sort(τs)) + maximum(max.((1:n) ./ n .- F, F .- (0:(n - 1)) ./ n)) + end + ks_critical = 1.95 / sqrt(N) # significance level 0.001 + + function test_results(results) + τs = first.(results) + @test all(last, results) + @test allunique(τs) + @test ks_distance(τs) < ks_critical + end + + tspan = (0.0, 10.0) + cases = [("ODE", ODEProblem(f!, [1.0, 0.0], tspan), Tsit5(), (;)), + ("SDE", SDEProblem(f!, g!, [1.0, 0.0], tspan), SOSRI(), (;)), + ("SDE", SDEProblem(f!, g!, [1.0, 0.0], tspan), EM(), (; dt = 0.01))] + for (name, prob, alg, kwargs) in cases + @testset "$name, $(nameof(typeof(alg)))" begin + jprob = JumpProblem(prob, vrj; vr_aggregator = VR_FRM()) + jprob2 = remake(jprob; u0 = [2.0, 0.0]) + @test all(iszero, jprob2.prob.u0.jump_u) + # `remake` also zeroed the original problem's `jump_u`, which it aliases + @test all(iszero, jprob.prob.u0.jump_u) + rng = StableRNG(1234) + + @testset "solve remade problem" begin + test_results([check_sol(solve(jprob2, alg; rng, kwargs...)) for _ in 1:N]) + end + + @testset "solve original problem" begin + test_results([check_sol(solve(jprob, alg; rng, kwargs...)) for _ in 1:N]) + end + + @testset "reinit! integrator of remade problem" begin + integrator = init(jprob2, alg; rng, kwargs...) + results = map(1:N) do _ + reinit!(integrator) + solve!(integrator) + check_sol(integrator.sol) + end + test_results(results) + end + end + end +end From b1b7d01a88144a5297c55b7ab53d8ec3938d0699 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sun, 11 Oct 2026 09:32:59 -0400 Subject: [PATCH 50/52] Restore final-value checks in the ODE + VR RNG stream tests Commit 33b86e9c dropped the final-value uniqueness checks from the serial ensemble testset, and swapped them for first jump times in the different-seeds testset, on the grounds that discrete final states could collide by chance. In this ODE model the state decays continuously and jumps subtract one, so the final value depends continuously on every jump time; collisions are no more likely than for first jump times. Unlike first jump times, final values also cover the jumps after the first, so both checks are kept, as in the sequential solves testset. Co-Authored-By: Claude Opus 5.5 (1M context) --- test/ensemble_problems.jl | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/test/ensemble_problems.jl b/test/ensemble_problems.jl index d0ad24fff..d04c81e0f 100644 --- a/test/ensemble_problems.jl +++ b/test/ensemble_problems.jl @@ -55,6 +55,8 @@ first_jump_time(traj) = traj.t[findfirst(>(traj.t[1]), traj.t)] trajectories = 3, rng = StableRNG(12345)) times = [first_jump_time(sol.u[i]) for i in 1:3] @test allunique(times) + finals = [sol.u[i].u[end][1] for i in 1:3] + @test allunique(finals) end # EM() uses a fixed time grid so jump event times aren't directly visible @@ -127,6 +129,8 @@ end sols = [solve(jprob, Tsit5(); rng = StableRNG(s)) for s in (100, 200, 300)] times = [first_jump_time(s) for s in sols] @test allunique(times) + finals = [s.u[end][1] for s in sols] + @test allunique(finals) end end From 847f195ad77ab57e55537b3409b56e56110329c3 Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sun, 11 Oct 2026 12:45:27 -0400 Subject: [PATCH 51/52] Prepare the transitional 10.0.0 release Bump the version to 10.0.0. StochasticDiffEq, ModelingToolkit, and Catalyst do not support JumpProcesses 10 yet, so environments with them keep JumpProcesses 9.x until their updates and 10.0.1 are released. The README, the docs home page, and HISTORY say so, and that jump problems coupled to DDEs are not supported in 10.0 at this time. The temporary test and docs pins stay, now at exact commits: the StochasticDiffEq branch (SciML/OrdinaryDiffEq.jl#4836) for StochasticDiffEqCore and, since registered versions of them do not admit JumpProcesses 10, StochasticDiffEqImplicit and StochasticDiffEqLeaping as test extras; and the Catalyst (SciML/Catalyst.jl#1572) and ModelingToolkitBase (SciML/ModelingToolkit.jl#5269) branches for the docs, whose JumpProcesses compat becomes "10". Those branches now admit JumpProcesses 10. The pins only affect running the tests and building the docs; they are removed in 10.0.1. Co-Authored-By: Claude Opus 5.5 (1M context) --- HISTORY.md | 8 +++++--- Project.toml | 19 +++++++++++++------ README.md | 9 +++++++++ docs/Project.toml | 15 ++++++++------- docs/src/assets/Project.toml | 15 ++++++++------- docs/src/index.md | 11 +++++++++++ 6 files changed, 54 insertions(+), 23 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 9e44c846a..020d3efb0 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -2,9 +2,11 @@ ## 10.0 (Breaking) -This section describes the forthcoming 10.0 release. During integration, -`Project.toml` retains version 9.33.1 for development and testing; these breaking -changes are not intended for a 9.x release. +JumpProcesses 10.0 is a transitional release. StochasticDiffEq, ModelingToolkit, and +Catalyst do not yet support it, so environments that include them keep JumpProcesses +9.x until their updates, and JumpProcesses 10.0.1, are released. Jump problems coupled +to delay differential equations (DDEs) are not supported in JumpProcesses 10.0 at this +time. - **Breaking**: The `rng` keyword argument has been removed from `JumpProblem`. Pass `rng` to `solve` or `init` instead: diff --git a/Project.toml b/Project.toml index 4ebce5804..566eaa892 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "JumpProcesses" uuid = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5" authors = ["Chris Rackauckas "] -version = "9.33.1" +version = "10.0.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" @@ -65,16 +65,21 @@ StaticArrays = "1.9.18" Statistics = "1" StochasticDiffEq = "7.0.0, 7" StochasticDiffEqCore = "2.2.3" +StochasticDiffEqImplicit = "2.2.1" +StochasticDiffEqLeaping = "2.1.0" SymbolicIndexingInterface = "0.3.48" Test = "1" julia = "1.10" [sources] -# Temporary: tests use the StochasticDiffEqCore branch for JumpProcesses 10 +# Temporary: tests use a commit of the StochasticDiffEq branch for JumpProcesses 10 # (OrdinaryDiffEq.jl#4836), which removes the need for a JumpProcesses StochasticDiffEqCore -# extension. Remove when it is released. Julia 1.10 ignores [sources]; see -# test/test_groups.toml. -StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "sdecore/jumpprocesses-v10", subdir = "lib/StochasticDiffEqCore"} +# extension. Implicit and Leaping are pinned too, since their registered versions do not +# admit JumpProcesses 10. Remove when those releases are registered. Julia 1.10 ignores +# [sources]; see test/test_groups.toml. +StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "1de403a5e59474fab9a43bf60d376391734eb65b", subdir = "lib/StochasticDiffEqCore"} +StochasticDiffEqImplicit = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "1de403a5e59474fab9a43bf60d376391734eb65b", subdir = "lib/StochasticDiffEqImplicit"} +StochasticDiffEqLeaping = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "1de403a5e59474fab9a43bf60d376391734eb65b", subdir = "lib/StochasticDiffEqLeaping"} [extras] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" @@ -92,7 +97,9 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" StochasticDiffEqCore = "19c5a474-6cd1-4a5f-be79-46dc34e54d7f" +StochasticDiffEqImplicit = "5080b986-4c76-4669-b5dc-373a41579d5b" +StochasticDiffEqLeaping = "aefaaa88-39f2-4e89-b162-d61b7e8cc81b" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Adapt", "FastBroadcast", "ForwardDiff", "KernelAbstractions", "LinearAlgebra", "LinearSolve", "OrdinaryDiffEq", "OrdinaryDiffEqFunctionMap", "Pkg", "SafeTestsets", "SciMLTesting", "StableRNGs", "Statistics", "StochasticDiffEq", "StochasticDiffEqCore", "Test"] +test = ["Adapt", "FastBroadcast", "ForwardDiff", "KernelAbstractions", "LinearAlgebra", "LinearSolve", "OrdinaryDiffEq", "OrdinaryDiffEqFunctionMap", "Pkg", "SafeTestsets", "SciMLTesting", "StableRNGs", "Statistics", "StochasticDiffEq", "StochasticDiffEqCore", "StochasticDiffEqImplicit", "StochasticDiffEqLeaping", "Test"] diff --git a/README.md b/README.md index cb0e09963..6729e4c8c 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,15 @@ [![ColPrac: Contributor's Guide on Collaborative Practices for Community Packages](https://img.shields.io/badge/ColPrac-Contributor%27s%20Guide-blueviolet)](https://github.com/SciML/ColPrac) [![SciML Code Style](https://img.shields.io/static/v1?label=code%20style&message=SciML&color=9558b2&labelColor=389826)](https://github.com/SciML/SciMLStyle) +> [!WARNING] +> **JumpProcesses 10.0 is a transitional release.** StochasticDiffEq, ModelingToolkit, and +> Catalyst do not yet support it and cannot be installed alongside it, so environments that +> include them will keep JumpProcesses 9.x. We recommend waiting for JumpProcesses 10.0.1, +> which will follow once those packages are updated, unless you use JumpProcesses on its +> own. Jump problems coupled to delay differential equations (DDEs) are not supported in +> JumpProcesses 10.0 at this time. See [HISTORY.md](HISTORY.md) for the breaking changes in +> 10.0. + JumpProcesses.jl provides methods for simulating jump processes, known as stochastic simulation algorithms (SSAs), Doob's method, Gillespie methods, or Kinetic Monte Carlo methods across different fields of science. It also enables the diff --git a/docs/Project.toml b/docs/Project.toml index 747ee7252..d6d928740 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -24,7 +24,7 @@ DiffEqCallbacks = "4.17.0" Distributions = "0.25" Documenter = "1.1" Graphs = "1.9" -JumpProcesses = "9.13.2" +JumpProcesses = "10" ModelingToolkitBase = "1.77" SciMLBase = "3" OrdinaryDiffEqFunctionMap = "1, 2" @@ -37,9 +37,10 @@ StochasticDiffEqCore = "2.2.3" StochasticDiffEqHighOrder = "2" [sources] -# Temporary: build the docs against the JumpProcesses 10 branches of Catalyst, -# ModelingToolkitBase, and StochasticDiffEqCore (SciML/OrdinaryDiffEq.jl#4836). Pin exact -# commits for the 10.0 release build, and remove once those releases are registered. -Catalyst = {url = "https://github.com/isaacsas/Catalyst.jl.git", rev = "catalyst/jumpprocesses-v10"} -ModelingToolkitBase = {url = "https://github.com/isaacsas/ModelingToolkit.jl.git", rev = "mtkbase/jumpprocesses-v10", subdir = "lib/ModelingToolkitBase"} -StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "sdecore/jumpprocesses-v10", subdir = "lib/StochasticDiffEqCore"} +# Temporary: build the 10.0 docs against commits of the JumpProcesses 10 branches of +# Catalyst (SciML/Catalyst.jl#1572), ModelingToolkitBase (SciML/ModelingToolkit.jl#5269), +# and StochasticDiffEqCore (SciML/OrdinaryDiffEq.jl#4836). Remove once those releases are +# registered. +Catalyst = {url = "https://github.com/isaacsas/Catalyst.jl.git", rev = "6efbc9b8ea8aff6d1bd64904f0027f562b8e2a64"} +ModelingToolkitBase = {url = "https://github.com/isaacsas/ModelingToolkit.jl.git", rev = "13113b635fb8f1307e7e8511e0c707dac7171272", subdir = "lib/ModelingToolkitBase"} +StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "1de403a5e59474fab9a43bf60d376391734eb65b", subdir = "lib/StochasticDiffEqCore"} diff --git a/docs/src/assets/Project.toml b/docs/src/assets/Project.toml index 747ee7252..d6d928740 100644 --- a/docs/src/assets/Project.toml +++ b/docs/src/assets/Project.toml @@ -24,7 +24,7 @@ DiffEqCallbacks = "4.17.0" Distributions = "0.25" Documenter = "1.1" Graphs = "1.9" -JumpProcesses = "9.13.2" +JumpProcesses = "10" ModelingToolkitBase = "1.77" SciMLBase = "3" OrdinaryDiffEqFunctionMap = "1, 2" @@ -37,9 +37,10 @@ StochasticDiffEqCore = "2.2.3" StochasticDiffEqHighOrder = "2" [sources] -# Temporary: build the docs against the JumpProcesses 10 branches of Catalyst, -# ModelingToolkitBase, and StochasticDiffEqCore (SciML/OrdinaryDiffEq.jl#4836). Pin exact -# commits for the 10.0 release build, and remove once those releases are registered. -Catalyst = {url = "https://github.com/isaacsas/Catalyst.jl.git", rev = "catalyst/jumpprocesses-v10"} -ModelingToolkitBase = {url = "https://github.com/isaacsas/ModelingToolkit.jl.git", rev = "mtkbase/jumpprocesses-v10", subdir = "lib/ModelingToolkitBase"} -StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "sdecore/jumpprocesses-v10", subdir = "lib/StochasticDiffEqCore"} +# Temporary: build the 10.0 docs against commits of the JumpProcesses 10 branches of +# Catalyst (SciML/Catalyst.jl#1572), ModelingToolkitBase (SciML/ModelingToolkit.jl#5269), +# and StochasticDiffEqCore (SciML/OrdinaryDiffEq.jl#4836). Remove once those releases are +# registered. +Catalyst = {url = "https://github.com/isaacsas/Catalyst.jl.git", rev = "6efbc9b8ea8aff6d1bd64904f0027f562b8e2a64"} +ModelingToolkitBase = {url = "https://github.com/isaacsas/ModelingToolkit.jl.git", rev = "13113b635fb8f1307e7e8511e0c707dac7171272", subdir = "lib/ModelingToolkitBase"} +StochasticDiffEqCore = {url = "https://github.com/isaacsas/OrdinaryDiffEq.jl.git", rev = "1de403a5e59474fab9a43bf60d376391734eb65b", subdir = "lib/StochasticDiffEqCore"} diff --git a/docs/src/index.md b/docs/src/index.md index e8a28b74c..224cdd8eb 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -1,5 +1,16 @@ # JumpProcesses.jl: Stochastic Simulation Algorithms for Jump Processes, Jump-ODEs, and Jump-Diffusions +!!! warning "JumpProcesses 10.0 is a transitional release" + + StochasticDiffEq, ModelingToolkit, and Catalyst do not yet support JumpProcesses 10 + and cannot be installed alongside it, so environments that include them will keep + JumpProcesses 9.x. We recommend waiting for JumpProcesses 10.0.1, which will follow + once those packages are updated, unless you use JumpProcesses on its own. Jump + problems coupled to delay differential equations (DDEs) are not supported in + JumpProcesses 10.0 at this time. See the + [release notes](https://github.com/SciML/JumpProcesses.jl/blob/master/HISTORY.md) + for the breaking changes in 10.0. + JumpProcesses.jl provides methods for simulating jump and point processes. Across different fields of science, such methods are also known as stochastic simulation algorithms (SSAs), Doob's method, Gillespie methods, Kinetic Monte From eea393512e06782f3c5e5db6bdc2a92a1087204a Mon Sep 17 00:00:00 2001 From: Sam Isaacson Date: Sun, 11 Oct 2026 12:54:55 -0400 Subject: [PATCH 52/52] Skip the lts ThreadSafety job until the StochasticDiffEq pins go At 10.0.0 the test environment cannot resolve on Julia 1.10: it ignores [sources], and registered StochasticDiffEq packages do not admit JumpProcesses 10. The lts ThreadSafety job failed for this reason, so skip it like the lts InterfaceI and InterfaceII jobs, whose comments now give the same reason. Restore all three when the pins are removed. Co-Authored-By: Claude Opus 5.5 (1M context) --- .github/workflows/ThreadSafety.yml | 4 +++- test/test_groups.toml | 12 ++++++------ 2 files changed, 9 insertions(+), 7 deletions(-) diff --git a/.github/workflows/ThreadSafety.yml b/.github/workflows/ThreadSafety.yml index 45939c730..bfd4b4e90 100644 --- a/.github/workflows/ThreadSafety.yml +++ b/.github/workflows/ThreadSafety.yml @@ -23,7 +23,9 @@ jobs: fail-fast: false matrix: version: - - "lts" + # Restore "lts" when the temporary StochasticDiffEq pins are removed: Julia 1.10 + # ignores [sources], and registered StochasticDiffEq packages do not admit + # JumpProcesses 10, so the test environment cannot resolve. - "1" - "pre" runs-on: ubuntu-latest diff --git a/test/test_groups.toml b/test/test_groups.toml index c953a3782..abb9b127d 100644 --- a/test/test_groups.toml +++ b/test/test_groups.toml @@ -1,7 +1,7 @@ [InterfaceI] -# Restore "lts" when the temporary StochasticDiffEqCore pin is removed: Julia 1.10 -# ignores [sources], and registered StochasticDiffEqCore lacks the stochastic `init` -# dispatch these tests check. +# Restore "lts" when the temporary StochasticDiffEq pins are removed: Julia 1.10 +# ignores [sources], and registered StochasticDiffEq packages do not admit +# JumpProcesses 10, so the test environment cannot resolve. versions = ["1", "pre"] ["InterfaceI 32-bit"] @@ -11,9 +11,9 @@ os = ["ubuntu-latest"] arch = "x86" [InterfaceII] -# Restore "lts" when the temporary StochasticDiffEqCore pin is removed: Julia 1.10 -# ignores [sources], and registered StochasticDiffEqCore lacks the stochastic `init` -# dispatch these tests check. +# Restore "lts" when the temporary StochasticDiffEq pins are removed: Julia 1.10 +# ignores [sources], and registered StochasticDiffEq packages do not admit +# JumpProcesses 10, so the test environment cannot resolve. versions = ["1", "pre"] [QA]