Document JumpProcesses public API

Project.toml

@@ -53,7 +53,7 @@ RecursiveArrayTools = "4.3.0, 4"
 Reexport = "1.2.2"
 SafeTestsets = "0.1"
 SciMLBase = "3.18.0, 3.1"
-SciMLTesting = "1.6"
+SciMLTesting = "2.1"
 StableRNGs = "1"
 StaticArrays = "1.9.18"
 Statistics = "1"

docs/src/api.md

@@ -7,8 +7,11 @@ CurrentModule = JumpProcesses
 ## Core Types
 
 ```@docs
+ExtendedJumpArray
 JumpProblem
+PureLeaping
 SSAStepper
+SplitCoupledJumpProblem
 reset_aggregated_jumps!
 ```
 
@@ -29,20 +32,50 @@ Aggregators are the underlying algorithms used for sampling
 [`VariableRateJump`](@ref)s.
 
 ```@docs
+BracketData
+CCNRM
 Coevolve
 Direct
 DirectCR
+DirectCRDirect
+DirectFW
 FRM
+FRMFW
 NRM
+NSM
 RDirect
 RSSA
 RSSACR
 SortingDirect
+get_num_majumps
+needs_depgraph
+needs_vartojumps_map
 ```
 
-# Private API Functions
+## Variable Rate Aggregators
 
 ```@docs
-ExtendedJumpArray
-SSAIntegrator
+VariableRateAggregator
+VR_Direct
+VR_DirectFW
+VR_FRM
+```
+
+## Tau-Leaping Algorithms
+
+```@docs
+EnsembleGPUKernel
+SimpleExplicitTauLeaping
+SimpleTauLeaping
+```
+
+## Spatial Jump APIs
+
+```@docs
+CartesianGrid
+CartesianGridRej
+SpatialMassActionJump
+neighbors
+num_sites
+outdegree
 ```

src/JumpProcesses.jl

@@ -113,8 +113,35 @@ include("variable_rate.jl")
 export VariableRateAggregator, VR_FRM, VR_Direct, VR_DirectFW
 
 """
-Aggregator to indicate that individual jumps should also be handled via the leaping
-algorithm that is passed to solve.
+    PureLeaping()
+
+Request that all jumps in a [`JumpProblem`](@ref) are handled by the leaping algorithm
+passed to `solve`, instead of being converted into callback-based SSA aggregators during
+problem construction.
+
+## Returns
+
+  - A stateless aggregator marker for `JumpProblem(prob, PureLeaping(), jumps; kwargs...)`.
+
+## Notes
+
+  - `PureLeaping` is currently intended for tau-leaping algorithms such as
+    [`SimpleTauLeaping`](@ref) and [`SimpleExplicitTauLeaping`](@ref).
+  - Spatial jump problems are not supported by the `PureLeaping` construction path.
+
+## Examples
+
+```julia
+using JumpProcesses, DiffEqBase
+
+rate!(out, u, p, t) = (out[1] = 0.2 * u[1])
+affect!(du, u, p, t, counts, mark) = (du[1] = -counts[1])
+rj = RegularJump(rate!, affect!, 1)
+
+prob = DiscreteProblem([10], (0.0, 1.0))
+jprob = JumpProblem(prob, PureLeaping(), rj)
+sol = solve(jprob, SimpleTauLeaping(); dt = 0.1)
+```
 """
 struct PureLeaping <: AbstractAggregatorAlgorithm end
 export PureLeaping

src/aggregators/aggregators.jl

