diff --git a/.github/workflows/Downgrade.yml b/.github/workflows/Downgrade.yml index 26c77b953..d5b91b32f 100644 --- a/.github/workflows/Downgrade.yml +++ b/.github/workflows/Downgrade.yml @@ -13,7 +13,11 @@ 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 + # 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 secrets: "inherit" diff --git a/.github/workflows/ThreadSafety.yml b/.github/workflows/ThreadSafety.yml new file mode 100644 index 000000000..bfd4b4e90 --- /dev/null +++ b/.github/workflows/ThreadSafety.yml @@ -0,0 +1,67 @@ +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: + # 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 + steps: + - uses: actions/checkout@v7 + + - name: "Setup Julia ${{ matrix.version }}" + uses: julia-actions/setup-julia@v3 + 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@v7 + 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/HISTORY.md b/HISTORY.md index 2929b9f5b..020d3efb0 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,6 +1,178 @@ # Breaking updates and feature summaries across releases -## JumpProcesses unreleased (master branch) +## 10.0 (Breaking) + +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: + ```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)) + ``` + - **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**: 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, 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 + `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 + 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**: 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. + - 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 + `merge_callbacks = false`. `init` now combines stored and call-level callbacks + exactly as `solve` does. + - 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 + state types" in the solver documentation. + - **Breaking**: Multi-trajectory `EnsembleGPUKernel` solves reject `alias` values + 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!`). + - **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 + # 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`) 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 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. ## 9.33.1 @@ -28,7 +200,6 @@ ```julia using JumpProcesses - rate(u, p, t) = p[1] * u[1] / (1 + u[2]) function affect!(integrator) integrator.u[1] -= 1 @@ -41,10 +212,8 @@ 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]] @@ -67,7 +236,7 @@ - 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] @@ -82,10 +251,11 @@ `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/Project.toml b/Project.toml index 9c002e4fc..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" @@ -52,22 +52,35 @@ 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" SafeTestsets = "0.1" -SciMLBase = "3.18.0, 3.1" +SciMLBase = "3.34" SciMLTesting = "2.8" SimpleNonlinearSolve = "2, 3" StableRNGs = "1" 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 a commit of the StochasticDiffEq branch for JumpProcesses 10 +# (OrdinaryDiffEq.jl#4836), which removes the need for a JumpProcesses 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" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" @@ -83,7 +96,10 @@ 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" +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", "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 e3606dbb3..6729e4c8c 100644 --- a/README.md +++ b/README.md @@ -6,20 +6,29 @@ + + [![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) +> [!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 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 @@ -38,32 +47,25 @@ 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). ## 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 @@ -131,7 +133,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) @@ -167,8 +169,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 6a59991e5..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, fixed rng for reproducibility) +# Solves (SSAStepper, seeded RNG) # ============================================================================= SUITE["solve"] = BenchmarkGroup() -jprob_maj = JumpProblem(dprob, Direct(), maj; rng = StableRNG(12345)) -jprob_cr = JumpProblem(dprob, Direct(), jump1, jump2; rng = StableRNG(12345)) +jprob_maj = jump_problem(dprob, Direct(), maj) +jprob_cr = jump_problem(dprob, Direct(), jump1, jump2) -SUITE["solve"]["massaction"] = @benchmarkable solve($jprob_maj, SSAStepper()) -SUITE["solve"]["constantrate"] = @benchmarkable solve($jprob_cr, SSAStepper()) +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()) 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 = 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()) -SUITE["aggregators"]["SortingDirect"] = @benchmarkable solve( - $jprob_sorting, SSAStepper() -) -SUITE["aggregators"]["NRM"] = @benchmarkable solve($jprob_nrm, SSAStepper()) +SUITE["aggregators"]["RDirect"] = benchmark_solve(jprob_rdirect) +SUITE["aggregators"]["SortingDirect"] = benchmark_solve(jprob_sorting) +SUITE["aggregators"]["NRM"] = benchmark_solve(jprob_nrm) diff --git a/docs/Project.toml b/docs/Project.toml index c29a7a37d..d6d928740 100644 --- a/docs/Project.toml +++ b/docs/Project.toml @@ -1,33 +1,46 @@ [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" +ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839" 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" +StochasticDiffEqCore = "19c5a474-6cd1-4a5f-be79-46dc34e54d7f" +StochasticDiffEqHighOrder = "0520c28c-50fd-4d16-9c96-902fc80b3bab" [compat] Catalyst = "16" DensityInterface = "0.4" -DifferentialEquations = "7.11, 8.0" +DiffEqCallbacks = "4.17.0" Distributions = "0.25" Documenter = "1.1" Graphs = "1.9" -JumpProcesses = "9.13.2" +JumpProcesses = "10" +ModelingToolkitBase = "1.77" 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" +StochasticDiffEqCore = "2.2.3" +StochasticDiffEqHighOrder = "2" + +[sources] +# 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/make.jl b/docs/make.jl index 9fd328ff4..ba37a4a6d 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,8 +21,11 @@ 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, + # 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/pages.jl b/docs/pages.jl index 1f973f241..53cbdd932 100644 --- a/docs/pages.jl +++ b/docs/pages.jl @@ -5,10 +5,12 @@ 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"], "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 c89a8bec7..454de5b4d 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 @@ -122,3 +123,216 @@ ODE/SDE integrators a `JumpProblem` can be paired with. Anything else from SciMLBase -- the BVP, DAE, DDE, nonlinear and optimization problem 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`, 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 + +JumpProcesses supports controlling the random number generator (RNG) used for +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 + +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)) +``` + +### Resolution priority + +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()` | + +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, the ODE/DAE pathway, and CPU tau-leaping +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 | + +!!! 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. + 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 + +```@docs +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 | +| 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 (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 | + +For **SSAStepper**, `rng` is resolved via `resolve_rng` in `SSA_stepper.jl`'s +`__init` and stored on the [`SSAIntegrator`](@ref). + +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 **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`, 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 +them again, and discard `merge_callbacks`. The `solve` path merges once in `__solve`. +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, +`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` +directly with the `rng` and `seed` kwargs from the `solve` call. diff --git a/docs/src/applications/advanced_point_process.md b/docs/src/applications/advanced_point_process.md index f88121875..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 @@ -389,10 +388,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 @@ -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) @@ -833,11 +832,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/assets/Project.toml b/docs/src/assets/Project.toml index 546453a95..d6d928740 100644 --- a/docs/src/assets/Project.toml +++ b/docs/src/assets/Project.toml @@ -1,29 +1,46 @@ [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" +ModelingToolkitBase = "7771a370-6774-4173-bd38-47e70ca0b839" 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" +StochasticDiffEqCore = "19c5a474-6cd1-4a5f-be79-46dc34e54d7f" +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" +JumpProcesses = "10" +ModelingToolkitBase = "1.77" 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" +StochasticDiffEqCore = "2.2.3" +StochasticDiffEqHighOrder = "2" + +[sources] +# 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/faq.md b/docs/src/faq.md index d2d0be3be..e2508c49a 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`? @@ -55,21 +70,31 @@ 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)`. -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, `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? @@ -122,11 +147,26 @@ 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 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 @@ -160,10 +200,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 @@ -186,9 +250,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/index.md b/docs/src/index.md index dbd5be6b3..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 @@ -7,9 +18,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 +128,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 @@ -146,8 +149,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_solve.md b/docs/src/jump_solve.md index 41ac44124..c9d51aceb 100644 --- a/docs/src/jump_solve.md +++ b/docs/src/jump_solve.md @@ -62,6 +62,38 @@ 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). + 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/jump_types.md b/docs/src/jump_types.md index 181964d91..b018421bd 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 @@ -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** @@ -124,7 +134,7 @@ MassActionJump(reactant_stoich, net_stoich; scale_rates = true, param_idxs = not `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) ``` @@ -132,31 +142,32 @@ MassActionJump(reactant_stoich, net_stoich; scale_rates = true, param_idxs = not - 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]] ``` @@ -267,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 @@ -417,9 +443,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. As such, only the new problem generated by -`remake` should be used for subsequent simulations. +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: @@ -478,29 +511,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/migration.md b/docs/src/migration.md new file mode 100644 index 000000000..a0d84997a --- /dev/null +++ b/docs/src/migration.md @@ -0,0 +1,243 @@ +# [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. + +## 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 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 +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. + +## 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/docs/src/tutorials/discrete_stochastic_example.md b/docs/src/tutorials/discrete_stochastic_example.md index 99db5078d..8fa9a0e3d 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 @@ -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) ``` @@ -477,6 +481,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. @@ -684,6 +693,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 @@ -720,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 @@ -854,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 @@ -926,6 +1030,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/ensembles_and_problem_reuse.md b/docs/src/tutorials/ensembles_and_problem_reuse.md new file mode 100644 index 000000000..72e75a2f0 --- /dev/null +++ b/docs/src/tutorials/ensembles_and_problem_reuse.md @@ -0,0 +1,291 @@ +# [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. +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 +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. +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 +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()).u[end][1] for _ 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) +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 chunk + Threads.@spawn begin + local_prob = deepcopy(jprob) # one independent copy per task + [solve(local_prob, SSAStepper()).u[end][1] for _ in chunk] + 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 +Random.seed!(3) # hide +integrator1 = init(deepcopy(jprob), SSAStepper()) +integrator2 = init(deepcopy(jprob), SSAStepper()) +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 +Random.seed!(4) # hide +prob_u0 = deepcopy(jprob) +solve(prob_u0, SSAStepper(); 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 StochasticDiffEqHighOrder +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/docs/src/tutorials/jump_diffusion.md b/docs/src/tutorials/jump_diffusion.md index 8b73b8df6..6a10f5d66 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 @@ -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 ae8e0955e..e5f27caf3 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 @@ -317,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. @@ -371,29 +370,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. 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 - -```@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/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 5df7b3876..948c8740c 100644 --- a/ext/JumpProcessesKernelAbstractionsExt.jl +++ b/ext/JumpProcessesKernelAbstractionsExt.jl @@ -5,6 +5,51 @@ 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 + +# 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, 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) || 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 && + 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 + include("ssa_stepper.jl") include("explicit_tau_leaping.jl") @@ -13,15 +58,18 @@ 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 + reject_kernel_alias_kwargs(ensembleprob.prob, kwargs) + seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob @@ -35,7 +83,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 +132,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 @@ -117,7 +164,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] @@ -152,7 +199,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 +246,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/JumpProcessesOrdinaryDiffEqCoreExt.jl b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl index 77f2631e4..ceeceb0c9 100644 --- a/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl +++ b/ext/JumpProcessesOrdinaryDiffEqCoreExt.jl @@ -1,27 +1,24 @@ module JumpProcessesOrdinaryDiffEqCoreExt using JumpProcesses -import DiffEqBase import SciMLBase import OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, DAEAlgorithm, - StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm + 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 function SciMLBase.__init( _jump_prob::JumpProcesses.JumpProblem{IIP, P}, - alg::Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm}; - kwargs...) where {IIP, P} - _jump_init(_jump_prob, alg; kwargs...) -end - -function SciMLBase.__init( - _jump_prob::JumpProcesses.JumpProblem{IIP, P}, - alg::Union{StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; + alg::Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm, + StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm}; kwargs...) where {IIP, P} + # Cover the full intersection with OrdinaryDiffEqCore's initializer. SDE/RODE + # initialization belongs to StochasticDiffEqCore, whose per-algorithm methods are + # more specific. _jump_init(_jump_prob, alg; kwargs...) end diff --git a/ext/explicit_tau_leaping.jl b/ext/explicit_tau_leaping.jl index 38bc6a524..908877d84 100644 --- a/ext/explicit_tau_leaping.jl +++ b/ext/explicit_tau_leaping.jl @@ -213,10 +213,19 @@ 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. 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. @@ -226,24 +235,24 @@ 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 + reject_kernel_alias_kwargs(ensembleprob.prob, kwargs) + seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob jump_prob isa JumpProblem || @@ -289,7 +298,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..f40c49935 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 @@ -54,12 +68,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 +282,19 @@ 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. 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. @@ -276,6 +304,7 @@ function SciMLBase.__solve(ensembleprob::SciMLBase.AbstractEnsembleProblem, ensemblealg::EnsembleGPUKernel; trajectories, seed = nothing, + rng = nothing, saveat = nothing, save_start = true, save_end = true, @@ -283,16 +312,20 @@ 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) + 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(ensembleprob.prob, kwargs) + seed_kernel_backend!(backend, seed, rng, kwargs) jump_prob = ensembleprob.prob jump_prob isa JumpProblem || @@ -340,7 +373,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/JumpProcesses.jl b/src/JumpProcesses.jl index 80acf38f1..be553da70 100644 --- a/src/JumpProcesses.jl +++ b/src/JumpProcesses.jl @@ -4,7 +4,7 @@ using DiffEqBase # 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 @@ -32,13 +32,14 @@ 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 using DiffEqBase: DiffEqBase, DAEFunction, DDEFunction, isinplace using SimpleNonlinearSolve: SimpleNonlinearSolve, SimpleNewtonRaphson using ADTypes: ADTypes, AutoFiniteDiff -using SciMLBase: SciMLBase, DEIntegrator, NonlinearProblem +using SciMLBase: SciMLBase, DEIntegrator, NonlinearProblem, get_rng # The SciML common interface that JumpProcesses reexports (see the second `export` # block below), so that `using JumpProcesses` on its own is enough to build the problem @@ -61,8 +62,6 @@ abstract type AbstractJumpAggregator end abstract type AbstractSSAIntegrator{Alg, IIP, U, T} <: DEIntegrator{Alg, IIP, U, T} end -const DEFAULT_RNG = Random.default_rng() - # thresholds for auto-alg below which the listed alg is used # see select_aggregator for details const USE_DIRECT_THRESHOLD = 20 @@ -181,7 +180,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") @@ -210,13 +209,16 @@ include("coupling.jl") affect! = integrator -> (integrator.u[1] += 1) jump = ConstantRateJump(rate, affect!) prob = DiscreteProblem([10.0], (0.0, 1.0)) - jump_prob = JumpProblem(prob, Direct(), jump; - rng = Random.MersenneTwister(12345), save_positions = (false, false)) + jump_prob = JumpProblem(prob, Direct(), jump; save_positions = (false, false)) @compile_workload begin integrator = init(jump_prob, SSAStepper()) step!(integrator) solve(jump_prob, SSAStepper()) + + seeded_integrator = init(jump_prob, SSAStepper(); seed = 12345) + step!(seeded_integrator) + solve(jump_prob, SSAStepper(); seed = 12345) end end diff --git a/src/SSA_stepper.jl b/src/SSA_stepper.jl index a13b93c62..8af965090 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. @@ -13,6 +63,26 @@ 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`. + - 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)`, @@ -56,62 +126,232 @@ 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 """ $(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, UP, TC} <: 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 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 direction time is changing in (must be positive, indicating time is increasing)""" + """ + 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 - """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. + """ + 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 +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 + +# 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 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 + return recursivecopy(u0), recursivecopy(u0) + end +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 -(integrator::SSAIntegrator)(t) = copy(integrator.u) -(integrator::SSAIntegrator)(out, t) = (out .= integrator.u) +# 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 @@ -155,27 +395,35 @@ 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) && - integrator.saveat[integrator.cur_saveat] < integrator.t - push!(integrator.sol.t, integrator.saveat[integrator.cur_saveat]) - push!(integrator.sol.u, copy(integrator.u)) + integrator.saveat[integrator.cur_saveat] < integrator.t + 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)) + # 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 && + (isempty(integrator.sol.t) || integrator.sol.t[end] != end_time) + save_current_state!(integrator, end_time) end end @@ -216,19 +464,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, - alias_jump = Threads.threadid() == 1, + rng = nothing, + 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.") @@ -242,25 +529,21 @@ function SciMLBase.__init(jump_prob::JumpProblem, # Check for continuous callbacks passed via kwargs (from JumpProblem constructor or solve) check_continuous_callback_error(callback) - 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 + + _rng = resolve_rng(rng, seed) + + # 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)[] @@ -294,22 +577,27 @@ 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) + # 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, tmp_cache) cb.initialize(cb, integrator.u, prob.tspan[1], integrator) DiffEqBase.initialize!(opts.callback, integrator.u, prob.tspan[1], integrator) if save_start @@ -368,6 +656,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) @@ -386,9 +675,8 @@ 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 - push!(integrator.sol.t, integrator.saveat[integrator.cur_saveat]) - push!(integrator.sol.u, copy(integrator.u)) + integrator.saveat[integrator.cur_saveat] < integrator.t + save_current_state!(integrator, integrator.saveat[integrator.cur_saveat]) integrator.cur_saveat += 1 end end @@ -425,8 +713,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 @@ -446,10 +733,13 @@ function should_continue_solve(integrator::SSAIntegrator) integrator.keep_stepping && (has_jump || has_tstop) end -function reset_aggregated_jumps!(integrator::SSAIntegrator, uprev = nothing; - update_jump_params = true, kwargs...) - reset_aggregated_jumps!(integrator, uprev, integrator.cb; - update_jump_params, kwargs...) +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 @@ -459,9 +749,8 @@ 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} + <:SSAStepper}) where {T, N, uType, uType2, DType, tType, rateType, discType, P} sol.dense end 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/aggregators.jl b/src/aggregators/aggregators.jl index ebaa13fc6..6bd27d992 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 @@ -58,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 @@ -71,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 @@ -153,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 @@ -164,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/aggregators/bracketing.jl b/src/aggregators/bracketing.jl index 9b56418b4..23f584aa1 100644 --- a/src/aggregators/bracketing.jl +++ b/src/aggregators/bracketing.jl @@ -79,8 +79,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 """ @@ -90,7 +90,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 p.brackets[rx - num_majumps](p.ulow, p.uhigh, u, params, t) end @@ -144,7 +144,8 @@ function validate_brackets(p::AbstractSSAJumpAggregator, u, params, t) lrate, urate = p.cur_rate_low[rx], p.cur_rate_high[rx] (zero(lrate) <= lrate <= urate < Inf) || error("Invalid rate bounds for jump $rx: expected 0 <= lrate <= urate < Inf, got ($lrate, $urate).") - rate = calculate_jump_rate(p.ma_jumps, num_majumps, p.rates, u, params, t, rx) + rate = calculate_jump_rate(p.ma_jumps, num_majumps, p.rates, u, params, t, rx, + p.maj_rates) (lrate <= rate <= urate) || error("Rate bounds for jump $rx do not bracket the rate at the initial state: $rate ∉ [$lrate, $urate].") end diff --git a/src/aggregators/ccnrm.jl b/src/aggregators/ccnrm.jl index cdf7d5900..24517da44 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 @@ -12,18 +12,19 @@ mutable struct CCNRMJumpAggregation{T, S, F1, F2, RNG, 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} - 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}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), + 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 +46,32 @@ 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, + crs, sr, maj, maj_rates, rs, affs!, sps, - rng, dg, ptt) + dg, ptt) end ############################# 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::Number) p.end_time = integrator.sol.prob.tspan[2] - initialize_rates_and_times!(p, u, params, t) + 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) nothing end @@ -79,7 +82,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 +92,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 +106,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) @@ -114,7 +118,7 @@ function update_dependent_rates!(p::CCNRMJumpAggregation, u, params, t) # 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,7 +129,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 +138,16 @@ 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)) + maj_rates = p.maj_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] + cur_rates[i] = evalrxrate(u, i, majumps, maj_rates) + pttdata[i] = t + randexp(rng) / cur_rates[i] end # Initialize next-reaction times for the constant rates @@ -150,11 +155,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 c9cfa8016..1bdf270be 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 @@ -10,10 +10,10 @@ mutable struct CoevolveJumpAggregation{T, S, F1, F2, RNG, 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 - 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 +24,11 @@ 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}; + 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, 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 +50,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), - typeof(pq)}(nj, nj, njt, et, crs, sr, maj, - rs, affs!, sps, rng, dg, pq, + CoevolveJumpAggregation{T, S, F1, affecttype, typeof(dg), + typeof(pq)}(nj, nj, njt, et, crs, sr, maj, maj_rates, + rs, affs!, sps, dg, pq, lrates, urates, rateintervals, haslratevec, cur_lrates) end @@ -98,7 +99,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)}} @@ -136,20 +137,23 @@ 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, rng; - u, dep_graph, lrates, urates, rateintervals, haslratevec, + ma_jumps, rates, affects!, save_positions; + 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::Number) p.end_time = integrator.sol.prob.tspan[2] - fill_rates_and_get_times!(p, u, params, t) + 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) nothing end @@ -160,7 +164,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 +183,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 +231,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 @@ -237,7 +243,7 @@ function update_dependent_rates!(p::CoevolveJumpAggregation, u, params, t) 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) @@ -256,8 +262,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 +306,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 d1045794a..061318c0a 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 @@ -7,46 +7,48 @@ mutable struct DirectJumpAggregation{T, S, F1, F2, RNG} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} 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}; + 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, RNG}(nj, nj, njt, et, crs, sr, maj, rs, - affs!, sps, rng) + DirectJumpAggregation{T, S, F1, affecttype}(nj, nj, njt, et, crs, sr, maj, maj_rates, + 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 function initialize!(p::DirectJumpAggregation, integrator, u, params, t::Number) 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 @@ -60,9 +62,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,16 +73,17 @@ 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 # 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 @@ -96,7 +100,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 # Unrolled at compile time: recursive splatting over `rates...` is only inferred for @@ -112,16 +116,17 @@ 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 # 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 @@ -137,5 +142,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 de84459a4..0ba564647 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 @@ -20,10 +20,10 @@ mutable struct DirectCRJumpAggregation{T, S, F1, F2, RNG, DEPGR, U <: PriorityTa cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} 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 +32,12 @@ 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}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), + 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 +64,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, + maj_rates, rs, affs!, sps, dg, minrate, maxrate, rt, ratetogroup) end @@ -74,19 +75,20 @@ 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 # set up a new simulation and calculate the first jump / jump time function initialize!(p::DirectCRJumpAggregation, integrator, u, params, t::Number) 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) @@ -113,10 +115,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 @@ -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 6a3c3bddb..8a4c7e2bf 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 @@ -7,48 +7,50 @@ mutable struct FRMJumpAggregation{T, S, F1, F2, RNG} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} 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}; + 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, RNG}(nj, nj, njt, et, crs, sr, maj, rs, - affs!, sps, rng) + FRMJumpAggregation{T, S, F1, affecttype}(nj, nj, njt, et, crs, sr, maj, maj_rates, + 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 function initialize!(p::FRMJumpAggregation, integrator, u, params, t::Number) 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 @@ -62,8 +64,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 +82,14 @@ 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 + maj_rates = p.maj_rates @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] + p.cur_rates[i] = evalrxrate(u, i, majumps, maj_rates) + dt = randexp(rng) / p.cur_rates[i] if dt < ttnj ttnj = dt nextrx = i @@ -95,15 +99,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 +119,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 95975f1d2..dc4c3ca26 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 @@ -10,18 +10,19 @@ mutable struct NRMJumpAggregation{T, S, F1, F2, RNG, 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} - 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}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), + 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 +41,31 @@ 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, - crs, sr, maj, + NRMJumpAggregation{T, S, F1, affecttype, typeof(dg), typeof(pq)}(nj, nj, njt, et, + crs, sr, maj, maj_rates, 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::Number) p.end_time = integrator.sol.prob.tspan[2] - fill_rates_and_get_times!(p, u, params, t) + 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) nothing end @@ -73,7 +76,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 +91,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) @@ -97,7 +101,7 @@ function update_dependent_rates!(p::NRMJumpAggregation, u, params, t) # 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) @@ -108,7 +112,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,15 +122,16 @@ 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 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) - pqdata[i] = t + randexp(p.rng) / cur_rates[i] + cur_rates[i] = evalrxrate(u, i, majumps, maj_rates) + pqdata[i] = t + randexp(rng) / cur_rates[i] end # constant rates @@ -134,7 +139,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/prioritytable.jl b/src/aggregators/prioritytable.jl index b7af359ff..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 @@ -269,7 +289,7 @@ end return gid - mingid + 2 end -@inline function sample(pg::PriorityGroup, priorities, rng = DEFAULT_RNG) +@inline function sample(pg::PriorityGroup, priorities, rng) (; maxpriority, numpids, pids) = pg pididx = 0 @@ -289,7 +309,7 @@ end pid end -function sample(pt::PriorityTable, priorities, rng = DEFAULT_RNG) +function sample(pt::PriorityTable, priorities, rng) (; groups, gsum, gsums) = pt # return id zero if total priority is zero diff --git a/src/aggregators/rdirect.jl b/src/aggregators/rdirect.jl index 7ccc7be4c..463a4ad13 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 @@ -11,10 +11,10 @@ mutable struct RDirectJumpAggregation{T, S, F1, F2, RNG, DEPGR} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 save_positions::Tuple{Bool, Bool} - rng::RNG dep_gr::DEPGR max_rate::T counter::Int @@ -22,10 +22,11 @@ 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}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), 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 +43,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, - crs, sr, maj, rs, - affs!, sps, rng, + return RDirectJumpAggregation{T, S, F1, affecttype, typeof(dg)}(nj, nj, njt, et, + crs, sr, maj, maj_rates, rs, + affs!, sps, dg, max_rate, 0, Int(counter_threshold)) end @@ -53,19 +54,20 @@ 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 # set up a new simulation and calculate the first jump / jump time function initialize!(p::RDirectJumpAggregation, integrator, u, params, t::Number) 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) @@ -93,7 +95,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 +108,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 @@ -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 e14c9420b..caa3b7c27 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, CB, VJMAP, JVMAP, BD, U} <: - AbstractSSAJumpAggregator{T, S, F1, F2, RNG} +mutable struct RSSAJumpAggregation{T, S, F1, F2, CB, VJMAP, JVMAP, BD, U} <: + AbstractSSAJumpAggregator{T, S, F1, F2} next_jump::Int prev_jump::Int next_jump_time::T @@ -14,11 +14,11 @@ mutable struct RSSAJumpAggregation{T, S, F1, F2, RNG, CB, VJMAP, JVMAP, BD, U} < cur_rate_high::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} rates::F1 affects!::F2 brackets::CB save_positions::Tuple{Bool, Bool} - rng::RNG vartojumps_map::VJMAP jumptovars_map::JVMAP bracket_data::BD @@ -27,10 +27,11 @@ mutable struct RSSAJumpAggregation{T, S, F1, F2, RNG, CB, 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, brackets, vartojumps_map = nothing, + maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), + u::U, brackets, 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) @@ -64,10 +65,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(brackets), typeof(vtoj_map), + RSSAJumpAggregation{T, S, F1, affecttype, typeof(brackets), typeof(vtoj_map), typeof(jtov_map), typeof(bd), U}(nj, nj, njt, et, crl_bnds, - crh_bnds, sr, maj, rs, affs!, brackets, sps, - rng, vtoj_map, jtov_map, bd, ulow, + crh_bnds, sr, maj, maj_rates, rs, affs!, brackets, sps, + vtoj_map, jtov_map, bd, ulow, uhigh) end @@ -75,20 +76,21 @@ 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) brackets = get_jump_bracket_fwrappers(u, p, t, constant_jumps, aggregator) build_jump_aggregation(RSSAJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; u = u, brackets, + rates, affects!, save_positions; u, brackets, kwargs...) end # set up a new simulation and calculate the first jump / jump time function initialize!(p::RSSAJumpAggregation, integrator, u, params, t::Number) 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 @@ -110,7 +112,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) @@ -122,7 +125,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) @@ -134,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 ######################### """ @@ -161,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 @@ -187,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 664ad4b15..9bf3f034f 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, CB, U, VJMAP, JVMAP, BD, +mutable struct RSSACRJumpAggregation{F, S, F1, F2, CB, 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 @@ -15,11 +15,11 @@ mutable struct RSSACRJumpAggregation{F, S, F1, F2, RNG, CB, U, VJMAP, JVMAP, BD, cur_rate_high::Vector{F} sum_rate::F ma_jumps::S + maj_rates::Vector{F} rates::F1 affects!::F2 brackets::CB save_positions::Tuple{Bool, Bool} - rng::RNG vartojumps_map::VJMAP jumptovars_map::JVMAP bracket_data::BD @@ -32,11 +32,12 @@ mutable struct RSSACRJumpAggregation{F, S, F1, F2, RNG, CB, 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, brackets, - 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, brackets, 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) @@ -81,10 +82,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, typeof(brackets), U, + RSSACRJumpAggregation{typeof(njt), S, F1, affecttype, typeof(brackets), 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!, brackets, sps, rng, vtoj_map, + sum_rate, maj, maj_rates, rs, affs!, brackets, sps, vtoj_map, jtov_map, bd, ulow, uhigh, minrate, maxrate, rt, ratetogroup) end @@ -93,19 +94,20 @@ 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) brackets = get_jump_bracket_fwrappers(u, p, t, constant_jumps, aggregator) build_jump_aggregation(RSSACRJumpAggregation, u, p, t, end_time, ma_jumps, - rates, affects!, save_positions, rng; u = u, brackets, kwargs...) + rates, affects!, save_positions; u, brackets, kwargs...) end # set up a new simulation and calculate the first jump / jump time function initialize!(p::RSSACRJumpAggregation, integrator, u, params, t::Number) 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 @@ -136,7 +138,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) @@ -147,7 +150,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) @@ -179,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]) @@ -207,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/sortingdirect.jl b/src/aggregators/sortingdirect.jl index 6a44a0244..541d427f1 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 @@ -11,19 +11,20 @@ mutable struct SortingDirectJumpAggregation{T, S, F1, F2, RNG, DEPGR} <: cur_rates::Vector{T} sum_rate::T ma_jumps::S + maj_rates::Vector{T} 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}; + maj_rates = Vector{T}(undef, get_num_majumps(maj)), + 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 +43,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, - crs, sr, maj, rs, - affs!, sps, rng, + SortingDirectJumpAggregation{T, S, F1, affecttype, typeof(dg)}(nj, nj, njt, et, + crs, sr, maj, maj_rates, rs, + affs!, sps, dg, jtoidx, zero(Int)) end @@ -53,19 +54,20 @@ 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 # set up a new simulation and calculate the first jump / jump time function initialize!(p::SortingDirectJumpAggregation, integrator, u, params, t::Number) 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 @@ -92,14 +94,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 c9789e897..cc6a74d51 100644 --- a/src/aggregators/ssajump.jl +++ b/src/aggregators/ssajump.jl @@ -13,13 +13,12 @@ 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) @@ -112,12 +111,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 @@ -129,12 +128,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, rng; kwargs...) + majumps, rates, affects!, save_positions; maj_rates, kwargs...) end """ @@ -148,8 +148,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 @@ -171,9 +172,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 @@ -196,7 +198,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 @@ -298,7 +300,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] @@ -307,7 +309,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/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 48ac3e925..a1f7e645d 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,64 +515,95 @@ 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. + + - `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, reaction rates are rescaled when constructing the `MassActionJump` - 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. + - 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 - """The (scaled) reaction rate constants.""" + """ + 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.""" + """ + 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 `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, 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) @@ -596,7 +679,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) + 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.") @@ -637,34 +721,45 @@ get_num_majumps(nothing) == 0 @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} -# 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 +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) + +which is called by [`fill_scaled_rates!`](@ref) during aggregator initialization +and reinitialization to populate the working rate vector from parameters. -# Note this is unscaled -function (ratemap::MassActionJumpParamMapper{U})(params) where {U <: Int} - params[ratemap.param_idxs] +`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 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]] +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 -function to_collection(ratemap::MassActionJumpParamMapper{Int}) - MassActionJumpParamMapper([ratemap.param_idxs]) +function to_collection(ratemap::MassActionJumpParamMapper) + MassActionJumpParamMapper(copy(ratemap.param_idxs)) end function Base.merge!(pmap1::MassActionJumpParamMapper{U}, @@ -672,36 +767,28 @@ function Base.merge!(pmap1::MassActionJumpParamMapper{U}, 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 - """ -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) @@ -712,22 +799,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) @@ -737,13 +826,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} @@ -753,6 +850,16 @@ function JumpSet(vj::AbstractJump, cj::AbstractJump, rj::AbstractJump, 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 +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) @@ -775,6 +882,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(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 + 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.") @@ -849,10 +961,13 @@ 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 -function setup_majump_to_merge(sr::T, rs::AbstractVector{S}, ns::AbstractVector{U}, pmapper, +# 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 @@ -864,7 +979,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], @@ -930,9 +1055,26 @@ 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 uses_custom_mapper(maj) + using_params(maj) && !(maj.param_mapper isa MassActionJumpParamMapper) +end + function massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) + # 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, @@ -981,8 +1123,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), @@ -1012,8 +1155,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/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 008dab179..ce4e053eb 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 @@ -51,10 +51,9 @@ $(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. **NOTE**, - this only controls saving for non-`VariableRateJump`s. `VariableRateJump` saving is controlled + rates, specifies whether to save the system's state (before, after) the jump occurs. **NOTE**, + this only controls saving for non-`VariableRateJump`s. `VariableRateJump` saving is controlled at the individual jump level via the value of this kwarg passed to the individual `VariableRateJump`s constructor when it is created. - `spatial_system`, for spatial problems the underlying spatial structure. @@ -64,54 +63,158 @@ $(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`. + - `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 the DifferentialEquations.jl [docs](https://docs.sciml.ai/JumpProcesses/stable/) for usage examples and commonly asked questions. + +!!! 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, R, 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.""" +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. + """ 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 - """The random number generator to use.""" - rng::R - """kwargs to pass on to solve call.""" + """ + kwargs to pass on to solve call. + """ 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, - 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} + check_no_wrapped_callbacks(p) 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 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`. + +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 +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) @@ -132,9 +235,11 @@ 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) + # 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}) @@ -145,17 +250,15 @@ 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...) + 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 - # 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.") @@ -164,24 +267,11 @@ 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, - jprob.constant_jumps, jprob.variable_jumps, jprob.regular_jump, - jprob.massaction_jump, jprob.rng, 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 + jprob.constant_jumps, jprob.variable_jumps, jprob.regular_jump, + jprob.massaction_jump, jprob.kwargs) end DiffEqBase.isinplace(::JumpProblem{iip}) where {iip} = iip @@ -242,21 +332,24 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS vr_aggregator::VariableRateAggregator = VR_FRM(), save_positions = prob isa SciMLBase.AbstractDiscreteProblem ? (false, true) : (true, true), - rng = DEFAULT_RNG, scale_rates = true, useiszero = true, spatial_system = nothing, hopping_constants = nothing, callback = nothing, tstops = nothing, use_vrj_bounds = true, kwargs...) - - # 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 + 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 + 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 + maj = jumps.massaction_jump ## Spatial jumps handling if spatial_system !== nothing && hopping_constants !== nothing @@ -292,15 +385,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() @@ -313,55 +406,57 @@ function JumpProblem(prob, aggregator::AbstractAggregatorAlgorithm, jumps::JumpS 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 SciMLBase.AbstractDiscreteProblem ? (false, true) : (true, true), - rng = DEFAULT_RNG, 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 + 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 + 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) && 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 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) 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 -@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 4e25f6239..5faf26b2a 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. @@ -76,6 +77,9 @@ end SimpleExplicitTauLeaping(; epsilon = 0.05) = SimpleExplicitTauLeaping(epsilon) +SciMLBase.supports_solve_rng(::JumpProblem, ::SimpleTauLeaping) = true +SciMLBase.supports_solve_rng(::JumpProblem, ::SimpleExplicitTauLeaping) = true + """ SimpleImplicitTauLeaping(; epsilon = 0.05) @@ -111,19 +115,24 @@ $(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 SimpleImplicitTauLeaping(; epsilon = 0.05) = SimpleImplicitTauLeaping(epsilon) +SciMLBase.supports_solve_rng(::JumpProblem, ::SimpleImplicitTauLeaping) = true + """ SimpleTrapezoidalLeaping(; epsilon = 0.05) @@ -153,19 +162,24 @@ $(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 SimpleTrapezoidalLeaping(; epsilon = 0.05) = SimpleTrapezoidalLeaping(epsilon) +SciMLBase.supports_solve_rng(::JumpProblem, ::SimpleTrapezoidalLeaping) = true + """ $(TYPEDEF) @@ -201,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) @@ -211,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 @@ -231,6 +256,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 \ @@ -249,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(),...). \ @@ -273,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 @@ -283,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 @@ -301,13 +329,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) rj = jump_prob.regular_jump rate = rj.rate # rate function rate(out,u,p,t) @@ -348,7 +377,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 @@ -491,10 +520,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 @@ -568,27 +597,27 @@ 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) 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) + 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 @@ -627,7 +656,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) @@ -656,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 @@ -666,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) @@ -682,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 @@ -703,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 @@ -715,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 @@ -727,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 @@ -746,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] @@ -769,26 +798,26 @@ 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 - ) +) validate_pure_leaping_inputs(jump_prob, alg) || 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 @@ -824,7 +853,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 @@ -912,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 @@ -948,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 @@ -963,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 @@ -997,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] @@ -1020,26 +1049,26 @@ 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 - ) +) validate_pure_leaping_inputs(jump_prob, alg) || 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 @@ -1079,7 +1108,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, @@ -1110,6 +1139,33 @@ 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 + +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. + +## 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`, 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 - A `SciMLBase.EnsembleAlgorithm` value for use as the ensemble algorithm argument to diff --git a/src/solve.jl b/src/solve.jl index 8b17d80ac..b6a5553de 100644 --- a/src/solve.jl +++ b/src/solve.jl @@ -1,3 +1,24 @@ +""" + resolve_rng(rng, seed) + +Resolve which RNG to use for a jump simulation. + +Priority: `rng` > `seed` (creates `Xoshiro`) > `Random.default_rng()`. +""" +function resolve_rng(rng, seed) + if rng !== nothing + rng + elseif seed !== nothing + Random.Xoshiro(seed) + else + Random.default_rng() + end +end + +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; merge_callbacks = true, kwargs...) where {IIP, P} @@ -9,6 +30,7 @@ function SciMLBase.__solve(jump_prob::JumpProblem{IIP, P}, integrator.sol end +#Ambiguity Fix function SciMLBase.__solve(jump_prob::JumpProblem{IIP, P}, alg::Union{SciMLBase.AbstractRODEAlgorithm, SciMLBase.AbstractSDEAlgorithm}; merge_callbacks = true, kwargs...) where {IIP, P} @@ -20,6 +42,10 @@ function SciMLBase.__solve(jump_prob::JumpProblem{IIP, P}, integrator.sol end +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 function SciMLBase.__solve(jump_prob::JumpProblem{IIP, P}; @@ -31,60 +57,66 @@ 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...) - + alg::SciMLBase.AbstractDEAlgorithm; merge_callbacks = true, kwargs...) where { + IIP, P} __jump_init(_jump_prob, alg; kwargs...) -end - -function __jump_init(_jump_prob::SciMLBase.AbstractJumpProblem{P}, alg; - callback = nothing, seed = nothing, - alias_jump = Threads.threadid() == 1, - kwargs...) where {P} - if alias_jump - jump_prob = _jump_prob - reset_jump_problem!(jump_prob, seed) - else - jump_prob = resetted_jump_problem(_jump_prob, seed) - end +end - # DDEProblems do not have a recompile_flag argument - if jump_prob.prob isa SciMLBase.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 +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 -# 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). -# Truncate salt to native UInt so hash(::UInt64, ::UInt) matches on 32-bit Julia. -const _JUMP_SEED_SALT = 0x4a756d7050726f63 % UInt # "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)) +# 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 - jump_prob + nothing 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 +# 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; 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 + +function reset_jump_problem!(jump_prob, seed = nothing) + nothing end diff --git a/src/spatial/directcrdirect.jl b/src/spatial/directcrdirect.jl index 627bc2c8a..657df3058 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 @@ -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..e412fe5a0 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) @@ -49,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/nsm.jl b/src/spatial/nsm.jl index 0daa3fbd8..48993a7b4 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,9 +125,11 @@ 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 + fill_scaled_rates!(rx_rates.maj_rates, rx_rates.ma_jumps, integrator.p) reset!(rx_rates) reset!(hop_rates) @@ -153,30 +155,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/reaction_rates.jl b/src/spatial/reaction_rates.jl index 737cc5c98..b4854f44c 100644 --- a/src/spatial/reaction_rates.jl +++ b/src/spatial/reaction_rates.jl @@ -4,14 +4,25 @@ 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 + """ + maj_rates::Vector{F} end """ @@ -22,7 +33,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 +67,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 +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) 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 ebba52acd..77c18009f 100644 --- a/src/spatial/spatial_massaction_jump.jl +++ b/src/spatial/spatial_massaction_jump.jl @@ -152,6 +152,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, @@ -161,6 +166,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/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 370e86015..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 @@ -23,18 +29,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 3d73faafc..11f0d7d0e 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::SciMLBase.AbstractDiscreteProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::SciMLBase.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::SciMLBase.AbstractODEProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::SciMLBase.AbstractODEProblem, jumps) _f = SciMLBase.unwrapped_f(prob.f) if isinplace(prob) @@ -97,13 +96,13 @@ function extend_problem(prob::SciMLBase.AbstractODEProblem, jumps; rng = DEFAULT 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::SciMLBase.AbstractSDEProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::SciMLBase.AbstractSDEProblem, jumps) _f = SciMLBase.unwrapped_f(prob.f) if isinplace(prob) @@ -133,13 +132,13 @@ function extend_problem(prob::SciMLBase.AbstractSDEProblem, jumps; rng = DEFAULT 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::SciMLBase.AbstractDDEProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::SciMLBase.AbstractDDEProblem, jumps) _f = SciMLBase.unwrapped_f(prob.f) if isinplace(prob) @@ -159,14 +158,14 @@ function extend_problem(prob::SciMLBase.AbstractDDEProblem, jumps; rng = DEFAULT 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::SciMLBase.AbstractDAEProblem, jumps; rng = DEFAULT_RNG) +function extend_problem(prob::SciMLBase.AbstractDAEProblem, jumps) _f = SciMLBase.unwrapped_f(prob.f) if isinplace(prob) @@ -186,16 +185,15 @@ function extend_problem(prob::SciMLBase.AbstractDAEProblem, jumps; rng = DEFAULT 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,19 +202,21 @@ 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)) derivative_discontinuity!(integrator, true) 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, @@ -227,15 +227,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) + 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) @@ -267,21 +267,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()) @@ -333,21 +333,19 @@ sol = solve(jprob, Tsit5()) """ 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 @@ -355,14 +353,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 @@ -372,16 +369,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 @@ -399,8 +397,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 @@ -428,16 +426,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 @@ -465,8 +463,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 @@ -511,8 +508,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)) @@ -539,7 +536,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/allocations.jl b/test/allocations.jl index 250bb5c19..e76112efa 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 @@ -124,8 +124,7 @@ let jumps = [VariableRateJump((u, p, t) -> p[1] * (1 + u[i]) * t, integ -> (integ.u[i] += 1; nothing)) for i in 1:n] oprob = ODEProblem(f!, zeros(n), (0.0, 1.0), (0.5,)) - jprob = JumpProblem(oprob, Direct(), jumps...; vr_aggregator = VR_Direct(), - rng = StableRNG(1)) + jprob = JumpProblem(oprob, Direct(), jumps...; vr_aggregator = VR_Direct()) cache = jprob.jump_callback.continuous_callbacks[1].condition u, p, t = oprob.u0, oprob.p, 0.3 @test JumpProcesses.total_variable_rate(cache, u, p, t) ≈ n * 0.5 * 0.3 @@ -143,10 +142,9 @@ let for agg in (Direct(), FRM()) nallocs = map((10.0, 100.0)) do T dprob = DiscreteProblem(zeros(Int, n), (0.0, T), (0.5,)) - jprob = JumpProblem(dprob, agg, jumps...; save_positions = (false, false), - rng = StableRNG(1)) - solve(jprob, SSAStepper()) - @allocations solve(jprob, SSAStepper()) + jprob = JumpProblem(dprob, agg, jumps...; save_positions = (false, false)) + solve(jprob, SSAStepper(); rng = StableRNG(1)) + @allocations solve(jprob, SSAStepper(); rng = StableRNG(1)) end @test nallocs[1] == nallocs[2] end @@ -160,10 +158,8 @@ let net_stoich = [[1 => 1], [1 => -1]] maj = MassActionJump(reactant_stoich, net_stoich; param_idxs = rateidxs) dprob = DiscreteProblem([10], (0.0, 1.0), [1.0, 0.1]) - jprob = JumpProblem( - dprob, RDirect(), maj; save_positions = (false, false), - rng = StableRNG(1), counter_threshold = Int32(2) - ) + jprob = JumpProblem(dprob, RDirect(), maj; save_positions = (false, false), + counter_threshold = Int32(2)) agg = jprob.discrete_jump_aggregation @test fieldtype(typeof(agg), :counter_threshold) === Int @test agg.counter_threshold isa Int diff --git a/test/bimolerx_test.jl b/test/bimolerx_test.jl index d855b7f58..c1d725c2a 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 + solve(jump_prob, SSAStepper(); rng) + else + 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 45d7b4041..cec65800f 100644 --- a/test/bracketing.jl +++ b/test/bracketing.jl @@ -43,10 +43,12 @@ 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] == +@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)[2] == +@test JP.get_majump_brackets(ulow, uhigh, reaction_index, majump, maj_rates_work)[2] == majump_rates[1] * uhigh[1] # high # constant rate @@ -61,13 +63,14 @@ ucur = [5] ### Aggregator ### mutable struct DummyAggregator{T, M, R, CB, 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} cur_rate_high::Vector{T} sum_rate::T ma_jumps::M + maj_rates::Vector{T} rates::R brackets::CB bracket_data::BD @@ -78,7 +81,7 @@ cur_rate_high = [0.0, 0.0] sum_rate = 0.0 brackets = JP.get_jump_bracket_fwrappers(ulow, params, t, (crj,), RSSA()) p = DummyAggregator([0], [0], cur_rate_low, cur_rate_high, sum_rate, majump, - [rate], brackets, bd) + maj_rates_work, [rate], brackets, bd) u = [100] JP.update_u_brackets!(p, u) @@ -95,7 +98,7 @@ reaction_index = 2 @test JP.get_jump_brackets(reaction_index, p, u, params, t)[2] == rate(p.ulow, params, t) p = DummyAggregator([0], [0], cur_rate_low, cur_rate_high, sum_rate, majump, - [rate], brackets, bd) + maj_rates_work, [rate], brackets, 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 @@ -158,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/callbacks.jl b/test/callbacks.jl index 8f70bffeb..af5c20864 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) @@ -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,20 +457,20 @@ 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 @testset "SDE + jump callback not duplicated" begin @@ -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,10 +508,180 @@ 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) @test n_continuous == 1 end + +# 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[] + 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) + 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(), (;)), + ("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) + end + + @testset "Direct __init discards merge_callbacks" begin + 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 not merge them again. + 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 + +# 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 diff --git a/test/constant_rate.jl b/test/constant_rate.jl index 833d6b981..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) @@ -16,37 +17,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/degenerate_rx_cases.jl b/test/degenerate_rx_cases.jl index 4658562bb..c0eb89274 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 @@ -97,18 +97,20 @@ 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; 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))) 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 55d7127be..d04c81e0f 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 @@ -44,7 +44,7 @@ first_jump_time(traj) = traj.t[findfirst(>(traj.t[1]), traj.t)] @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 @@ -52,7 +52,7 @@ first_jump_time(traj) = traj.t[findfirst(>(traj.t[1]), traj.t)] @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] @@ -77,13 +77,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] @@ -92,17 +94,46 @@ end 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. +# 3. rng kwarg reproducibility: same rng seed → identical trajectory, +# different rng seeds → different trajectories # ========================================================================== +@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 + + @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 + + @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)] + 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 + @testset "EnsembleThreads: no data race" begin @testset "SSAStepper" begin jprob = make_ssa_jump_prob() @@ -113,8 +144,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.u) == 4 @@ -122,8 +151,6 @@ 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. sol = solve(EnsembleProblem(jprob), EM(), EnsembleThreads(); trajectories = 4, dt = 0.01, save_everystep = false) @test length(sol.u) == 4 @@ -133,53 +160,23 @@ 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. has_rng / get_rng / set_rng! interface on SSAIntegrator # ========================================================================== -@testset "resetted_jump_problem: different seeds → different streams" begin +@testset "SSAIntegrator RNG interface" begin jprob = make_ssa_jump_prob() - seeds = UInt64[100, 200, 300] + integrator = init(jprob, SSAStepper(); rng = StableRNG(42)) - # 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 - 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 + @test SciMLBase.has_rng(integrator) + rng = SciMLBase.get_rng(integrator) + @test rng isa StableRNG -@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) - end - @test allunique(draws) -end + new_rng = StableRNG(99) + SciMLBase.set_rng!(integrator, new_rng) + @test SciMLBase.get_rng(integrator) === new_rng -@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)) + # mismatched RNG type should throw + @test_throws ArgumentError SciMLBase.set_rng!(integrator, Random.Xoshiro(123)) end # ========================================================================== @@ -187,13 +184,43 @@ 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. +# 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 + # 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) + trajectories = 3, rng = StableRNG(12345)) event_times = [first_jump_time(sol.u[i]) for i in 1:3] @test allunique(event_times) + init_thresholds = [sol.u[i].u[2].jump_u[1] for i in 1:3] + @test allunique(init_thresholds) +end + +# ========================================================================== +# 6. JumpProblem rng kwarg throws ArgumentError +# ========================================================================== + +@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)) + 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 72fb853ab..0390e8bca 100644 --- a/test/ensemble_uniqueness.jl +++ b/test/ensemble_uniqueness.jl @@ -7,36 +7,20 @@ 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, ctx) - 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 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 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 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 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 eltype(sol.u[1].u[1]) == Int diff --git a/test/explicit_tau_kernel_tests.jl b/test/explicit_tau_kernel_tests.jl index 8e0904416..b7b8d8a5e 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 @@ -26,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) @@ -159,5 +161,9 @@ 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) + run_kernel_alias_tests(backend, SimpleExplicitTauLeaping(), rng_prob; saveat = 1.0) return nothing end diff --git a/test/extended_jump_array.jl b/test/extended_jump_array.jl index 861f4c763..3e8d95c3c 100644 --- a/test/extended_jump_array.jl +++ b/test/extended_jump_array.jl @@ -211,9 +211,9 @@ 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 @@ -255,9 +255,8 @@ end u0 = ArrayPartition([1.0], [1.0]) ode = ODEProblem(f!, u0, (0.0, 1.0)) jump = VariableRateJump((u, p, t) -> 10 * u[1], affect!) - jump_prob = JumpProblem(ode, Direct(), jump; vr_aggregator = agg, - rng = StableRNG(57)) - sol = solve(jump_prob, Tsit5(); abstol = 1e-11, reltol = 1e-11) + jump_prob = JumpProblem(ode, Direct(), jump; vr_aggregator = agg) + sol = solve(jump_prob, Tsit5(); abstol = 1e-11, reltol = 1e-11, rng = StableRNG(57)) @test sol.retcode == ReturnCode.Success @test jumps[] > 0 @test isapprox(sol.u[end].x[1][1], exp(-1); rtol = 1e-9) diff --git a/test/extended_jump_array_remake.jl b/test/extended_jump_array_remake.jl index 7f2168d5e..c64325ebe 100644 --- a/test/extended_jump_array_remake.jl +++ b/test/extended_jump_array_remake.jl @@ -1,12 +1,11 @@ # 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, StochasticDiffEq, SciMLBase, 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)) @@ -17,20 +16,25 @@ 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) + # 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 should be resampled (different from original) - @test prob2.prob.u0.jump_u != original_jump_u + @test all(iszero, prob2.prob.u0.jump_u) + # 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 @@ -47,34 +51,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 + @test all(iszero, prob2.prob.u0.jump_u) + # 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 - 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 + @test all(iszero, prob2.prob.u0.jump_u) + 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 - 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 + @test all(iszero, prob2.prob.u0.jump_u) + 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 @@ -88,25 +92,105 @@ 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 + @test all(iszero, prob2.prob.u0.jump_u) + 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 + +# `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 diff --git a/test/extinction_test.jl b/test/extinction_test.jl index afd077951..9f3b6ba6b 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 06b4cf166..041a958ea 100644 --- a/test/fp_unknowns.jl +++ b/test/fp_unknowns.jl @@ -34,11 +34,10 @@ 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, - vartojumps_map = vtoj, jumptovars_map = jtov, dep_graph = dg, - scale_rates = false) + jprob = JumpProblem(dprob, agg, maj; save_positions = (false, false), + vartojumps_map = vtoj, jumptovars_map = jtov, dep_graph = dg) 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 9a44031d3..e110e35b2 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 fafbec33b..497c724f5 100644 --- a/test/geneexpr_test.jl +++ b/test/geneexpr_test.jl @@ -22,19 +22,23 @@ 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 + solve(jump_prob, SSAStepper(); rng) + else + 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 = init(vrjprob, Tsit5(); saveat = tsave, rng) solve!(integrator) Psamp[1] = integrator.sol[3, end] for i in 2:Nsims @@ -94,8 +98,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 +110,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 +122,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 +132,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 +143,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 +191,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..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 @@ -72,7 +69,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 1fc1a16a1..a468581ca 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,13 @@ 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/jprob_symbol_indexing.jl b/test/jprob_symbol_indexing.jl index 45ec4345c..6b071b979 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/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 new file mode 100644 index 000000000..eb437eab6 --- /dev/null +++ b/test/kernel_contract_tests.jl @@ -0,0 +1,208 @@ +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_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) + + # 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. + 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) + + @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]) + 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 + + 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) + run_kernel_alias_tests(backend, SimpleTauLeaping(), jp; dt = 0.25) +end diff --git a/test/linearreaction_test.jl b/test/linearreaction_test.jl index 37d43bf30..4601b6459 100644 --- a/test/linearreaction_test.jl +++ b/test/linearreaction_test.jl @@ -3,7 +3,7 @@ using DiffEqBase, JumpProcesses, SciMLBase, Statistics using Test using StableRNGs -rng = StableRNG(12345) +const rng = StableRNG(12345) # using BenchmarkTools # dobenchmark = true @@ -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 69b039096..052d304d9 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/massaction_state.jl b/test/massaction_state.jl new file mode 100644 index 000000000..61cb3da8f --- /dev/null +++ b/test/massaction_state.jl @@ -0,0 +1,222 @@ +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) + @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) + @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) + 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) + @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) + @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 + +# 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) + 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/monte_carlo_test.jl b/test/monte_carlo_test.jl index 42566b3ad..2141013c8 100644 --- a/test/monte_carlo_test.jl +++ b/test/monte_carlo_test.jl @@ -8,28 +8,28 @@ 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) first_event(traj) = traj.t[findfirst(>(traj.t[1]), traj.t)] @test first_event(sol.u[1]) != first_event(sol.u[2]) != first_event(sol.u[3]) -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/precompile_workload.jl b/test/precompile_workload.jl index 91bb0bb05..28811ff77 100644 --- a/test/precompile_workload.jl +++ b/test/precompile_workload.jl @@ -1,18 +1,28 @@ using JumpProcesses, DiffEqBase using Random +using SciMLBase: get_rng using Test rate = (u, p, t) -> u[1] affect! = integrator -> (integrator.u[1] += 1) jump = ConstantRateJump(rate, affect!) prob = DiscreteProblem([10.0], (0.0, 1.0)) -jump_prob = JumpProblem(prob, Direct(), jump; - rng = MersenneTwister(12345), save_positions = (false, false)) +jump_prob = JumpProblem(prob, Direct(), jump; save_positions = (false, false)) -integrator = init(jump_prob, SSAStepper()) -step!(integrator) -@test integrator.u[1] == 11.0 +@testset "Precompile workload seed=$seed" for seed in (nothing, 12345) + kwargs = seed === nothing ? (;) : (; seed) + integrator = init(jump_prob, SSAStepper(); kwargs...) + if seed === nothing + @test get_rng(integrator) isa Random.TaskLocalRNG + else + @test get_rng(integrator) isa Random.Xoshiro + end + step!(integrator) + @test integrator.t > 0.0 + @test integrator.u[1] >= 10.0 -sol = solve(jump_prob, SSAStepper()) -@test sol.t[end] == 1.0 -@test sol.u[end][1] > 10.0 + sol = solve(jump_prob, SSAStepper(); kwargs...) + @test sol.retcode == ReturnCode.Success + @test sol.t[end] == 1.0 + @test sol.u[end][1] >= 10.0 +end 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 6f25e5c2b..faba3a5c9 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 @@ -10,14 +10,14 @@ 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; 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 @@ -56,16 +190,17 @@ 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) 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 @@ -80,45 +215,47 @@ 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) + 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 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) + 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) 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) @@ -152,16 +289,17 @@ 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) 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 @@ -177,45 +315,47 @@ 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) + 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 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) + 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) 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) @@ -238,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) @@ -247,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 @@ -263,13 +403,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 +424,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 +437,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 @@ -307,84 +447,86 @@ 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); 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 @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); 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 - + # 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); 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 - + # 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); 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 - + # 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; 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 @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); rng, - spatial_system = spatial_sys) - @test_throws ErrorException JumpProblem(prob, PureLeaping(), JumpSet(maj); rng, - hopping_constants = hopping_consts) - + @test_throws ErrorException JumpProblem(prob, PureLeaping(), JumpSet(maj); + spatial_system = spatial_sys) + @test_throws ErrorException JumpProblem(prob, PureLeaping(), JumpSet(maj); + hopping_constants = hopping_consts) + # Test MassActionJump with parameter mapping maj_params = MassActionJump(reactant_stoich, net_stoich; param_idxs = [1, 2]) - jp_params = JumpProblem(prob, PureLeaping(), JumpSet(maj_params); rng) - scaled_rates = [p[1], p[2]/2] - @test jp_params.massaction_jump.scaled_rates == scaled_rates + jp_params = JumpProblem(prob, PureLeaping(), JumpSet(maj_params)) + @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 @@ -398,23 +540,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 @@ -466,7 +608,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/remake_test.jl b/test/remake_test.jl index 7952512e1..0fd6fbfac 100644 --- a/test/remake_test.jl +++ b/test/remake_test.jl @@ -1,4 +1,4 @@ -using JumpProcesses, DiffEqBase, OrdinaryDiffEq, SciMLBase +using JumpProcesses, DiffEqBase, OrdinaryDiffEq, SciMLBase, 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,20 +74,20 @@ 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, :] t = sol.t first_nontstart = findfirst(>(t[1]), t) @@ -97,7 +96,7 @@ let 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 @@ -109,24 +108,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, :] t = sol.t first_nontstart = findfirst(>(t[1]), t) diff --git a/test/reversible_binding.jl b/test/reversible_binding.jl index da74e6252..a781d01b2 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 = 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 new file mode 100644 index 000000000..11f91b82f --- /dev/null +++ b/test/rng_kwarg_tests.jl @@ -0,0 +1,327 @@ +using JumpProcesses, OrdinaryDiffEq, StochasticDiffEq, SciMLBase, Test +using StableRNGs, Random + +# ========================================================================== +# 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() + 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) +end + +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) +end + +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) +end + +# Helpers +first_jump_time(traj) = traj.t[findfirst(>(traj.t[1]), traj.t)] + +# ========================================================================== +# 1. JumpProblem(; rng=...) throws ArgumentError +# ========================================================================== +@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/init +# ========================================================================== +@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 solve rng +# ========================================================================== +@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: different seeds → different trajectories +# ========================================================================== +@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 solve/init +# ========================================================================== +@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: reproducibility via solve rng +# ========================================================================== +@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 + +# ========================================================================== +# 7. ODE + VR: different seeds → different trajectories +# ========================================================================== +@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 = [first_jump_time(s) for s in sols] + @test allunique(times) +end + +# ========================================================================== +# 8. SDE + VR: rng via solve/init +# ========================================================================== +@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 + +# ========================================================================== +# 9. SDE + VR: reproducibility via solve rng +# ========================================================================== +@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 + +# ========================================================================== +# 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) + 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, rng = StableRNG(42)) + sol2 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(42)) + @test sol1.u == sol2.u +end + +# ========================================================================== +# 12. SimpleTauLeaping: different seeds → different trajectories +# ========================================================================== +@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)) + jprob = JumpProblem(dprob, PureLeaping(), rj) + sol1 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(42)) + sol2 = solve(jprob, SimpleTauLeaping(); dt = 0.01, rng = StableRNG(99)) + @test sol1.u != sol2.u +end + +# ========================================================================== +# 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() + 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 + + # mismatched RNG type should throw + @test_throws ArgumentError SciMLBase.set_rng!(integrator, Random.Xoshiro(123)) +end + +# ========================================================================== +# 16. 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 + + @testset "SDE + VR" begin + jprob = make_sde_vr_jump_prob() + sol = solve(jprob, EM(); dt = 0.01) + @test sol.retcode == ReturnCode.Success + end +end + +# ========================================================================== +# 17. seed kwarg: creates Xoshiro from integer seed +# ========================================================================== +@testset "seed kwarg creates Xoshiro" begin + @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 + +# ========================================================================== +# 18. 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 + +# ========================================================================== +# 19. 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 + +# ========================================================================== +# 20. 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/runtests.jl b/test/runtests.jl index c9f655b9e..dea382bb0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -23,6 +23,8 @@ end @time @safetestset "Monte Carlo Tests" begin include("monte_carlo_test.jl") 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 @@ -33,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 @@ -40,9 +43,12 @@ 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 + @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 "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 @@ -61,6 +67,11 @@ end @time @safetestset "Explicit tau-leaping kernel on the CPU backend" begin include("kernelabstractions_explicit_tau.jl") end 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" activate_gpu_env() @time @safetestset "GPU Tau Leaping test" begin include("gpu/regular_jumps.jl") end diff --git a/test/save_positions.jl b/test/save_positions.jl index d23e209b8..caf6d43ff 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 d05cce8e7..a31674606 100644 --- a/test/saveat_regression.jl +++ b/test/saveat_regression.jl @@ -11,12 +11,12 @@ using Random, StableRNGs n0 = [1, 1, 0] tspan = (0, 0.2) dprob = DiscreteProblem(n0, tspan) - jprob = JumpProblem(dprob, Direct(), maj; save_positions = (false, false), 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) + trajectories = Nsims, rng) for i in eachindex(sol.u) NA .+= sol.u[i][1, :] @@ -27,10 +27,10 @@ using Random, StableRNGs end NA = zeros(length(ts)) - jprob = JumpProblem(dprob, Direct(), maj; 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 eachindex(ts) @@ -74,19 +74,21 @@ end end prob = DiscreteProblem([0], (0.0, 1.0)) jump = ConstantRateJump(rate, affect!) - jprob = JumpProblem(prob, Direct(), jump; save_positions, rng) - # 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, alias_jump = true) + jprob = JumpProblem(prob, Direct(), jump; save_positions) + # 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 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) @@ -154,3 +158,74 @@ 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 + +# 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 diff --git a/test/scale_rates_field_test.jl b/test/scale_rates_field_test.jl index 4b84840cd..41ff1b0eb 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 @@ -76,27 +80,34 @@ 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 - @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,12 +124,11 @@ 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 -# 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]] @@ -149,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 @@ -161,3 +171,44 @@ end # Mismatched rescale_rates_on_update via JumpProblem varargs — should error @test_throws ErrorException JumpProblem(dprob, Direct(), maj_true, maj_false) end + +# 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]) + 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]) + 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 + 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 +@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/sir_model.jl b/test/sir_model.jl index 40bf4ac98..af6cac618 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/sortingdirect_reuse.jl b/test/sortingdirect_reuse.jl new file mode 100644 index 000000000..ff3d6551d --- /dev/null +++ b/test/sortingdirect_reuse.jl @@ -0,0 +1,57 @@ +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) + @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) + 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) + @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/spatial/ABC.jl b/test/spatial/ABC.jl index 8f7b7c89a..ee8027d3b 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/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/spatial/diffusion.jl b/test/spatial/diffusion.jl index 6407c8be4..50da107a8 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/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 8199dd39e..4dadca11a 100644 --- a/test/spatial/reaction_rates.jl +++ b/test/spatial/reaction_rates.jl @@ -36,10 +36,15 @@ 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 diff --git a/test/spatial/spatial_majump.jl b/test/spatial/spatial_majump.jl index b27b53034..f21364b50 100644 --- a/test/spatial/spatial_majump.jl +++ b/test/spatial/spatial_majump.jl @@ -67,26 +67,26 @@ empty_spatial_majump_with_mapper = SpatialMassActionJump(nothing, reactstoch, ne 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 @@ -112,8 +112,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) @@ -128,8 +128,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) @@ -141,8 +141,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) @@ -156,8 +156,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 2568f94f4..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) @@ -12,16 +13,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 +42,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 +76,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 +92,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_aggregator_smoke.jl b/test/ssa_aggregator_smoke.jl new file mode 100644 index 000000000..49952e5d2 --- /dev/null +++ b/test/ssa_aggregator_smoke.jl @@ -0,0 +1,102 @@ +using JumpProcesses, SciMLBase, Test + +# Small trajectory checks catch initialization and RNG plumbing failures without +# the ensembles used by the statistical correctness tests. +function check_ssa_smoke_solution(sol, prob) + @test SciMLBase.successful_retcode(sol) + @test first(sol.t) == first(prob.tspan) + @test last(sol.t) == last(prob.tspan) + @test issorted(sol.t) + @test all(u -> all(>=(0), u), sol.u) + @test all(u -> sum(u) == sum(prob.u0), sol.u) +end + +function check_ssa_smoke_problem(jprob) + # 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) + + sol_seeded = solve(deepcopy(jprob), SSAStepper(); seed = 12345) + check_ssa_smoke_solution(sol_seeded, jprob.prob) + sol_repeated = solve(deepcopy(jprob), SSAStepper(); seed = 12345) + @test sol_seeded.t == sol_repeated.t + @test sol_seeded.u == sol_repeated.u +end + +function ssa_smoke_problem(aggregator, representation) + # A <-> B: equivalent mass-action and constant-rate representations. + p = [0.7, 0.4] + prob = DiscreteProblem([12, 8], (0.0, 0.5), p) + reactant_stoch = [[1 => 1], [2 => 1]] + net_stoch = [[1 => -1, 2 => 1], [1 => 1, 2 => -1]] + + if representation === :parameterized_massaction + maj = MassActionJump(reactant_stoch, net_stoch; param_idxs = [1, 2]) + return JumpProblem(prob, aggregator, maj) + elseif representation === :fixed_massaction + maj = MassActionJump(p, reactant_stoch, net_stoch) + return JumpProblem(prob, aggregator, maj) + end + + rate_ab(u, p, t) = p[1] * u[1] + 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_b_to_a!(integrator) + integrator.u[1] += 1 + integrator.u[2] -= 1 + nothing + end + 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]] + JumpProblem(prob, aggregator, jumps...; dep_graph, vartojumps_map, jumptovars_map) +end + +@testset "Nonspatial SSA aggregator smoke" begin + @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 + + @testset "$(typeof(aggregator)): fixed mass action" for aggregator in (RSSA(), RSSACR()) + check_ssa_smoke_problem(ssa_smoke_problem(aggregator, :fixed_massaction)) + end +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)) + 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] + net_stoch = [[1 => -1, 2 => 1] for _ in 1:num_reactions] + maj = MassActionJump(reactant_stoch, net_stoch; + param_idxs = collect(1:num_reactions)) + jprob = JumpProblem(prob, JumpSet(maj)) + @test jprob.aggregator isa expected_aggregator + check_ssa_smoke_problem(jprob) + end +end + +@testset "Spatial SSA aggregator smoke" begin + @testset "$(typeof(aggregator))" for aggregator in (NSM(), DirectCRDirect()) + prob = DiscreteProblem([12 0; 0 8], (0.0, 0.5), [0.7, 0.4]) + maj = MassActionJump([[1 => 1], [2 => 1]], + [[1 => -1, 2 => 1], [1 => 1, 2 => -1]]; param_idxs = [1, 2]) + spatial_system = CartesianGrid((2,)) + hopping_constants = [0.2, 0.2] + jprob = JumpProblem(prob, aggregator, maj; spatial_system, hopping_constants) + check_ssa_smoke_problem(jprob) + end +end diff --git a/test/ssa_callback_test.jl b/test/ssa_callback_test.jl index 5852617f4..535e028ad 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,41 @@ 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,30 +104,36 @@ 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 -# 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), rng = rng) -@test all(jprob.massaction_jump.scaled_rates .== [0.5]) -jprob = JumpProblem(dprob, Direct(), maj5, save_positions = (false, false), rng = rng, scale_rates = false) -@test all(jprob.massaction_jump.scaled_rates .== [1.0]) +jprob = JumpProblem(dprob, Direct(), maj5, save_positions = (false, false)) +@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]]; 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 +144,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 +167,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 +194,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 +214,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,46 +223,41 @@ 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 end -# test that reset_aggregated_jumps! with update_jump_params kwarg dispatches correctly -# for SSAIntegrator (https://github.com/SciML/JumpProcesses.jl/issues/562) -let - rate1(u, p, t) = p[1] * u[1] * u[2] - affect1!(integrator) = (integrator.u[1] -= 1; integrator.u[2] += 1) - jump1 = ConstantRateJump(rate1, affect1!) - - rate2(u, p, t) = p[2] * u[2] - affect2!(integrator) = (integrator.u[2] -= 1; integrator.u[3] += 1) - jump2 = ConstantRateJump(rate2, affect2!) - - p = (1.0, 1.0) - prob = DiscreteProblem([990, 0, 0], (0.0, 250.0), p) - jump_prob = JumpProblem(prob, Direct(), jump1, jump2; rng) - - # with update_jump_params = true, should behave same as without kwarg - int = init(jump_prob, SSAStepper()) - int[2] = 10 - reset_aggregated_jumps!(int; update_jump_params = true) - step!(int, 1000.0, true) - @test int.u[3] > 0 # at least some recovered, confirming jumps fired +# 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 - # with update_jump_params = false, should also reset jump aggregation - int2 = init(jump_prob, SSAStepper()) - int2[2] = 10 - reset_aggregated_jumps!(int2; update_jump_params = false) - step!(int2, 1000.0, true) - @test int2.u[3] > 0 +# 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 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 diff --git a/test/ssa_kernel_tests.jl b/test/ssa_kernel_tests.jl index 81328c0e1..c0a46d3f5 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. @@ -60,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) @@ -69,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) @@ -77,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) @@ -88,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) @@ -114,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) @@ -197,5 +203,10 @@ 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_kernel_alias_tests(backend, SSAStepper(), rng_prob; saveat = 1.0) + run_regular_kernel_rng_tests(backend) return nothing end 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/stochastic_dispatch.jl b/test/stochastic_dispatch.jl new file mode 100644 index 000000000..e7bc1d5dd --- /dev/null +++ b/test/stochastic_dispatch.jl @@ -0,0 +1,151 @@ +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 + +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 + +const SDECore = StochasticDiffEq.StochasticDiffEqCore +const ODECORE_EXT = Base.get_extension(JumpProcesses, :JumpProcessesOrdinaryDiffEqCoreExt) + +# Whether a method is a `__init`, `__solve`, or `supports_solve_rng` method for +# `JumpProblem`s. +function involves_jump_init(m) + 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 + # 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))) + @test method.module === SDECore + @test method.sig == Tuple{typeof(SciMLBase.__init), JumpProblem, A} + end + + # 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 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 + +@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 native alias settings" begin + for (kind, alg) in ((:SDE, EM()), (:RODE, RandomEM())), + alias_mode in (:default, :alias, :copy, :specifier_alias, :specifier_copy) + + jprob = stochastic_jump_problem(kind) + constructor = kind === :SDE ? SciMLBase.SDEAliasSpecifier : + SciMLBase.RODEAliasSpecifier + 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_jumps + if alias_mode in (:copy, :specifier_alias, :specifier_copy) + @test integrator.u !== jprob.prob.u0 + end + solve!(integrator) + @test successful_retcode(integrator.sol) + end +end diff --git a/test/table_test.jl b/test/table_test.jl index 555ffaf1e..1cf01f4c1 100644 --- a/test/table_test.jl +++ b/test/table_test.jl @@ -1,5 +1,6 @@ using DiffEqBase, JumpProcesses using Test +using StableRNGs const DJ = JumpProcesses # test data @@ -53,11 +54,12 @@ priorities[10] = 0.0 @test pt.groups[1].numpids == 2 # test sampling +rng = StableRNG(12345) cnt = 0 Nsamps = Int(1e7) for i in 1:Nsamps global cnt - pid = DJ.sample(pt, priorities) + pid = DJ.sample(pt, priorities, rng) (pid == 8) && (cnt += 1) end @test abs(cnt // Nsamps - 0.008968535978248484) / 0.008968535978248484 < 0.05 diff --git a/test/test_groups.toml b/test/test_groups.toml index bfc13e0c8..abb9b127d 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 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"] group = "InterfaceI" @@ -8,7 +11,10 @@ os = ["ubuntu-latest"] arch = "x86" [InterfaceII] -versions = ["lts", "1", "pre"] +# 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] versions = ["lts", "1"] diff --git a/test/thread_safety.jl b/test/thread_safety.jl index 8ab5629e1..473c8ad4c 100644 --- a/test/thread_safety.jl +++ b/test/thread_safety.jl @@ -1,7 +1,6 @@ using DiffEqBase, SciMLBase, Test -using JumpProcesses, OrdinaryDiffEq +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]]) @@ -9,10 +8,54 @@ 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) -solve(EnsembleProblem(jprob; safetycopy = true), SSAStepper(), EnsembleThreads(); - trajectories = 10) +jprob = JumpProblem(dprob, Direct(), maj) + +# 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.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.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 @@ -26,11 +69,30 @@ let for agg in (VR_FRM(), VR_Direct(), VR_DirectFW()) jump_prob = JumpProblem(ode_prob, Direct(), vrj; vr_aggregator = agg) - prob_func(prob, ctx) = 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.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 + +# 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 = 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)] @test allunique(firstrx_time) end end diff --git a/test/variable_rate.jl b/test/variable_rate.jl index 86b54f052..6dfd996fc 100644 --- a/test/variable_rate.jl +++ b/test/variable_rate.jl @@ -30,11 +30,11 @@ 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 = 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) function vrfrm_trajectory(default_seed) task = @task begin @@ -45,9 +45,8 @@ function vrfrm_trajectory(default_seed) integrator -> (integrator.u[1] += 1)) jump2 = VariableRateJump((u, p, t) -> 10.0, integrator -> (integrator.u[1] += 1)) - jump_prob = JumpProblem(ode_prob, jump1, jump2; vr_aggregator = VR_FRM(), - rng = local_rng) - solve(jump_prob, Tsit5(); saveat = 0.0:0.05:1.0).u + jump_prob = JumpProblem(ode_prob, jump1, jump2; vr_aggregator = VR_FRM()) + solve(jump_prob, Tsit5(); saveat = 0.0:0.05:1.0, rng = local_rng).u end schedule(task) return fetch(task) @@ -55,9 +54,9 @@ end @test vrfrm_trajectory(123) == vrfrm_trajectory(456) -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_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.u)]) <= 1e-12 @@ -67,10 +66,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.u)]) <= 1e-12 @test maximum([sol.u[i][3] for i in 1:length(sol.u)]) <= 1e-12 @@ -93,10 +92,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 @@ -107,10 +106,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] @@ -118,10 +117,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] @@ -129,10 +128,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] @@ -148,10 +147,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) @@ -165,10 +164,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 @@ -177,7 +176,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 @@ -275,12 +274,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 @@ -290,9 +289,9 @@ end # accuracy test based on # https://github.com/SciML/JumpProcesses.jl/issues/320 -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] @@ -309,8 +308,7 @@ function getmean(Nsims, prob, alg, tsave, seed) end let - seed = 12345 - rng = StableRNG(seed) + rng_seed = 12345 b = 2.0 d = 1.0 n0 = 1.0 @@ -342,12 +340,12 @@ 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) + 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 @@ -357,10 +355,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 @@ -458,10 +456,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 @@ -516,8 +514,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] @@ -547,8 +545,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 @@ -579,8 +577,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 @@ -636,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