@@ -179,6 +179,33 @@ const JUMP_AGGREGATORS = (Direct(), DirectFW(), DirectCR(), SortingDirect(), RSS
 struct NullAggregator <: AbstractAggregatorAlgorithm end
 
 # true if aggregator requires a jump dependency graph
+"""
+    needs_depgraph(aggregator) -> Bool
+
+Return whether an aggregator requires a reaction dependency graph when building a
+[`JumpProblem`](@ref).
+
+## Arguments
+
+  - `aggregator`: An `AbstractAggregatorAlgorithm` value such as [`DirectCR`](@ref),
+    [`SortingDirect`](@ref), [`NRM`](@ref), [`CCNRM`](@ref), [`RDirect`](@ref), or
+    [`Coevolve`](@ref).
+
+## Returns
+
+  - `true` when the aggregator needs a dependency graph from each reaction to the
+    reactions whose propensities must be updated after it fires.
+  - `false` for aggregators that update all rates or otherwise do not use this graph.
+
+## Examples
+
+```julia
+using JumpProcesses
+
+needs_depgraph(Direct()) == false
+needs_depgraph(NRM()) == true
+```
+"""
 needs_depgraph(aggregator::AbstractAggregatorAlgorithm) = false
 needs_depgraph(aggregator::DirectCR) = true
 needs_depgraph(aggregator::SortingDirect) = true
@@ -190,6 +217,30 @@ needs_depgraph(aggregator::Coevolve) = true
 # true if aggregator requires a map from solution variable to dependent jumps.
 # It is implicitly assumed these aggregators also require the reverse map, from
 # jumps to variables they update.
+"""
+    needs_vartojumps_map(aggregator) -> Bool
+
+Return whether an aggregator requires the species-to-reaction dependency map used by
+RSSA-style bracketing.
+
+## Arguments
+
+  - `aggregator`: An `AbstractAggregatorAlgorithm` value.
+
+## Returns
+
+  - `true` for [`RSSA`](@ref) and [`RSSACR`](@ref).
+  - `false` for other built-in aggregators.
+
+## Examples
+
+```julia
+using JumpProcesses
+
+needs_vartojumps_map(RSSA()) == true
+needs_vartojumps_map(Direct()) == false
+```
+"""
 needs_vartojumps_map(aggregator::AbstractAggregatorAlgorithm) = false
 needs_vartojumps_map(aggregator::RSSA) = true
 needs_vartojumps_map(aggregator::RSSACR) = true

src/aggregators/bracketing.jl

@@ -3,6 +3,38 @@
 #       large-scale reaction networks", Thanh et al, J. Chem. Phys., 2015
 # note, expects the type of the bracketing variables [ulow,uhigh] to be the
 # same as the fluct_rate and ushift.
+"""
+    BracketData(fluctrate, threshold, Δu)
+    BracketData{T1, T2}()
+
+Configure species-population brackets used by RSSA-based aggregators.
+
+For species population `u[i]`, the bracket is
+`[(1 - fluctrate) * u[i], (1 + fluctrate) * u[i]]` when `u[i] >= threshold`.
+For smaller populations, the bracket is `[max(u[i] - Δu, 0), u[i] + Δu]`.
+Each field may be either a scalar shared by all species or a vector indexed by species.
+
+## Fields
+
+  - `fluctrate`: Relative fluctuation width used for populations at or above `threshold`.
+  - `threshold`: Population threshold below which the absolute `Δu` bracket is used.
+  - `Δu`: Absolute bracket half-width used for populations below `threshold`.
+
+## Notes
+
+  - `BracketData{T1, T2}()` constructs `BracketData(T1(0.1), T2(25), T2(4))`.
+  - The bracketing rules follow the RSSA construction in Thanh et al., J. Chem. Phys. 142,
+    244106 (2015).
+
+## Examples
+
+```julia
+using JumpProcesses
+
+bd = BracketData(0.1, 25, 4)
+bd.fluctrate == 0.1
+```
+"""
 struct BracketData{T1, T2}
     fluctrate::T1         # interval should be [1-fluctrate,1+fluctrate] * u
     threshold::T2         # for u below threshold interval is:

src/jumps.jl

@@ -429,6 +429,30 @@ end
 using_params(maj::MassActionJump{T, S, U, Nothing}) where {T, S, U} = false
 using_params(maj::MassActionJump) = true
 using_params(maj::Nothing) = false
+"""
+    get_num_majumps(jumps) -> Int
+
+Return the number of mass-action jumps represented by a jump container.
+
+## Arguments
+
+  - `jumps`: A [`MassActionJump`](@ref), [`SpatialMassActionJump`](@ref),
+    [`JumpSet`](@ref), or `nothing`.
+
+## Returns
+
+  - The number of mass-action reaction channels. `nothing` returns `0`.
+
+## Examples
+
+```julia
+using JumpProcesses
+
+maj = MassActionJump([1.0], [[1 => 1]], [[1 => -1]])
+get_num_majumps(maj) == 1
+get_num_majumps(nothing) == 0
+```
+"""
 @inline get_num_majumps(maj::MassActionJump) = length(maj.net_stoch)
 @inline get_num_majumps(maj::Nothing) = 0
 

src/simple_regular_solve.jl

@@ -1,5 +1,75 @@
+"""
+    SimpleTauLeaping()
+
+Fixed-step tau-leaping algorithm for pure [`RegularJump`](@ref) problems.
+
+Use `SimpleTauLeaping` with `JumpProblem(prob, PureLeaping(), regular_jump)` and pass
+the timestep through the `dt` keyword to `solve`.
+
+## Keyword Arguments
+
+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.
+  - `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.
+
+## Returns
+
+  - A stateless `SciMLBase.AbstractDEAlgorithm` value.
+
+## Examples
+
+```julia
+using JumpProcesses, DiffEqBase
+
+rate!(out, u, p, t) = (out[1] = 0.1 * u[1])
+affect!(du, u, p, t, counts, mark) = (du[1] = -counts[1])
+rj = RegularJump(rate!, affect!, 1)
+
+prob = DiscreteProblem([20], (0.0, 2.0))
+jprob = JumpProblem(prob, PureLeaping(), rj)
+sol = solve(jprob, SimpleTauLeaping(); dt = 0.1)
+```
+"""
 struct SimpleTauLeaping <: SciMLBase.AbstractDEAlgorithm end
 
+"""
+    SimpleExplicitTauLeaping(; epsilon = 0.05)
+    SimpleExplicitTauLeaping(epsilon)
+
+Adaptive explicit tau-leaping algorithm for pure [`MassActionJump`](@ref) problems.
+
+Use `SimpleExplicitTauLeaping` with `JumpProblem(prob, PureLeaping(), mass_action_jump)`.
+The algorithm computes step sizes from the mass-action propensities and the error-control
+parameter `epsilon`.
+
+## Arguments
+
+  - `epsilon`: Positive floating-point error-control parameter. Smaller values generally
+    produce smaller steps.
+
+## Fields
+
+  - `epsilon`: Stored error-control parameter used by the adaptive leaping step selector.
+
+## Returns
+
+  - A `SciMLBase.AbstractDEAlgorithm` value.
+
+## Examples
+
+```julia
+using JumpProcesses, DiffEqBase
+
+maj = MassActionJump([0.1], [[1 => 1]], [[1 => -1]])
+prob = DiscreteProblem([20], (0.0, 2.0))
+jprob = JumpProblem(prob, PureLeaping(), maj)
+sol = solve(jprob, SimpleExplicitTauLeaping())
+```
+"""
 struct SimpleExplicitTauLeaping{T <: AbstractFloat} <: SciMLBase.AbstractDEAlgorithm
     epsilon::T  # Error control parameter
 end
@@ -405,6 +475,35 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping;
     return sol
 end
 
+"""
+    EnsembleGPUKernel()
+    EnsembleGPUKernel(backend)
+
+Ensemble algorithm marker for GPU execution of tau-leaping ensemble simulations.
+
+## Arguments
+
+  - `backend`: Optional KernelAbstractions-compatible backend. `nothing` requests the
+    default backend selected by the extension.
+
+## Fields
+
+  - `backend`: Backend object used by the GPU extension.
+  - `cpu_offload`: Fraction of trajectories to offload to CPU execution.
+
+## Returns
+
+  - A `SciMLBase.EnsembleAlgorithm` value for use as the ensemble algorithm argument to
+    `solve`.
+
+## Examples
+
+```julia
+using JumpProcesses
+
+ensemble_alg = EnsembleGPUKernel()
+```
+"""
 struct EnsembleGPUKernel{Backend} <: SciMLBase.EnsembleAlgorithm
     backend::Backend
     cpu_offload::Float64