From b36bc13c095782876176849762070eceb5435d24 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Thu, 9 Jul 2026 04:55:33 -0400 Subject: [PATCH 1/5] Document JumpProcesses public API Co-Authored-By: Chris Rackauckas --- docs/src/api.md | 39 ++- src/JumpProcesses.jl | 43 +++- src/aggregators/aggregators.jl | 63 ++++- src/aggregators/bracketing.jl | 61 +++-- src/jumps.jl | 338 ++++++++++++++++--------- src/simple_regular_solve.jl | 174 ++++++++++--- src/spatial/spatial_massaction_jump.jl | 267 +++++++++++++------ src/spatial/topology.jl | 134 +++++++++- src/variable_rate.jl | 183 ++++++++----- test/public_api_docs.jl | 64 +++++ test/runtests.jl | 1 + 11 files changed, 1043 insertions(+), 324 deletions(-) create mode 100644 test/public_api_docs.jl diff --git a/docs/src/api.md b/docs/src/api.md index d8fd944b2..9302b9d81 100644 --- a/docs/src/api.md +++ b/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 ``` diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl index d6ce8bfe9..ce270344f 100644 --- a/src/JumpProcesses.jl +++ b/src/JumpProcesses.jl @@ -36,10 +36,10 @@ import SymbolicIndexingInterface as SII # Import additional types and functions from DiffEqBase and SciMLBase using DiffEqBase: DiffEqBase, CallbackSet, ContinuousCallback, DAEFunction, - DDEFunction, DiscreteProblem, ODEFunction, ODEProblem, - ODESolution, ReturnCode, SDEFunction, SDEProblem, add_tstop!, - deleteat!, isinplace, remake, savevalues!, step!, - derivative_discontinuity! + DDEFunction, DiscreteProblem, ODEFunction, ODEProblem, + ODESolution, ReturnCode, SDEFunction, SDEProblem, add_tstop!, + deleteat!, isinplace, remake, savevalues!, step!, + derivative_discontinuity! using SciMLBase: SciMLBase, DEIntegrator abstract type AbstractJump end @@ -47,7 +47,7 @@ abstract type AbstractMassActionJump <: AbstractJump end abstract type AbstractAggregatorAlgorithm end abstract type AbstractJumpAggregator end abstract type AbstractSSAIntegrator{Alg, IIP, U, T} <: - DEIntegrator{Alg, IIP, U, T} end +DEIntegrator{Alg, IIP, U, T} end const DEFAULT_RNG = Random.default_rng() @@ -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 @@ -129,7 +156,7 @@ export init, solve, solve! include("SSA_stepper.jl") export SSAStepper -# leaping: +# leaping: include("simple_regular_solve.jl") export SimpleTauLeaping, SimpleExplicitTauLeaping, EnsembleGPUKernel diff --git a/src/aggregators/aggregators.jl b/src/aggregators/aggregators.jl index b8dd8dd5b..1e3dcd503 100644 --- a/src/aggregators/aggregators.jl +++ b/src/aggregators/aggregators.jl @@ -172,13 +172,42 @@ algorithm with optimal binning, Journal of Chemical Physics 143, 074108 """ struct DirectCRDirect <: AbstractAggregatorAlgorithm end -const JUMP_AGGREGATORS = (Direct(), DirectFW(), DirectCR(), SortingDirect(), RSSA(), FRM(), - FRMFW(), NRM(), RSSACR(), RDirect(), Coevolve(), CCNRM()) +const JUMP_AGGREGATORS = ( + Direct(), DirectFW(), DirectCR(), SortingDirect(), RSSA(), FRM(), + FRMFW(), NRM(), RSSACR(), RDirect(), Coevolve(), CCNRM(), +) # For JumpProblem construction without an aggregator 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 +219,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 @@ -204,9 +257,11 @@ is_spatial(aggregator::NSM) = true is_spatial(aggregator::DirectCRDirect) = true # return the fastest aggregator out of the available ones -function select_aggregator(jumps::JumpSet; vartojumps_map = nothing, +function select_aggregator( + jumps::JumpSet; vartojumps_map = nothing, jumptovars_map = nothing, dep_graph = nothing, spatial_system = nothing, - hopping_constants = nothing) + hopping_constants = nothing + ) # detect if a spatial SSA should be used !isnothing(spatial_system) && !isnothing(hopping_constants) && return DirectCRDirect diff --git a/src/aggregators/bracketing.jl b/src/aggregators/bracketing.jl index 18bb9ece3..e6abfdd03 100644 --- a/src/aggregators/bracketing.jl +++ b/src/aggregators/bracketing.jl @@ -1,12 +1,39 @@ -# bracketing intervals for RSSA-based aggregators -# see: "On the rejection-based algorithm for simulation and analysis of -# 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: - Δu::T2 # [max(u-Δu,0),u+Δu] + fluctrate::T1 + threshold::T2 + Δu::T2 end # # suggested defaults @@ -26,7 +53,7 @@ BracketData{T1, T2}() where {T1, T2} = BracketData(T1(0.1), T2(25), T2(4)) @inline getΔu(bd::BracketData{T1, T2}, i) where {T1, T2 <: Number} = bd.Δu @inline function delta_bracket(u::Integer, δ) - (trunc(typeof(u), (one(δ) - δ) * u), trunc(typeof(u), (one(δ) + δ) * u)) + return (trunc(typeof(u), (one(δ) - δ) * u), trunc(typeof(u), (one(δ) + δ) * u)) end @inline delta_bracket(u, δ) = ((one(δ) - δ) * u), ((one(δ) + δ) * u) @@ -48,7 +75,7 @@ end # Get propensity brackets of massaction jump k. @inline function get_majump_brackets(ulow, uhigh, k, majumps) - evalrxrate(ulow, k, majumps), evalrxrate(uhigh, k, majumps) + return evalrxrate(ulow, k, majumps), evalrxrate(uhigh, k, majumps) end # for constant rate jumps we must check the ordering of the bracket values @@ -68,8 +95,10 @@ get brackets for the rate of reaction rx by first checking if the reaction is a if rx <= num_majumps return get_majump_brackets(p.ulow, p.uhigh, rx, ma_jumps) else - @inbounds return get_cjump_brackets(p.ulow, p.uhigh, p.rates[rx - num_majumps], - params, t) + @inbounds return get_cjump_brackets( + p.ulow, p.uhigh, p.rates[rx - num_majumps], + params, t + ) end end @@ -79,7 +108,7 @@ end @inbounds for (i, uval) in enumerate(u) ulow[i], uhigh[i] = get_spec_brackets(p.bracket_data, i, uval) end - nothing + return nothing end @inline function update_u_brackets!(p::AbstractSSAJumpAggregator, u::SVector) @@ -88,12 +117,12 @@ end p.ulow = setindex(p.ulow, ulow, i) p.uhigh = setindex(p.uhigh, uhigh, i) end - nothing + return nothing end # For ExtendedJumpArray, only iterate over the species portion (u.u), not the jump tracking portion @inline function update_u_brackets!(p::AbstractSSAJumpAggregator, u::ExtendedJumpArray) - update_u_brackets!(p, u.u) + return update_u_brackets!(p, u.u) end # Set up bracketing. The aggregator must have fields @@ -122,5 +151,5 @@ function set_bracketing!(p::AbstractSSAJumpAggregator, u, params, t) end p.sum_rate = sum_rate - nothing + return nothing end diff --git a/src/jumps.jl b/src/jumps.jl index 07185d954..013b8ba4e 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -171,15 +171,17 @@ function VariableRateJump(rate, affect!; lrate = nothing, urate = nothing, abstol = 1e-12, reltol = 0) ``` """ -function VariableRateJump(rate, affect!; +function VariableRateJump( + rate, affect!; lrate = nothing, urate = nothing, rateinterval = nothing, rootfind = true, idxs = nothing, save_positions = (false, true), interp_points = 10, - abstol = 1e-12, reltol = 0) + abstol = 1.0e-12, reltol = 0 + ) if !(urate !== nothing && rateinterval !== nothing) && - !(urate === nothing && rateinterval === nothing) + !(urate === nothing && rateinterval === nothing) error("`urate` and `rateinterval` must both be `nothing`, or must both be defined.") end @@ -188,8 +190,10 @@ function VariableRateJump(rate, affect!; error("If a lower bound rate, `lrate`, is given than an upper bound rate, `urate`, and rate interval, `rateinterval`, must also be provided.") end - VariableRateJump(rate, affect!, lrate, urate, rateinterval, idxs, rootfind, - interp_points, save_positions, abstol, reltol) + return VariableRateJump( + rate, affect!, lrate, urate, rateinterval, idxs, rootfind, + interp_points, save_positions, abstol, reltol + ) end """ @@ -241,14 +245,14 @@ struct RegularJump{iip, R, C, MD} """ 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) + return new{iip, typeof(rate), typeof(c), typeof(mark_dist)}(rate, c, numjumps, mark_dist) end end DiffEqBase.isinplace(::RegularJump{iip, R, C, MD}) where {iip, R, C, MD} = iip function RegularJump(rate, c, numjumps::Int; kwargs...) - RegularJump{DiffEqBase.isinplace(rate, 4)}(rate, c, numjumps; kwargs...) + return RegularJump{DiffEqBase.isinplace(rate, 4)}(rate, c, numjumps; kwargs...) end # deprecate old call @@ -256,9 +260,9 @@ function RegularJump(rate, c, dc::AbstractMatrix; constant_c = false, mark_dist @warn("The RegularJump interface has changed to be matrix-free. See the documentation for more details.") function _c(du, u, p, t, counts, mark) c(dc, u, p, t, mark) - mul!(du, dc, counts) + return mul!(du, dc, counts) end - RegularJump{true}(rate, _c, size(dc, 2); mark_dist = mark_dist) + return RegularJump{true}(rate, _c, size(dc, 2); mark_dist = mark_dist) end """ @@ -336,10 +340,12 @@ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump """Whether `update_parameters!` should apply stoichiometric scaling to rates.""" rescale_rates_on_update::Bool - function MassActionJump{T, S, U, V}(rates::T, rs_in::S, ns::U, pmapper::V, + 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) @@ -351,67 +357,91 @@ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump if scale_rates && !isempty(sr) scalerates!(sr, rs) end - new(sr, rs, ns, pmapper, rescale_rates_on_update) + return new(sr, rs, ns, pmapper, rescale_rates_on_update) end - function MassActionJump{Nothing, Vector{S}, - Vector{U}, V}(::Nothing, rs_in::Vector{S}, + function MassActionJump{ + Nothing, Vector{S}, + Vector{U}, V, + }( + ::Nothing, rs_in::Vector{S}, ns::Vector{U}, pmapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool, - rescale_rates_on_update::Bool = scale_rates) where {S <: AbstractVector, - U <: AbstractVector, V} + rescale_rates_on_update::Bool = scale_rates + ) where { + S <: AbstractVector, + U <: AbstractVector, V, + } rs = nocopy ? rs_in : copy(rs_in) for i in eachindex(rs) if useiszero && (length(rs[i]) == 1) && iszero(rs[i][1][1]) rs[i] = typeof(rs[i])() end end - new(nothing, rs, ns, pmapper, rescale_rates_on_update) + return new(nothing, rs, ns, pmapper, rescale_rates_on_update) end - function MassActionJump{T, S, U, V}(rate::T, rs_in::S, ns::U, pmapper::V, + function MassActionJump{T, S, U, V}( + rate::T, rs_in::S, ns::U, pmapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool, - rescale_rates_on_update::Bool = scale_rates) where {T <: Number, S, U, V} + rescale_rates_on_update::Bool = scale_rates + ) where {T <: Number, S, U, V} rs = rs_in if useiszero && (length(rs) == 1) && iszero(rs[1][1]) rs = typeof(rs)() end sr = scale_rates ? scalerate(rate, rs) : rate - new(sr, rs, ns, pmapper, rescale_rates_on_update) + return new(sr, rs, ns, pmapper, rescale_rates_on_update) end - function MassActionJump{Nothing, S, U, V}(::Nothing, rs_in::S, ns::U, pmapper::V, + function MassActionJump{Nothing, S, U, V}( + ::Nothing, rs_in::S, ns::U, pmapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool, - rescale_rates_on_update::Bool = scale_rates) where {S, U, V} + rescale_rates_on_update::Bool = scale_rates + ) where {S, U, V} rs = rs_in if useiszero && (length(rs) == 1) && iszero(rs[1][1]) rs = typeof(rs)() end - new(nothing, rs, ns, pmapper, rescale_rates_on_update) + return new(nothing, rs, ns, pmapper, rescale_rates_on_update) end end -function MassActionJump(usr::T, rs::S, ns::U, pmapper::V; scale_rates = true, +function MassActionJump( + usr::T, rs::S, ns::U, pmapper::V; scale_rates = true, useiszero = true, nocopy = false, - rescale_rates_on_update = scale_rates) where {T, S, U, V} - MassActionJump{T, S, U, V}(usr, rs, ns, pmapper, scale_rates, useiszero, nocopy, - rescale_rates_on_update) -end -function MassActionJump(usr::T, rs, ns; scale_rates = true, useiszero = true, - nocopy = false, rescale_rates_on_update = scale_rates) where {T <: AbstractVector} - MassActionJump(usr, rs, ns, nothing; scale_rates, useiszero, nocopy, - rescale_rates_on_update) -end -function MassActionJump(usr::T, rs, ns; scale_rates = true, useiszero = true, - nocopy = false, rescale_rates_on_update = scale_rates) where {T <: Number} - MassActionJump(usr, rs, ns, nothing; scale_rates, useiszero, nocopy, - rescale_rates_on_update) + rescale_rates_on_update = scale_rates + ) where {T, S, U, V} + return MassActionJump{T, S, U, V}( + usr, rs, ns, pmapper, scale_rates, useiszero, nocopy, + rescale_rates_on_update + ) +end +function MassActionJump( + usr::T, rs, ns; scale_rates = true, useiszero = true, + nocopy = false, rescale_rates_on_update = scale_rates + ) where {T <: AbstractVector} + return MassActionJump( + usr, rs, ns, nothing; scale_rates, useiszero, nocopy, + rescale_rates_on_update + ) +end +function MassActionJump( + usr::T, rs, ns; scale_rates = true, useiszero = true, + nocopy = false, rescale_rates_on_update = scale_rates + ) where {T <: Number} + return MassActionJump( + usr, rs, ns, nothing; scale_rates, useiszero, nocopy, + rescale_rates_on_update + ) end # with parameter indices or mapping, multiple jump case -function MassActionJump(rs, ns; param_idxs = nothing, param_mapper = nothing, +function MassActionJump( + rs, ns; param_idxs = nothing, param_mapper = nothing, scale_rates = true, useiszero = true, nocopy = false, - rescale_rates_on_update = scale_rates) + rescale_rates_on_update = scale_rates + ) if param_mapper === nothing (param_idxs === nothing) && error("If no parameter indices are given via param_idxs, an explicit parameter mapping must be passed in via param_mapper.") @@ -422,13 +452,39 @@ function MassActionJump(rs, ns; param_idxs = nothing, param_mapper = nothing, pmapper = param_mapper end - MassActionJump(nothing, nocopy ? rs : copy(rs), ns, pmapper; scale_rates, useiszero, - nocopy = true, rescale_rates_on_update) + return MassActionJump( + nothing, nocopy ? rs : copy(rs), ns, pmapper; scale_rates, useiszero, + nocopy = true, rescale_rates_on_update + ) end using_params(maj::MassActionJump{T, S, U, Nothing}) where {T, S, U} = false 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 @@ -439,43 +495,53 @@ end # create the initial parameter vector for use in a MassActionJump # Note these are unscaled function (ratemap::MassActionJumpParamMapper{U})(params) where {U <: AbstractArray} - [params[pidx] for pidx in ratemap.param_idxs] + return [params[pidx] for pidx in ratemap.param_idxs] end # Note this is unscaled function (ratemap::MassActionJumpParamMapper{U})(params) where {U <: Int} - params[ratemap.param_idxs] + return params[ratemap.param_idxs] end # update a maj with parameter vectors -function (ratemap::MassActionJumpParamMapper{U})(maj::MassActionJump, newparams; +function (ratemap::MassActionJumpParamMapper{U})( + maj::MassActionJump, newparams; scale_rates, - kwargs...) where {U <: AbstractArray} + kwargs... + ) where {U <: AbstractArray} for i in 1:get_num_majumps(maj) maj.scaled_rates[i] = newparams[ratemap.param_idxs[i]] end scale_rates && scalerates!(maj.scaled_rates, maj.reactant_stoch) - nothing + return nothing end function to_collection(ratemap::MassActionJumpParamMapper{Int}) - MassActionJumpParamMapper([ratemap.param_idxs]) + return MassActionJumpParamMapper([ratemap.param_idxs]) end -function Base.merge!(pmap1::MassActionJumpParamMapper{U}, - pmap2::MassActionJumpParamMapper{U}) where {U <: AbstractVector} - append!(pmap1.param_idxs, pmap2.param_idxs) +function Base.merge!( + pmap1::MassActionJumpParamMapper{U}, + pmap2::MassActionJumpParamMapper{U} + ) where {U <: AbstractVector} + return 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) +function Base.merge!( + pmap1::MassActionJumpParamMapper{U}, + pmap2::MassActionJumpParamMapper{V} + ) where { + U <: AbstractVector, + V <: Int, + } + return push!(pmap1.param_idxs, pmap2.param_idxs) end -function Base.merge(pmap1::MassActionJumpParamMapper{Int}, - pmap2::MassActionJumpParamMapper{Int}) - MassActionJumpParamMapper([pmap1.param_idxs, pmap2.param_idxs]) +function Base.merge( + pmap1::MassActionJumpParamMapper{Int}, + pmap2::MassActionJumpParamMapper{Int} + ) + return MassActionJumpParamMapper([pmap1.param_idxs, pmap2.param_idxs]) end """ @@ -494,7 +560,7 @@ Notes: 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) + return maj.param_mapper(maj, newparams; scale_rates, kwargs) end """ @@ -542,22 +608,24 @@ struct JumpSet{T1, T2, T3, T4} <: AbstractJump massaction_jump::T4 end function JumpSet(vj, cj, rj, maj::MassActionJump{S, T, U, V}) where {S <: Number, T, U, V} - JumpSet(vj, cj, rj, check_majump_type(maj)) + return JumpSet(vj, cj, rj, check_majump_type(maj)) end JumpSet(jump::ConstantRateJump) = JumpSet((), (jump,), nothing, nothing) JumpSet(jump::VariableRateJump) = JumpSet((jump,), (), nothing, nothing) JumpSet(jump::RegularJump) = JumpSet((), (), jump, nothing) JumpSet(jump::AbstractMassActionJump) = JumpSet((), (), nothing, jump) -function JumpSet(; variable_jumps = (), constant_jumps = (), - regular_jumps = nothing, massaction_jumps = nothing) - JumpSet(variable_jumps, constant_jumps, regular_jumps, massaction_jumps) +function JumpSet(; + variable_jumps = (), constant_jumps = (), + regular_jumps = nothing, massaction_jumps = nothing + ) + return JumpSet(variable_jumps, constant_jumps, regular_jumps, massaction_jumps) end JumpSet(jb::Nothing) = JumpSet() # For Varargs, use recursion to make it type-stable function JumpSet(jumps::AbstractJump...) - JumpSet(split_jumps((), (), nothing, nothing, jumps...)...) + return JumpSet(split_jumps((), (), nothing, nothing, jumps...)...) end # handle vector of mass action jumps @@ -570,32 +638,34 @@ function JumpSet(vjs, cjs, rj, majv::Vector{T}) where {T <: MassActionJump} if !all(m -> m.rescale_rates_on_update == sr_val, majv) error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") end - maj = setup_majump_to_merge(majv[1].scaled_rates, majv[1].reactant_stoch, - majv[1].net_stoch, majv[1].param_mapper, sr_val) + maj = setup_majump_to_merge( + majv[1].scaled_rates, majv[1].reactant_stoch, + majv[1].net_stoch, majv[1].param_mapper, sr_val + ) for i in 2:length(majv) massaction_jump_combine(maj, majv[i]) end - JumpSet(vjs, cjs, rj, maj) + return JumpSet(vjs, cjs, rj, maj) end @inline get_num_majumps(jset::JumpSet) = get_num_majumps(jset.massaction_jump) @inline num_majumps(jset::JumpSet) = get_num_majumps(jset) @inline function num_crjs(jset::JumpSet) - (jset.constant_jumps !== nothing) ? length(jset.constant_jumps) : 0 + return (jset.constant_jumps !== nothing) ? length(jset.constant_jumps) : 0 end @inline function num_vrjs(jset::JumpSet) - (jset.variable_jumps !== nothing) ? length(jset.variable_jumps) : 0 + return (jset.variable_jumps !== nothing) ? length(jset.variable_jumps) : 0 end @inline function num_bndvrjs(jset::JumpSet) - (jset.variable_jumps !== nothing) ? count(isbounded, jset.variable_jumps) : 0 + return (jset.variable_jumps !== nothing) ? count(isbounded, jset.variable_jumps) : 0 end @inline function num_continvrjs(jset::JumpSet) - (jset.variable_jumps !== nothing) ? count(!isbounded, jset.variable_jumps) : 0 + return (jset.variable_jumps !== nothing) ? count(!isbounded, jset.variable_jumps) : 0 end num_jumps(jset::JumpSet) = num_majumps(jset) + num_crjs(jset) + num_vrjs(jset) @@ -604,22 +674,24 @@ num_cdiscretejumps(jset::JumpSet) = num_majumps(jset) + num_crjs(jset) @inline split_jumps(vj, cj, rj, maj) = vj, cj, rj, maj @inline function split_jumps(vj, cj, rj, maj, v::VariableRateJump, args...) - split_jumps((vj..., v), cj, rj, maj, args...) + return split_jumps((vj..., v), cj, rj, maj, args...) end @inline function split_jumps(vj, cj, rj, maj, c::ConstantRateJump, args...) - split_jumps(vj, (cj..., c), rj, maj, args...) + return split_jumps(vj, (cj..., c), rj, maj, args...) end @inline function split_jumps(vj, cj, rj, maj, c::RegularJump, args...) - split_jumps(vj, cj, regular_jump_combine(rj, c), maj, args...) + return split_jumps(vj, cj, regular_jump_combine(rj, c), maj, args...) end @inline function split_jumps(vj, cj, rj, maj, c::MassActionJump, args...) - split_jumps(vj, cj, rj, massaction_jump_combine(maj, c), args...) + return split_jumps(vj, cj, rj, massaction_jump_combine(maj, c), args...) end @inline function split_jumps(vj, cj, rj, maj, j::JumpSet, args...) - split_jumps((vj..., j.variable_jumps...), + return split_jumps( + (vj..., j.variable_jumps...), (cj..., j.constant_jumps...), regular_jump_combine(rj, j.regular_jump), - massaction_jump_combine(maj, j.massaction_jump), args...) + massaction_jump_combine(maj, j.massaction_jump), args... + ) end regular_jump_combine(rj1::RegularJump, rj2::Nothing) = rj1 @@ -631,43 +703,65 @@ end # functionality to merge two mass action jumps together function check_majump_type(maj::MassActionJump{S, T, U, V}) where {S <: Number, T, U, V} - setup_majump_to_merge(maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, - maj.param_mapper, maj.rescale_rates_on_update) + return setup_majump_to_merge( + maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, + maj.param_mapper, maj.rescale_rates_on_update + ) end function check_majump_type(maj::MassActionJump{Nothing, T, U, V}) where {T, U, V} - setup_majump_to_merge(maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, - maj.param_mapper, maj.rescale_rates_on_update) + return setup_majump_to_merge( + maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, + maj.param_mapper, maj.rescale_rates_on_update + ) end # if given containers of rates and stoichiometry directly create a jump -function setup_majump_to_merge(sr::T, rs::AbstractVector{S}, ns::AbstractVector{U}, pmapper, - rescale_rates_on_update::Bool) where {T <: AbstractVector, S <: AbstractArray, - U <: AbstractArray} - MassActionJump(sr, rs, ns, pmapper; scale_rates = false, rescale_rates_on_update) +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, + } + return MassActionJump(sr, rs, ns, pmapper; scale_rates = false, rescale_rates_on_update) end # if just given the data for one jump (and not in a container) wrap in a vector -function setup_majump_to_merge(sr::S, rs::T, ns::U, pmapper, - rescale_rates_on_update::Bool) where {S <: Number, T <: AbstractArray, - U <: AbstractArray} - MassActionJump([sr], [rs], [ns], +function setup_majump_to_merge( + sr::S, rs::T, ns::U, pmapper, + rescale_rates_on_update::Bool + ) where { + S <: Number, T <: AbstractArray, + U <: AbstractArray, + } + return MassActionJump( + [sr], [rs], [ns], (pmapper === nothing) ? pmapper : to_collection(pmapper); - scale_rates = false, rescale_rates_on_update) + scale_rates = false, rescale_rates_on_update + ) end # if no rate field setup yet -function setup_majump_to_merge(::Nothing, rs::T, ns::U, pmapper, - rescale_rates_on_update::Bool) where {T <: AbstractArray, U <: AbstractArray} - MassActionJump(nothing, [rs], [ns], +function setup_majump_to_merge( + ::Nothing, rs::T, ns::U, pmapper, + rescale_rates_on_update::Bool + ) where {T <: AbstractArray, U <: AbstractArray} + return MassActionJump( + nothing, [rs], [ns], (pmapper === nothing) ? pmapper : to_collection(pmapper); - scale_rates = false, rescale_rates_on_update) + scale_rates = false, rescale_rates_on_update + ) end # when given a collection of reactions to add to maj -function majump_merge!(maj::MassActionJump{U, <:AbstractVector{V}, <:AbstractVector{W}, X}, +function majump_merge!( + maj::MassActionJump{U, <:AbstractVector{V}, <:AbstractVector{W}, X}, sr::U, rs::AbstractVector{V}, ns::AbstractVector{W}, - param_mapper) where {U <: Union{AbstractVector, Nothing}, - V <: AbstractVector, W <: AbstractVector, X} + param_mapper + ) where { + U <: Union{AbstractVector, Nothing}, + V <: AbstractVector, W <: AbstractVector, X, + } (U <: AbstractVector) && append!(maj.scaled_rates, sr) append!(maj.reactant_stoch, rs) append!(maj.net_stoch, ns) @@ -677,16 +771,20 @@ function majump_merge!(maj::MassActionJump{U, <:AbstractVector{V}, <:AbstractVec else merge!(maj.param_mapper, param_mapper) end - maj + return maj end # when given a single jump's worth of data to add to maj -function majump_merge!(maj::MassActionJump{U, V, W, X}, sr::T, rs::S1, ns::S2, - param_mapper) where {T <: Union{Number, Nothing}, +function majump_merge!( + maj::MassActionJump{U, V, W, X}, sr::T, rs::S1, ns::S2, + param_mapper + ) where { + T <: Union{Number, Nothing}, S1 <: AbstractArray, S2 <: AbstractArray, U <: Union{AbstractVector{T}, Nothing}, V <: AbstractVector{S1}, - W <: AbstractVector{S2}, X} + W <: AbstractVector{S2}, X, + } (T <: Number) && push!(maj.scaled_rates, sr) push!(maj.reactant_stoch, rs) push!(maj.net_stoch, ns) @@ -697,26 +795,34 @@ function majump_merge!(maj::MassActionJump{U, V, W, X}, sr::T, rs::S1, ns::S2, merge!(maj.param_mapper, param_mapper) end - maj + return maj end # when maj only stores a single jump's worth of data (and not in a collection) # create a new jump with the merged data stored in vectors -function majump_merge!(maj::MassActionJump{T, S, U, V}, sr::T, rs::S, ns::U, - param_mapper::V) where {T <: Union{Number, Nothing}, +function majump_merge!( + maj::MassActionJump{T, S, U, V}, sr::T, rs::S, ns::U, + param_mapper::V + ) where { + T <: Union{Number, Nothing}, S <: AbstractArray{<:Pair}, - U <: AbstractArray{<:Pair}, V} + U <: AbstractArray{<:Pair}, V, + } rates = (T <: Nothing) ? nothing : [maj.scaled_rates, sr] if maj.param_mapper === nothing (param_mapper === nothing) || error("Error, trying to merge a MassActionJump with a parameter mapping to one without a parameter mapping.") - return MassActionJump(rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], + return MassActionJump( + rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], param_mapper; scale_rates = false, - rescale_rates_on_update = maj.rescale_rates_on_update) + rescale_rates_on_update = maj.rescale_rates_on_update + ) else - return MassActionJump(rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], + return MassActionJump( + rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], merge(maj.param_mapper, param_mapper); scale_rates = false, - rescale_rates_on_update = maj.rescale_rates_on_update) + rescale_rates_on_update = maj.rescale_rates_on_update + ) end end @@ -726,8 +832,10 @@ massaction_jump_combine(maj1::Nothing, maj2::Nothing) = maj1 function massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) (maj1.rescale_rates_on_update == maj2.rescale_rates_on_update) || error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") - majump_merge!(maj1, maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, - maj2.param_mapper) + return majump_merge!( + maj1, maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, + maj2.param_mapper + ) end ##### helper methods for unpacking rates and affects! from constant jumps ##### @@ -740,12 +848,14 @@ function get_jump_info_tuples(jumps) affects! = () end - rates, affects! + return rates, affects! end function get_jump_info_fwrappers(u, p, t, jumps) - RateWrapper = FunctionWrappers.FunctionWrapper{typeof(t), - Tuple{typeof(u), typeof(p), typeof(t)}} + RateWrapper = FunctionWrappers.FunctionWrapper{ + typeof(t), + Tuple{typeof(u), typeof(p), typeof(t)}, + } if (jumps !== nothing) && !isempty(jumps) rates = [RateWrapper(c.rate) for c in jumps] @@ -755,5 +865,5 @@ function get_jump_info_fwrappers(u, p, t, jumps) affects! = Any[] end - rates, affects! + return rates, affects! end diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index 2445237b7..4e08378c2 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -1,7 +1,77 @@ +""" + 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 + epsilon::T end SimpleExplicitTauLeaping(; epsilon = 0.05) = SimpleExplicitTauLeaping(epsilon) @@ -12,7 +82,7 @@ function validate_pure_leaping_inputs(jump_prob::JumpProblem, alg) JumpProblem, i.e. call JumpProblem(::DiscreteProblem, PureLeaping(),...). \ Passing $(jump_prob.aggregator) is deprecated and will be removed in the next breaking release." end - isempty(jump_prob.jump_callback.continuous_callbacks) && + return isempty(jump_prob.jump_callback.continuous_callbacks) && isempty(jump_prob.jump_callback.discrete_callbacks) && isempty(jump_prob.constant_jumps) && isempty(jump_prob.variable_jumps) && @@ -26,7 +96,7 @@ function validate_pure_leaping_inputs(jump_prob::JumpProblem, alg::SimpleExplici JumpProblem, i.e. call JumpProblem(::DiscreteProblem, PureLeaping(),...). \ Passing $(jump_prob.aggregator) is deprecated and will be removed in the next breaking release." end - isempty(jump_prob.jump_callback.continuous_callbacks) && + return isempty(jump_prob.jump_callback.continuous_callbacks) && isempty(jump_prob.jump_callback.discrete_callbacks) && isempty(jump_prob.constant_jumps) && isempty(jump_prob.variable_jumps) && @@ -52,7 +122,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 @@ -69,9 +139,11 @@ function _process_saveat(saveat, tspan, save_start, save_end) return saveat_vec, _save_start, _save_end end -function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping; +function DiffEqBase.solve( + jump_prob::JumpProblem, alg::SimpleTauLeaping; seed = nothing, dt = error("dt is required for SimpleTauLeaping."), - saveat = nothing, save_start = nothing, save_end = nothing) + 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.") @@ -142,9 +214,11 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping; push!(tsave, tspan[2]) end - sol = SciMLBase.build_solution(prob, alg, tsave, usave, + return sol = SciMLBase.build_solution( + prob, alg, tsave, usave, calculate_error = false, - interp = SciMLBase.ConstantInterpolation(tsave, usave)) + interp = SciMLBase.ConstantInterpolation(tsave, usave) + ) end # Compute the highest order of reaction (HOR) for each reaction j, as per Cao et al. (2006), Section IV. @@ -155,7 +229,8 @@ function compute_hor(reactant_stoch, numjumps) hor = zeros(stoch_type, numjumps) for j in 1:numjumps order = sum( - stoch for (spec_idx, stoch) in reactant_stoch[j]; init = zero(stoch_type)) + stoch for (spec_idx, stoch) in reactant_stoch[j]; init = zero(stoch_type) + ) if order > 3 error("Reaction $j has order $order, which is not supported (maximum order is 3).") end @@ -211,19 +286,19 @@ function compute_gi(u, max_hor, max_stoich, i, t) return 2 * one_max_hor else # if max_stoich[i] == 2 return u[i] > one_max_hor ? - 2 * one_max_hor + one_max_hor / (u[i] - one_max_hor) : 2 * one_max_hor # Fallback to 2 if x_i <= 1 + 2 * one_max_hor + one_max_hor / (u[i] - one_max_hor) : 2 * one_max_hor # Fallback to 2 if x_i <= 1 end elseif max_hor[i] == 3 if max_stoich[i] == 1 return 3 * one_max_hor elseif max_stoich[i] == 2 return u[i] > one_max_hor ? - (3 * one_max_hor / 2) * - (2 * one_max_hor + one_max_hor / (u[i] - one_max_hor)) : 3 * one_max_hor # Fallback to 3 if x_i <= 1 + (3 * one_max_hor / 2) * + (2 * one_max_hor + one_max_hor / (u[i] - one_max_hor)) : 3 * one_max_hor # Fallback to 3 if x_i <= 1 else # if max_stoich[i] == 3 return u[i] > 2 * one_max_hor ? - 3 * one_max_hor + one_max_hor / (u[i] - one_max_hor) + - 2 * one_max_hor / (u[i] - 2 * one_max_hor) : 3 * one_max_hor # Fallback to 3 if x_i <= 2 + 3 * one_max_hor + one_max_hor / (u[i] - one_max_hor) + + 2 * one_max_hor / (u[i] - 2 * one_max_hor) : 3 * one_max_hor # Fallback to 3 if x_i <= 2 end end return one_max_hor # Default case @@ -235,7 +310,8 @@ end # mu_i(x) = sum_j nu_ij * a_j(x), sigma_i^2(x) = sum_j nu_ij^2 * a_j(x) # I_rs is the set of reactant species (assumed to be all species here, as critical reactions are not specified). function compute_tau( - u, rate_cache, nu, hor, p, t, epsilon, rate, dtmin, max_hor, max_stoich, numjumps) + u, rate_cache, nu, hor, p, t, epsilon, rate, dtmin, max_hor, max_stoich, numjumps + ) rate(rate_cache, u, p, t) if all(<=(0), rate_cache) # Handle case where all rates are zero or negative return dtmin @@ -270,7 +346,8 @@ function simple_explicit_tau_leaping_loop!( 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) + save_end + ) save_idx = 1 while t_current < t_end @@ -279,11 +356,13 @@ function simple_explicit_tau_leaping_loop!( t_current = t_end break end - tau = compute_tau(u_current, rate_cache, nu, hor, p, t_current, - epsilon, rate, dtmin, max_hor, max_stoich, numjumps) + tau = compute_tau( + u_current, rate_cache, nu, hor, p, t_current, + epsilon, rate, dtmin, max_hor, max_stoich, numjumps + ) tau = min(tau, 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 # Calculate Poisson random numbers only for positive rates @@ -315,7 +394,7 @@ function simple_explicit_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] @@ -328,16 +407,18 @@ function simple_explicit_tau_leaping_loop!( end # Save endpoint if requested and not already saved - if save_end && (isempty(tsave) || tsave[end] != t_end) + return if save_end && (isempty(tsave) || tsave[end] != t_end) push!(usave, copy(u_current)) push!(tsave, t_end) end end -function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; +function DiffEqBase.solve( + jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; seed = nothing, dtmin = nothing, - saveat = nothing, save_start = nothing, save_end = 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.") @@ -346,7 +427,7 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; tspan = prob.tspan if dtmin === nothing - dtmin = 1e-10 * one(typeof(tspan[2])) + dtmin = 1.0e-10 * one(typeof(tspan[2])) end (seed !== nothing) && seed!(rng, seed) @@ -391,29 +472,62 @@ function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; reactant_stoch = maj.reactant_stoch hor = compute_hor(reactant_stoch, numjumps) max_hor, max_stoich = precompute_reaction_conditions( - reactant_stoch, hor, length(u0), numjumps) + reactant_stoch, hor, length(u0), numjumps + ) simple_explicit_tau_leaping_loop!( 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) + save_end + ) - sol = SciMLBase.build_solution(prob, alg, tsave, usave, + sol = SciMLBase.build_solution( + prob, alg, tsave, usave, calculate_error = false, - interp = SciMLBase.ConstantInterpolation(tsave, usave)) + interp = SciMLBase.ConstantInterpolation(tsave, usave) + ) 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 end function EnsembleGPUKernel(backend) - EnsembleGPUKernel(backend, 0.0) + return EnsembleGPUKernel(backend, 0.0) end function EnsembleGPUKernel() - EnsembleGPUKernel(nothing, 0.0) + return EnsembleGPUKernel(nothing, 0.0) end diff --git a/src/spatial/spatial_massaction_jump.jl b/src/spatial/spatial_massaction_jump.jl index c9af9d1e2..cb1be14de 100644 --- a/src/spatial/spatial_massaction_jump.jl +++ b/src/spatial/spatial_massaction_jump.jl @@ -1,32 +1,87 @@ const AVecOrNothing = Union{AbstractVector, Nothing} const AMatOrNothing = Union{AbstractMatrix, Nothing} +""" + SpatialMassActionJump(uniform_rates, spatial_rates, reactant_stoch, net_stoch, + param_mapper = nothing; scale_rates = true, useiszero = true, nocopy = false) + SpatialMassActionJump(spatial_rates, reactant_stoch, net_stoch, param_mapper = nothing; + kwargs...) + SpatialMassActionJump(uniform_rates, reactant_stoch, net_stoch, param_mapper = nothing; + kwargs...) + SpatialMassActionJump(ma_jumps::MassActionJump; scale_rates = false, kwargs...) + +Represent mass-action reactions whose rate constants may be uniform across sites, vary by +site, or include both uniform and spatially varying reaction channels. + +Uniform reactions are ordered before spatially varying reactions. For a spatial rate matrix, +rows index reactions and columns index sites. + +## Arguments + + - `uniform_rates`: Vector of rate constants for reactions that use the same rate at every + site, or `nothing`. + - `spatial_rates`: Matrix of rate constants for site-dependent reactions, or `nothing`. + - `reactant_stoch`: Reactant stoichiometry for each reaction, using the same pair-vector + representation as [`MassActionJump`](@ref). + - `net_stoch`: Net stoichiometry for each reaction. + - `param_mapper`: Optional function mapping problem parameters to rate constants. + - `ma_jumps`: Existing [`MassActionJump`](@ref) to reinterpret as spatial mass-action + reactions. + +## Keyword Arguments + + - `scale_rates`: Whether to divide rates by factorial stoichiometry factors. Defaults to + `true` for raw rates and `false` when constructing from an existing `MassActionJump`. + - `useiszero`: Whether a single zero reactant entry is treated as an empty reactant list. + - `nocopy`: Whether to store input arrays directly instead of copying them. + +## Fields + + - `uniform_rates`: Uniform reaction rates, or `nothing`. + - `spatial_rates`: Site-dependent reaction-rate matrix, or `nothing`. + - `reactant_stoch`: Reactant stoichiometry by reaction. + - `net_stoch`: Net stoichiometry by reaction. + - `param_mapper`: Optional parameter-to-rate mapping. + +## Examples + +```julia +using JumpProcesses + +rates = [0.1 0.2 0.3] +reactant_stoch = [[1 => 1]] +net_stoch = [[1 => -1]] +smaj = SpatialMassActionJump(rates, reactant_stoch, net_stoch) +get_num_majumps(smaj) == 1 +``` +""" struct SpatialMassActionJump{A <: AVecOrNothing, B <: AMatOrNothing, S, U, V} <: - AbstractMassActionJump - uniform_rates::A # reactions that are uniform in space - spatial_rates::B # reactions whose rate depends on the site + AbstractMassActionJump + uniform_rates::A + spatial_rates::B reactant_stoch::S net_stoch::U param_mapper::V - """ - uniform rates go first in ordering - """ - function SpatialMassActionJump{A, B, S, U, V}(uniform_rates::A, spatial_rates::B, + function SpatialMassActionJump{A, B, S, U, V}( + uniform_rates::A, spatial_rates::B, reactant_stoch::S, net_stoch::U, param_mapper::V, scale_rates::Bool, useiszero::Bool, - nocopy::Bool) where {A <: AVecOrNothing, + nocopy::Bool + ) where { + A <: AVecOrNothing, B <: AMatOrNothing, - S, U, V} + S, U, V, + } uniform_rates = (nocopy || isnothing(uniform_rates)) ? uniform_rates : - copy(uniform_rates) + copy(uniform_rates) spatial_rates = (nocopy || isnothing(spatial_rates)) ? spatial_rates : - copy(spatial_rates) + copy(spatial_rates) reactant_stoch = nocopy ? reactant_stoch : copy(reactant_stoch) for i in eachindex(reactant_stoch) if useiszero && (length(reactant_stoch[i]) == 1) && - iszero(reactant_stoch[i][1][1]) + iszero(reactant_stoch[i][1][1]) reactant_stoch[i] = typeof(reactant_stoch[i])() end end @@ -37,102 +92,154 @@ struct SpatialMassActionJump{A <: AVecOrNothing, B <: AMatOrNothing, S, U, V} <: if scale_rates && !isnothing(spatial_rates) && !isempty(spatial_rates) scalerates!(spatial_rates, reactant_stoch[(num_unif_rates + 1):end]) end - new(uniform_rates, spatial_rates, reactant_stoch, net_stoch, param_mapper) + return new(uniform_rates, spatial_rates, reactant_stoch, net_stoch, param_mapper) end end ################ Constructors ################## -function SpatialMassActionJump(urates::A, srates::B, rs::S, ns::U, pmapper::V; +function SpatialMassActionJump( + urates::A, srates::B, rs::S, ns::U, pmapper::V; scale_rates = true, useiszero = true, - nocopy = false) where {A <: AVecOrNothing, - B <: AMatOrNothing, S, U, V} - SpatialMassActionJump{A, B, S, U, V}(urates, srates, rs, ns, pmapper, scale_rates, - useiszero, nocopy) -end -function SpatialMassActionJump(urates::A, srates::B, rs, ns; scale_rates = true, + nocopy = false + ) where { + A <: AVecOrNothing, + B <: AMatOrNothing, S, U, V, + } + return SpatialMassActionJump{A, B, S, U, V}( + urates, srates, rs, ns, pmapper, scale_rates, + useiszero, nocopy + ) +end +function SpatialMassActionJump( + urates::A, srates::B, rs, ns; scale_rates = true, useiszero = true, - nocopy = false) where {A <: AVecOrNothing, - B <: AMatOrNothing} - SpatialMassActionJump(urates, srates, rs, ns, nothing; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy) + nocopy = false + ) where { + A <: AVecOrNothing, + B <: AMatOrNothing, + } + return SpatialMassActionJump( + urates, srates, rs, ns, nothing; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy + ) end -function SpatialMassActionJump(srates::B, rs, ns, pmapper; scale_rates = true, +function SpatialMassActionJump( + srates::B, rs, ns, pmapper; scale_rates = true, useiszero = true, - nocopy = false) where {B <: AMatOrNothing} - SpatialMassActionJump(nothing, srates, rs, ns, pmapper; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy) -end -function SpatialMassActionJump(srates::B, rs, ns; scale_rates = true, useiszero = true, - nocopy = false) where {B <: AMatOrNothing} - SpatialMassActionJump(nothing, srates, rs, ns, nothing; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy) + nocopy = false + ) where {B <: AMatOrNothing} + return SpatialMassActionJump( + nothing, srates, rs, ns, pmapper; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy + ) +end +function SpatialMassActionJump( + srates::B, rs, ns; scale_rates = true, useiszero = true, + nocopy = false + ) where {B <: AMatOrNothing} + return SpatialMassActionJump( + nothing, srates, rs, ns, nothing; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy + ) end -function SpatialMassActionJump(urates::A, rs, ns, pmapper; scale_rates = true, +function SpatialMassActionJump( + urates::A, rs, ns, pmapper; scale_rates = true, useiszero = true, - nocopy = false) where {A <: AVecOrNothing} - SpatialMassActionJump(urates, nothing, rs, ns, pmapper; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy) -end -function SpatialMassActionJump(urates::A, rs, ns; scale_rates = true, useiszero = true, - nocopy = false) where {A <: AVecOrNothing} - SpatialMassActionJump(urates, nothing, rs, ns, nothing; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy) + nocopy = false + ) where {A <: AVecOrNothing} + return SpatialMassActionJump( + urates, nothing, rs, ns, pmapper; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy + ) +end +function SpatialMassActionJump( + urates::A, rs, ns; scale_rates = true, useiszero = true, + nocopy = false + ) where {A <: AVecOrNothing} + return SpatialMassActionJump( + urates, nothing, rs, ns, nothing; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy + ) 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, - useiszero = true, nocopy = false) where {T, S, U, V} - SpatialMassActionJump(ma_jumps.scaled_rates, ma_jumps.reactant_stoch, +function SpatialMassActionJump( + ma_jumps::MassActionJump{T, S, U, V}; scale_rates = false, + useiszero = true, nocopy = false + ) where {T, S, U, V} + return SpatialMassActionJump( + ma_jumps.scaled_rates, ma_jumps.reactant_stoch, ma_jumps.net_stoch, ma_jumps.param_mapper; - scale_rates = scale_rates, useiszero = useiszero, nocopy = nocopy) + scale_rates = scale_rates, useiszero = useiszero, nocopy = nocopy + ) end ############################################## -function get_num_majumps(smaj::SpatialMassActionJump{ - Nothing, Nothing, S, U, V}) where - {S, U, V} - 0 -end -function get_num_majumps(smaj::SpatialMassActionJump{ - Nothing, B, S, U, V}) where - {B, S, U, V} - size(smaj.spatial_rates, 1) -end -function get_num_majumps(smaj::SpatialMassActionJump{ - A, Nothing, S, U, V}) where - {A, S, U, V} - length(smaj.uniform_rates) -end -function get_num_majumps(smaj::SpatialMassActionJump{ - A, B, S, U, V}) where - {A <: AbstractVector, B <: AbstractMatrix, S, U, V} - length(smaj.uniform_rates) + size(smaj.spatial_rates, 1) +function get_num_majumps( + smaj::SpatialMassActionJump{ + Nothing, Nothing, S, U, V, + } + ) where + {S, U, V} + return 0 +end +function get_num_majumps( + smaj::SpatialMassActionJump{ + Nothing, B, S, U, V, + } + ) where + {B, S, U, V} + return size(smaj.spatial_rates, 1) +end +function get_num_majumps( + smaj::SpatialMassActionJump{ + A, Nothing, S, U, V, + } + ) where + {A, S, U, V} + return length(smaj.uniform_rates) +end +function get_num_majumps( + smaj::SpatialMassActionJump{ + A, B, S, U, V, + } + ) where + {A <: AbstractVector, B <: AbstractMatrix, S, U, V} + return length(smaj.uniform_rates) + size(smaj.spatial_rates, 1) end using_params(smaj::SpatialMassActionJump) = false -function rate_at_site(rx, site, - smaj::SpatialMassActionJump{Nothing, B, S, U, V}) where {B, S, U, V} - smaj.spatial_rates[rx, site] -end -function rate_at_site(rx, site, - smaj::SpatialMassActionJump{A, Nothing, S, U, V}) where {A, S, U, V} - smaj.uniform_rates[rx] -end -function rate_at_site(rx, site, - smaj::SpatialMassActionJump{A, B, S, U, V}) where - {A <: AbstractVector, B <: AbstractMatrix, S, U, V} +function rate_at_site( + rx, site, + smaj::SpatialMassActionJump{Nothing, B, S, U, V} + ) where {B, S, U, V} + return smaj.spatial_rates[rx, site] +end +function rate_at_site( + rx, site, + smaj::SpatialMassActionJump{A, Nothing, S, U, V} + ) where {A, S, U, V} + return smaj.uniform_rates[rx] +end +function rate_at_site( + rx, site, + smaj::SpatialMassActionJump{A, B, S, U, V} + ) where + {A <: AbstractVector, B <: AbstractMatrix, S, U, V} num_unif_rxs = length(smaj.uniform_rates) - rx <= num_unif_rxs ? smaj.uniform_rates[rx] : - smaj.spatial_rates[rx - num_unif_rxs, site] + return rx <= num_unif_rxs ? smaj.uniform_rates[rx] : + smaj.spatial_rates[rx - num_unif_rxs, site] end -function evalrxrate(speciesmat::AbstractMatrix{T}, rxidx::S, majump::SpatialMassActionJump, - site::Int) where {T, S} +function evalrxrate( + speciesmat::AbstractMatrix{T}, rxidx::S, majump::SpatialMassActionJump, + site::Int + ) where {T, S} val = one(T) @inbounds for specstoch in majump.reactant_stoch[rxidx] specpop = speciesmat[specstoch[1], site] diff --git a/src/spatial/topology.jl b/src/spatial/topology.jl index bf2d5deb7..3c1d2218d 100644 --- a/src/spatial/topology.jl +++ b/src/spatial/topology.jl @@ -3,6 +3,29 @@ A file with structs and functions for setting up and using the topology of the s """ ################### Graph ######################## +""" + num_sites(spatial_system) -> Int + +Return the number of sites in a spatial system. + +## Arguments + + - `spatial_system`: A `Graphs.AbstractGraph` or [`CartesianGridRej`](@ref)-compatible + grid. + +## Returns + + - The number of graph vertices or Cartesian grid cells. + +## Examples + +```julia +using JumpProcesses + +grid = CartesianGrid((2, 3)) +num_sites(grid) == 6 +``` +""" num_sites(graph::AbstractGraph) = Graphs.nv(graph) # neighbors(graph::AbstractGraph, site) = Graphs.neighbors(graph, site) # outdegree(graph::AbstractGraph, site) = Graphs.outdegree(graph, site) @@ -15,7 +38,7 @@ const offsets_2D = [ CartesianIndex(0, -1), CartesianIndex(-1, 0), CartesianIndex(1, 0), - CartesianIndex(0, 1) + CartesianIndex(0, 1), ] const offsets_3D = [ CartesianIndex(0, 0, -1), @@ -23,7 +46,7 @@ const offsets_3D = [ CartesianIndex(-1, 0, 0), CartesianIndex(1, 0, 0), CartesianIndex(0, 1, 0), - CartesianIndex(0, 0, 1) + CartesianIndex(0, 0, 1), ] """ @@ -43,6 +66,29 @@ end dimension(grid) = length(grid.dims) num_sites(grid) = prod(grid.dims) +""" + outdegree(grid, site) -> Int + +Return the number of valid nearest-neighbor sites adjacent to `site`. + +## Arguments + + - `grid`: A [`CartesianGridRej`](@ref)-compatible grid. + - `site`: Linear site index. + +## Returns + + - The number of in-domain neighbors for `site`. + +## Examples + +```julia +using JumpProcesses + +grid = CartesianGrid((2, 2)) +outdegree(grid, 1) == 2 +``` +""" outdegree(grid, site) = grid.nums_neighbors[site] """ @@ -61,7 +107,7 @@ function nth_nbr(grid, site, n) CI = grid.CI offsets = grid.offsets @inbounds I = CI[site] - @inbounds for off in offsets + return @inbounds for off in offsets nb = I + off if nb in CI n -= 1 @@ -79,7 +125,7 @@ function neighbors(grid, site) CI = grid.CI LI = grid.LI I = CI[site] - Iterators.map(off -> LI[off + I], Iterators.filter(off -> off + I in CI, grid.offsets)) + return Iterators.map(off -> LI[off + I], Iterators.filter(off -> off + I in CI, grid.offsets)) end """ @@ -97,12 +143,73 @@ function pad_hop_vec!(to_pad::AbstractVector, grid, site, hop_vec::AbstractVecto to_pad[i] = zero(eltype(to_pad)) end end - to_pad + return to_pad end -CartesianGrid(dims) = CartesianGridRej(dims) # use CartesianGridRej by default +""" + CartesianGrid(dims) + +Construct the default Cartesian grid topology for spatial jump simulations. + +`CartesianGrid` currently returns a [`CartesianGridRej`](@ref), which samples random +neighbors by rejection from the potential coordinate offsets. + +## Arguments + + - `dims`: Tuple or vector of side lengths. One-, two-, and three-dimensional grids are + supported by the built-in offset tables. + +## Returns + + - A [`CartesianGridRej`](@ref) with sites indexed by Julia linear indices. -# neighbor sampling is rejection-based +## Examples + +```julia +using JumpProcesses + +grid = CartesianGrid((3, 4)) +num_sites(grid) == 12 +collect(neighbors(grid, 1)) == [2, 4] +``` +""" +CartesianGrid(dims) = CartesianGridRej(dims) + +""" + CartesianGridRej(dims) + CartesianGridRej(dimension, linear_size::Int) + +Cartesian grid topology with rejection-based random neighbor sampling. + +Sites are represented by linear indices over `dims`. Neighbor relations use the nearest +Cartesian offsets in one, two, or three dimensions and reject offsets that would leave the +domain. + +## Arguments + + - `dims`: Tuple or vector of side lengths. + - `dimension`: Number of Cartesian dimensions for a hypercube grid. + - `linear_size`: Side length used in every dimension when constructing from + `(dimension, linear_size)`. + +## Fields + + - `dims`: Side lengths of the grid. + - `nums_neighbors`: Number of valid neighbors for each site. + - `CI`: Cartesian indices for the grid domain. + - `LI`: Linear indices for the grid domain. + - `offsets`: Candidate Cartesian offsets used to enumerate or sample neighbors. + +## Examples + +```julia +using JumpProcesses + +grid = CartesianGridRej((2, 2)) +outdegree(grid, 1) == 2 +collect(neighbors(grid, 1)) == [2, 3] +``` +""" struct CartesianGridRej{N, T} "dimensions (side lengths) of the grid" dims::NTuple{N, Int} @@ -127,11 +234,11 @@ function CartesianGridRej(dims::Tuple) LI = LinearIndices(dims) offsets = potential_offsets(dim) nums_neighbors = Int8[count(x -> x + CI[site] in CI, offsets) for site in 1:prod(dims)] - CartesianGridRej(dims, nums_neighbors, CI, LI, offsets) + return CartesianGridRej(dims, nums_neighbors, CI, LI, offsets) end CartesianGridRej(dims) = CartesianGridRej(Tuple(dims)) function CartesianGridRej(dimension, linear_size::Int) - CartesianGridRej([linear_size for i in 1:dimension]) + return CartesianGridRej([linear_size for i in 1:dimension]) end function rand_nbr(rng, grid::CartesianGridRej, site::Int) CI = grid.CI @@ -141,9 +248,12 @@ function rand_nbr(rng, grid::CartesianGridRej, site::Int) @inbounds nb = rand(rng, offsets) + I @inbounds nb in CI && return grid.LI[nb] end + return end -function Base.show(io::IO, ::MIME"text/plain", - grid::CartesianGridRej) - println(io, "A Cartesian grid with dimensions $(grid.dims)") +function Base.show( + io::IO, ::MIME"text/plain", + grid::CartesianGridRej + ) + return println(io, "A Cartesian grid with dimensions $(grid.dims)") end diff --git a/src/variable_rate.jl b/src/variable_rate.jl index 7dff190ae..8abf0140a 100644 --- a/src/variable_rate.jl +++ b/src/variable_rate.jl @@ -58,8 +58,10 @@ sol = solve(jprob, Tsit5()) """ struct VR_FRM <: VariableRateAggregator end -function configure_jump_problem(prob, vr_aggregator::VR_FRM, jumps, cvrjs; - rng = DEFAULT_RNG) +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) return new_prob, variable_jump_callback @@ -84,7 +86,7 @@ function extend_problem(prob::SciMLBase.AbstractODEProblem, jumps; rng = DEFAULT jump_f = let _f = _f function (du::ExtendedJumpArray, u::ExtendedJumpArray, p, t) _f(du.u, u.u, p, t) - update_jumps!(du, u, p, t, length(u.u), jumps...) + return update_jumps!(du, u, p, t, length(u.u), jumps...) end end else @@ -98,9 +100,11 @@ function extend_problem(prob::SciMLBase.AbstractODEProblem, jumps; rng = DEFAULT end u0 = extend_u0(prob, length(jumps), rng) - f = ODEFunction{isinplace(prob)}(jump_f; sys = prob.f.sys, - observed = prob.f.observed) - remake(prob; f, u0) + f = ODEFunction{isinplace(prob)}( + jump_f; sys = prob.f.sys, + observed = prob.f.observed + ) + return remake(prob; f, u0) end function extend_problem(prob::SciMLBase.AbstractSDEProblem, jumps; rng = DEFAULT_RNG) @@ -110,7 +114,7 @@ function extend_problem(prob::SciMLBase.AbstractSDEProblem, jumps; rng = DEFAULT jump_f = let _f = _f function (du::ExtendedJumpArray, u::ExtendedJumpArray, p, t) _f(du.u, u.u, p, t) - update_jumps!(du, u, p, t, length(u.u), jumps...) + return update_jumps!(du, u, p, t, length(u.u), jumps...) end end else @@ -125,18 +129,20 @@ function extend_problem(prob::SciMLBase.AbstractSDEProblem, jumps; rng = DEFAULT if prob.noise_rate_prototype === nothing jump_g = function (du, u, p, t) - prob.g(du.u, u.u, p, t) + return prob.g(du.u, u.u, p, t) end else jump_g = function (du, u, p, t) - prob.g(du, u.u, p, t) + return prob.g(du, u.u, p, t) end end u0 = extend_u0(prob, length(jumps), rng) - f = SDEFunction{isinplace(prob)}(jump_f, jump_g; sys = prob.f.sys, - observed = prob.f.observed) - remake(prob; f, g = jump_g, u0) + f = SDEFunction{isinplace(prob)}( + jump_f, jump_g; sys = prob.f.sys, + observed = prob.f.observed + ) + return remake(prob; f, g = jump_g, u0) end function extend_problem(prob::SciMLBase.AbstractDDEProblem, jumps; rng = DEFAULT_RNG) @@ -146,7 +152,7 @@ function extend_problem(prob::SciMLBase.AbstractDDEProblem, jumps; rng = DEFAULT jump_f = let _f = _f function (du::ExtendedJumpArray, u::ExtendedJumpArray, h, p, t) _f(du.u, u.u, h, p, t) - update_jumps!(du, u, p, t, length(u.u), jumps...) + return update_jumps!(du, u, p, t, length(u.u), jumps...) end end else @@ -160,9 +166,11 @@ function extend_problem(prob::SciMLBase.AbstractDDEProblem, jumps; rng = DEFAULT end u0 = extend_u0(prob, length(jumps), rng) - f = DDEFunction{isinplace(prob)}(jump_f; sys = prob.f.sys, - observed = prob.f.observed) - remake(prob; f, u0) + f = DDEFunction{isinplace(prob)}( + jump_f; sys = prob.f.sys, + observed = prob.f.observed + ) + return remake(prob; f, u0) end # Not sure if the DAE one is correct: Should be a residual of sorts @@ -173,7 +181,7 @@ function extend_problem(prob::SciMLBase.AbstractDAEProblem, jumps; rng = DEFAULT jump_f = let _f = _f function (out, du::ExtendedJumpArray, u::ExtendedJumpArray, h, p, t) _f(out, du.u, u.u, h, p, t) - update_jumps!(out, u, p, t, length(u.u), jumps...) + return update_jumps!(out, u, p, t, length(u.u), jumps...) end end else @@ -187,9 +195,11 @@ function extend_problem(prob::SciMLBase.AbstractDAEProblem, jumps; rng = DEFAULT end u0 = extend_u0(prob, length(jumps), rng) - f = DAEFunction{isinplace(prob)}(jump_f, sys = prob.f.sys, - observed = prob.f.observed) - remake(prob; f, u0) + f = DAEFunction{isinplace(prob)}( + jump_f, sys = prob.f.sys, + observed = prob.f.observed + ) + return remake(prob; f, u0) end struct VR_FRMEventCallback{F, RNG} @@ -205,48 +215,50 @@ end function (c::VR_FRMEventCallback)(integrator) c.affect!(integrator) integrator.u.jump_u[c.idx] = -randexp(c.rng, typeof(integrator.t)) - nothing + return 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)) derivative_discontinuity!(integrator, true) - nothing + return nothing end function wrap_jump_in_callback(idx, jump; rng = DEFAULT_RNG) cb_functor = VR_FRMEventCallback(idx, jump.affect!, rng) - ContinuousCallback(cb_functor, cb_functor; + return ContinuousCallback( + cb_functor, cb_functor; initialize = cb_functor, idxs = jump.idxs, rootfind = jump.rootfind, interp_points = jump.interp_points, save_positions = jump.save_positions, abstol = jump.abstol, - reltol = jump.reltol) + reltol = jump.reltol + ) end function build_variable_callback(cb, idx, jump, jumps...; rng = DEFAULT_RNG) idx += 1 new_cb = wrap_jump_in_callback(idx, jump; rng) - build_variable_callback(CallbackSet(cb, new_cb), idx, jumps...; rng = DEFAULT_RNG) + return build_variable_callback(CallbackSet(cb, new_cb), idx, jumps...; rng = DEFAULT_RNG) end function build_variable_callback(cb, idx, jump; rng = DEFAULT_RNG) idx += 1 - CallbackSet(cb, wrap_jump_in_callback(idx, jump; rng)) + return CallbackSet(cb, wrap_jump_in_callback(idx, jump; rng)) end @inline function update_jumps!(du, u, p, t, idx, jump) idx += 1 - du[idx] = jump.rate(u.u, p, t) + return du[idx] = jump.rate(u.u, p, t) end @inline function update_jumps!(du, u, p, t, idx, jump, jumps...) idx += 1 du[idx] = jump.rate(u.u, p, t) - update_jumps!(du, u, p, t, idx, jumps...) + return update_jumps!(du, u, p, t, idx, jumps...) end ################################### VR_Direct and VR_DirectFW #################################### @@ -295,6 +307,43 @@ sol = solve(jprob, Tsit5()) - `VR_Direct` and `VR_DirectFW` are expected to generally be more performant than `VR_FRM`. """ struct VR_Direct <: VariableRateAggregator end + +""" +$(TYPEDEF) + +Function-wrapper variant of [`VR_Direct`](@ref) for simulations with many +[`VariableRateJump`](@ref)s. + +`VR_DirectFW` uses the same direct-method callback strategy as `VR_Direct`, but stores the +rate and affect functions in `FunctionWrappers`-based containers. + +## Returns + + - A stateless [`VariableRateAggregator`](@ref) value for the `vr_aggregator` keyword of + [`JumpProblem`](@ref). + +## Examples + +```julia +using JumpProcesses, OrdinaryDiffEq + +u0 = [1.0] +p = [10.0, 0.5] +tspan = (0.0, 10.0) + +birth_rate(u, p, t) = p[1] +birth_affect!(integrator) = (integrator.u[1] += 1; nothing) +birth_jump = VariableRateJump(birth_rate, birth_affect!) + +death_rate(u, p, t) = p[2] * u[1] +death_affect!(integrator) = (integrator.u[1] -= 1; nothing) +death_jump = VariableRateJump(death_rate, death_affect!) + +oprob = ODEProblem((du, u, p, t) -> du .= 0, u0, tspan, p) +jprob = JumpProblem(oprob, birth_jump, death_jump; vr_aggregator = VR_DirectFW()) +sol = solve(jprob, Tsit5()) +``` +""" struct VR_DirectFW <: VariableRateAggregator end mutable struct VR_DirectEventCache{T, RNG, F1, F2} @@ -310,7 +359,8 @@ mutable struct VR_DirectEventCache{T, RNG, F1, F2} end function VR_DirectEventCache( - jumps::JumpSet, ::VR_Direct, prob, ::Type{T}; rng = DEFAULT_RNG) where {T} + jumps::JumpSet, ::VR_Direct, prob, ::Type{T}; rng = DEFAULT_RNG + ) where {T} initial_threshold = randexp(rng, T) vjumps = jumps.variable_jumps @@ -319,13 +369,16 @@ function VR_DirectEventCache( cum_rate_sum = Vector{T}(undef, length(vjumps)) - VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), typeof(affect_funcs)}(zero(T), + return VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), typeof(affect_funcs)}( + zero(T), initial_threshold, zero(T), initial_threshold, zero(T), rng, rate_funcs, - affect_funcs, cum_rate_sum) + affect_funcs, cum_rate_sum + ) end function VR_DirectEventCache( - jumps::JumpSet, ::VR_DirectFW, prob, ::Type{T}; rng = DEFAULT_RNG) where {T} + jumps::JumpSet, ::VR_DirectFW, prob, ::Type{T}; rng = DEFAULT_RNG + ) where {T} initial_threshold = randexp(rng, T) vjumps = jumps.variable_jumps @@ -336,9 +389,11 @@ function VR_DirectEventCache( cum_rate_sum = Vector{T}(undef, length(vjumps)) - VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), Any}(zero(T), + return VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), Any}( + zero(T), initial_threshold, zero(T), initial_threshold, zero(T), rng, rate_funcs, - affect_funcs, cum_rate_sum) + affect_funcs, cum_rate_sum + ) end # Initialization function for VR_DirectEventCache @@ -349,23 +404,29 @@ function initialize_vr_direct_cache!(cache::VR_DirectEventCache, u, t, integrato cache.current_threshold = cache.prev_threshold cache.total_rate = zero(integrator.t) cache.cum_rate_sum .= 0 - nothing + return nothing end -@inline function concretize_vr_direct_affects!(cache::VR_DirectEventCache, - ::I) where {I <: SciMLBase.DEIntegrator} +@inline function concretize_vr_direct_affects!( + cache::VR_DirectEventCache, + ::I + ) where {I <: SciMLBase.DEIntegrator} if (cache.affect_funcs isa Vector) && - !(cache.affect_funcs isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}}) + !(cache.affect_funcs isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}}) AffectWrapper = FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}} - cache.affect_funcs = AffectWrapper[makewrapper(AffectWrapper, aff) - for aff in cache.affect_funcs] + cache.affect_funcs = AffectWrapper[ + makewrapper(AffectWrapper, aff) + for aff in cache.affect_funcs + ] end - nothing + return 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} - nothing +@inline function concretize_vr_direct_affects!( + cache::VR_DirectEventCache{T, RNG, F1, F2}, + ::I + ) where {T, RNG, F1, F2 <: Tuple, I <: SciMLBase.DEIntegrator} + return nothing end # Wrapper for initialize to match ContinuousCallback signature @@ -373,7 +434,7 @@ function initialize_vr_direct_wrapper(cb::ContinuousCallback, u, t, integrator) concretize_vr_direct_affects!(cb.condition, integrator) initialize_vr_direct_cache!(cb.condition, u, t, integrator) derivative_discontinuity!(integrator, false) - nothing + return nothing end # Merge callback parameters across all jumps for VR_Direct @@ -388,8 +449,10 @@ function build_variable_integcallback(cache::VR_DirectEventCache, jumps) reltol = min(reltol, jump.reltol) end - return ContinuousCallback(cache, cache; initialize = initialize_vr_direct_wrapper, - save_positions, abstol, reltol) + return ContinuousCallback( + cache, cache; initialize = initialize_vr_direct_wrapper, + save_positions, abstol, reltol + ) end function configure_jump_problem(prob, ::VR_Direct, jumps, cvrjs; rng = DEFAULT_RNG) @@ -409,7 +472,7 @@ end # recursively evaluate the cumulative sum of the rates for type stability @inline function cumsum_rates!(cum_rate_sum, u, p, t, rates::Tuple) cur_sum = zero(eltype(cum_rate_sum)) - cumsum_rates!(cum_rate_sum, u, p, t, 1, cur_sum, rates...) + return cumsum_rates!(cum_rate_sum, u, p, t, 1, cur_sum, rates...) end # loop-based version for Vector rate_funcs (avoids dynamic splatting) @@ -426,16 +489,18 @@ end new_sum = cur_sum + rate(u, p, t) @inbounds cum_rate_sum[idx] = new_sum idx += 1 - cumsum_rates!(cum_rate_sum, u, p, t, idx, new_sum, rates...) + return cumsum_rates!(cum_rate_sum, u, p, t, idx, new_sum, rates...) end @inline function cumsum_rates!(cum_rate_sum, u, p, t, idx, cur_sum, rate) - @inbounds cum_rate_sum[idx] = cur_sum + rate(u, p, t) + return @inbounds cum_rate_sum[idx] = cur_sum + rate(u, p, t) end function total_variable_rate( cache::VR_DirectEventCache{ - T, RNG, F1, F2}, u, p, t) where {T, RNG, F1, F2} + T, RNG, F1, F2, + }, u, p, t + ) where {T, RNG, F1, F2} (; cum_rate_sum, rate_funcs) = cache sum_rate = cumsum_rates!(cum_rate_sum, u, p, t, rate_funcs) return sum_rate @@ -480,18 +545,22 @@ 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} - quote +@generated function execute_affect!( + cache::VR_DirectEventCache{T, RNG, F1, F2}, + integrator::I, idx + ) where {T, RNG, F1, F2 <: Tuple, I <: SciMLBase.DEIntegrator} + return 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)) end end -@inline function execute_affect!(cache::VR_DirectEventCache, - integrator::I, idx) where {I <: SciMLBase.DEIntegrator} +@inline function execute_affect!( + cache::VR_DirectEventCache, + integrator::I, idx + ) where {I <: SciMLBase.DEIntegrator} (; affect_funcs) = cache - if affect_funcs isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}} + return if affect_funcs isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}} @inbounds affect_funcs[idx](integrator) else error("Error, invalid affect_funcs type. Expected a vector of function wrappers and got $(typeof(affect_funcs))") diff --git a/test/public_api_docs.jl b/test/public_api_docs.jl new file mode 100644 index 000000000..d7bbabb40 --- /dev/null +++ b/test/public_api_docs.jl @@ -0,0 +1,64 @@ +using JumpProcesses, Test + +const OWNED_PUBLIC_API = ( + :BracketData, + :CCNRM, + :CartesianGrid, + :CartesianGridRej, + :Coevolve, + :ConstantRateJump, + :Direct, + :DirectCR, + :DirectCRDirect, + :DirectFW, + :EnsembleGPUKernel, + :ExtendedJumpArray, + :FRM, + :FRMFW, + :JumpProblem, + :JumpSet, + :MassActionJump, + :NRM, + :NSM, + :PureLeaping, + :RDirect, + :RSSA, + :RSSACR, + :RegularJump, + :SSAStepper, + :SimpleExplicitTauLeaping, + :SimpleTauLeaping, + :SortingDirect, + :SpatialMassActionJump, + :SplitCoupledJumpProblem, + :VR_Direct, + :VR_DirectFW, + :VR_FRM, + :VariableRateAggregator, + :VariableRateJump, + :get_num_majumps, + :needs_depgraph, + :needs_vartojumps_map, + :neighbors, + :num_sites, + :outdegree, + :reset_aggregated_jumps!, +) + +const API_PAGE = read(joinpath(@__DIR__, "..", "docs", "src", "api.md"), String) + +@testset "owned public API has docstrings" begin + missing_docstrings = Symbol[] + for name in OWNED_PUBLIC_API + Docs.hasdoc(JumpProcesses, name) || push!(missing_docstrings, name) + end + @test missing_docstrings == Symbol[] +end + +@testset "owned public API is rendered in docs" begin + missing_entries = Symbol[] + for name in OWNED_PUBLIC_API + occursin("\n$(name)\n", API_PAGE) || push!(missing_entries, name) + end + @test missing_entries == Symbol[] +end diff --git a/test/runtests.jl b/test/runtests.jl index 9a2dfc246..0478e6506 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,6 +12,7 @@ end @time begin if GROUP == "QA" @time @safetestset "QA Tests" begin include("qa.jl") end + @time @safetestset "Public API Documentation" begin include("public_api_docs.jl") end end if GROUP == "All" || GROUP == "InterfaceI" From 1e44483a2a5043887018ca410671fbfd3900637f Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sat, 11 Jul 2026 02:00:19 -0400 Subject: [PATCH 2/5] Update SciMLTesting QA compat Co-Authored-By: Chris Rackauckas --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 8f737df06..f9b9d5f87 100644 --- a/Project.toml +++ b/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" From f0e4c3afdacfd0960caa13b16d1bd9d0e5817739 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sat, 11 Jul 2026 07:48:30 -0400 Subject: [PATCH 3/5] Use shared rendered API docs QA Co-Authored-By: Chris Rackauckas --- test/public_api_docs.jl | 64 ----------------------------------------- test/qa.jl | 9 +++++- test/runtests.jl | 1 - 3 files changed, 8 insertions(+), 66 deletions(-) delete mode 100644 test/public_api_docs.jl diff --git a/test/public_api_docs.jl b/test/public_api_docs.jl deleted file mode 100644 index d7bbabb40..000000000 --- a/test/public_api_docs.jl +++ /dev/null @@ -1,64 +0,0 @@ -using JumpProcesses, Test - -const OWNED_PUBLIC_API = ( - :BracketData, - :CCNRM, - :CartesianGrid, - :CartesianGridRej, - :Coevolve, - :ConstantRateJump, - :Direct, - :DirectCR, - :DirectCRDirect, - :DirectFW, - :EnsembleGPUKernel, - :ExtendedJumpArray, - :FRM, - :FRMFW, - :JumpProblem, - :JumpSet, - :MassActionJump, - :NRM, - :NSM, - :PureLeaping, - :RDirect, - :RSSA, - :RSSACR, - :RegularJump, - :SSAStepper, - :SimpleExplicitTauLeaping, - :SimpleTauLeaping, - :SortingDirect, - :SpatialMassActionJump, - :SplitCoupledJumpProblem, - :VR_Direct, - :VR_DirectFW, - :VR_FRM, - :VariableRateAggregator, - :VariableRateJump, - :get_num_majumps, - :needs_depgraph, - :needs_vartojumps_map, - :neighbors, - :num_sites, - :outdegree, - :reset_aggregated_jumps!, -) - -const API_PAGE = read(joinpath(@__DIR__, "..", "docs", "src", "api.md"), String) - -@testset "owned public API has docstrings" begin - missing_docstrings = Symbol[] - for name in OWNED_PUBLIC_API - Docs.hasdoc(JumpProcesses, name) || push!(missing_docstrings, name) - end - @test missing_docstrings == Symbol[] -end - -@testset "owned public API is rendered in docs" begin - missing_entries = Symbol[] - for name in OWNED_PUBLIC_API - occursin("\n$(name)\n", API_PAGE) || push!(missing_entries, name) - end - @test missing_entries == Symbol[] -end diff --git a/test/qa.jl b/test/qa.jl index 685ea29df..68860c3c1 100644 --- a/test/qa.jl +++ b/test/qa.jl @@ -1,4 +1,10 @@ -using SciMLTesting, JumpProcesses, Test +using SciMLTesting, JumpProcesses + +const REEXPORTED_API = Tuple( + name for name in names(JumpProcesses; all = false) if name !== :JumpProcesses && + isdefined(JumpProcesses, name) && + parentmodule(getfield(JumpProcesses, name)) !== JumpProcesses +) # The ExplicitImports ignore-lists below are names owned by other packages whose # released public API does not (yet) include them; each group is annotated with its @@ -7,6 +13,7 @@ using SciMLTesting, JumpProcesses, Test run_qa( JumpProcesses; explicit_imports = true, + api_docs_kwargs = (; rendered = true, rendered_ignore = REEXPORTED_API), aqua_kwargs = (; ambiguities = false, # TODO: fix ambiguities and enable piracies = false, # default solvers defined for AbstractJumpProblem diff --git a/test/runtests.jl b/test/runtests.jl index 0478e6506..9a2dfc246 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -12,7 +12,6 @@ end @time begin if GROUP == "QA" @time @safetestset "QA Tests" begin include("qa.jl") end - @time @safetestset "Public API Documentation" begin include("public_api_docs.jl") end end if GROUP == "All" || GROUP == "InterfaceI" From 23251da5b7b517e3e3daa9736bb8c87c013c5d6c Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sun, 12 Jul 2026 02:26:09 -0400 Subject: [PATCH 4/5] Remove unrelated formatting changes Co-Authored-By: Chris Rackauckas --- src/JumpProcesses.jl | 12 +- src/aggregators/aggregators.jl | 12 +- src/aggregators/bracketing.jl | 29 ++- src/jumps.jl | 314 +++++++++---------------- src/simple_regular_solve.jl | 75 +++--- src/spatial/spatial_massaction_jump.jl | 221 +++++++---------- src/spatial/topology.jl | 26 +- src/variable_rate.jl | 147 +++++------- 8 files changed, 323 insertions(+), 513 deletions(-) diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl index ce270344f..353072823 100644 --- a/src/JumpProcesses.jl +++ b/src/JumpProcesses.jl @@ -36,10 +36,10 @@ import SymbolicIndexingInterface as SII # Import additional types and functions from DiffEqBase and SciMLBase using DiffEqBase: DiffEqBase, CallbackSet, ContinuousCallback, DAEFunction, - DDEFunction, DiscreteProblem, ODEFunction, ODEProblem, - ODESolution, ReturnCode, SDEFunction, SDEProblem, add_tstop!, - deleteat!, isinplace, remake, savevalues!, step!, - derivative_discontinuity! + DDEFunction, DiscreteProblem, ODEFunction, ODEProblem, + ODESolution, ReturnCode, SDEFunction, SDEProblem, add_tstop!, + deleteat!, isinplace, remake, savevalues!, step!, + derivative_discontinuity! using SciMLBase: SciMLBase, DEIntegrator abstract type AbstractJump end @@ -47,7 +47,7 @@ abstract type AbstractMassActionJump <: AbstractJump end abstract type AbstractAggregatorAlgorithm end abstract type AbstractJumpAggregator end abstract type AbstractSSAIntegrator{Alg, IIP, U, T} <: -DEIntegrator{Alg, IIP, U, T} end + DEIntegrator{Alg, IIP, U, T} end const DEFAULT_RNG = Random.default_rng() @@ -156,7 +156,7 @@ export init, solve, solve! include("SSA_stepper.jl") export SSAStepper -# leaping: +# leaping: include("simple_regular_solve.jl") export SimpleTauLeaping, SimpleExplicitTauLeaping, EnsembleGPUKernel diff --git a/src/aggregators/aggregators.jl b/src/aggregators/aggregators.jl index 1e3dcd503..ebaa13fc6 100644 --- a/src/aggregators/aggregators.jl +++ b/src/aggregators/aggregators.jl @@ -172,10 +172,8 @@ algorithm with optimal binning, Journal of Chemical Physics 143, 074108 """ struct DirectCRDirect <: AbstractAggregatorAlgorithm end -const JUMP_AGGREGATORS = ( - Direct(), DirectFW(), DirectCR(), SortingDirect(), RSSA(), FRM(), - FRMFW(), NRM(), RSSACR(), RDirect(), Coevolve(), CCNRM(), -) +const JUMP_AGGREGATORS = (Direct(), DirectFW(), DirectCR(), SortingDirect(), RSSA(), FRM(), + FRMFW(), NRM(), RSSACR(), RDirect(), Coevolve(), CCNRM()) # For JumpProblem construction without an aggregator struct NullAggregator <: AbstractAggregatorAlgorithm end @@ -257,11 +255,9 @@ is_spatial(aggregator::NSM) = true is_spatial(aggregator::DirectCRDirect) = true # return the fastest aggregator out of the available ones -function select_aggregator( - jumps::JumpSet; vartojumps_map = nothing, +function select_aggregator(jumps::JumpSet; vartojumps_map = nothing, jumptovars_map = nothing, dep_graph = nothing, spatial_system = nothing, - hopping_constants = nothing - ) + hopping_constants = nothing) # detect if a spatial SSA should be used !isnothing(spatial_system) && !isnothing(hopping_constants) && return DirectCRDirect diff --git a/src/aggregators/bracketing.jl b/src/aggregators/bracketing.jl index e6abfdd03..d591444dd 100644 --- a/src/aggregators/bracketing.jl +++ b/src/aggregators/bracketing.jl @@ -1,3 +1,8 @@ +# bracketing intervals for RSSA-based aggregators +# see: "On the rejection-based algorithm for simulation and analysis of +# 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}() @@ -31,9 +36,9 @@ bd.fluctrate == 0.1 ``` """ struct BracketData{T1, T2} - fluctrate::T1 - threshold::T2 - Δu::T2 + fluctrate::T1 # interval should be [1-fluctrate,1+fluctrate] * u + threshold::T2 # for u below threshold interval is: + Δu::T2 # [max(u-Δu,0),u+Δu] end # # suggested defaults @@ -53,7 +58,7 @@ BracketData{T1, T2}() where {T1, T2} = BracketData(T1(0.1), T2(25), T2(4)) @inline getΔu(bd::BracketData{T1, T2}, i) where {T1, T2 <: Number} = bd.Δu @inline function delta_bracket(u::Integer, δ) - return (trunc(typeof(u), (one(δ) - δ) * u), trunc(typeof(u), (one(δ) + δ) * u)) + (trunc(typeof(u), (one(δ) - δ) * u), trunc(typeof(u), (one(δ) + δ) * u)) end @inline delta_bracket(u, δ) = ((one(δ) - δ) * u), ((one(δ) + δ) * u) @@ -75,7 +80,7 @@ end # Get propensity brackets of massaction jump k. @inline function get_majump_brackets(ulow, uhigh, k, majumps) - return evalrxrate(ulow, k, majumps), evalrxrate(uhigh, k, majumps) + evalrxrate(ulow, k, majumps), evalrxrate(uhigh, k, majumps) end # for constant rate jumps we must check the ordering of the bracket values @@ -95,10 +100,8 @@ get brackets for the rate of reaction rx by first checking if the reaction is a if rx <= num_majumps return get_majump_brackets(p.ulow, p.uhigh, rx, ma_jumps) else - @inbounds return get_cjump_brackets( - p.ulow, p.uhigh, p.rates[rx - num_majumps], - params, t - ) + @inbounds return get_cjump_brackets(p.ulow, p.uhigh, p.rates[rx - num_majumps], + params, t) end end @@ -108,7 +111,7 @@ end @inbounds for (i, uval) in enumerate(u) ulow[i], uhigh[i] = get_spec_brackets(p.bracket_data, i, uval) end - return nothing + nothing end @inline function update_u_brackets!(p::AbstractSSAJumpAggregator, u::SVector) @@ -117,12 +120,12 @@ end p.ulow = setindex(p.ulow, ulow, i) p.uhigh = setindex(p.uhigh, uhigh, i) end - return nothing + nothing end # For ExtendedJumpArray, only iterate over the species portion (u.u), not the jump tracking portion @inline function update_u_brackets!(p::AbstractSSAJumpAggregator, u::ExtendedJumpArray) - return update_u_brackets!(p, u.u) + update_u_brackets!(p, u.u) end # Set up bracketing. The aggregator must have fields @@ -151,5 +154,5 @@ function set_bracketing!(p::AbstractSSAJumpAggregator, u, params, t) end p.sum_rate = sum_rate - return nothing + nothing end diff --git a/src/jumps.jl b/src/jumps.jl index 013b8ba4e..f04080d6e 100644 --- a/src/jumps.jl +++ b/src/jumps.jl @@ -171,17 +171,15 @@ function VariableRateJump(rate, affect!; lrate = nothing, urate = nothing, abstol = 1e-12, reltol = 0) ``` """ -function VariableRateJump( - rate, affect!; +function VariableRateJump(rate, affect!; lrate = nothing, urate = nothing, rateinterval = nothing, rootfind = true, idxs = nothing, save_positions = (false, true), interp_points = 10, - abstol = 1.0e-12, reltol = 0 - ) + abstol = 1e-12, reltol = 0) if !(urate !== nothing && rateinterval !== nothing) && - !(urate === nothing && rateinterval === nothing) + !(urate === nothing && rateinterval === nothing) error("`urate` and `rateinterval` must both be `nothing`, or must both be defined.") end @@ -190,10 +188,8 @@ function VariableRateJump( error("If a lower bound rate, `lrate`, is given than an upper bound rate, `urate`, and rate interval, `rateinterval`, must also be provided.") end - return VariableRateJump( - rate, affect!, lrate, urate, rateinterval, idxs, rootfind, - interp_points, save_positions, abstol, reltol - ) + VariableRateJump(rate, affect!, lrate, urate, rateinterval, idxs, rootfind, + interp_points, save_positions, abstol, reltol) end """ @@ -245,14 +241,14 @@ struct RegularJump{iip, R, C, MD} """ A distribution for marks. Not currently used or supported. """ mark_dist::MD function RegularJump{iip}(rate, c, numjumps::Int; mark_dist = nothing) where {iip} - return new{iip, typeof(rate), typeof(c), typeof(mark_dist)}(rate, c, numjumps, mark_dist) + new{iip, typeof(rate), typeof(c), typeof(mark_dist)}(rate, c, numjumps, mark_dist) end end DiffEqBase.isinplace(::RegularJump{iip, R, C, MD}) where {iip, R, C, MD} = iip function RegularJump(rate, c, numjumps::Int; kwargs...) - return RegularJump{DiffEqBase.isinplace(rate, 4)}(rate, c, numjumps; kwargs...) + RegularJump{DiffEqBase.isinplace(rate, 4)}(rate, c, numjumps; kwargs...) end # deprecate old call @@ -260,9 +256,9 @@ function RegularJump(rate, c, dc::AbstractMatrix; constant_c = false, mark_dist @warn("The RegularJump interface has changed to be matrix-free. See the documentation for more details.") function _c(du, u, p, t, counts, mark) c(dc, u, p, t, mark) - return mul!(du, dc, counts) + mul!(du, dc, counts) end - return RegularJump{true}(rate, _c, size(dc, 2); mark_dist = mark_dist) + RegularJump{true}(rate, _c, size(dc, 2); mark_dist = mark_dist) end """ @@ -340,12 +336,10 @@ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump """Whether `update_parameters!` should apply stoichiometric scaling to rates.""" rescale_rates_on_update::Bool - function MassActionJump{T, S, U, V}( - rates::T, rs_in::S, ns::U, pmapper::V, + 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) @@ -357,91 +351,67 @@ struct MassActionJump{T, S, U, V} <: AbstractMassActionJump if scale_rates && !isempty(sr) scalerates!(sr, rs) end - return new(sr, rs, ns, pmapper, rescale_rates_on_update) + new(sr, rs, ns, pmapper, rescale_rates_on_update) end - function MassActionJump{ - Nothing, Vector{S}, - Vector{U}, V, - }( - ::Nothing, rs_in::Vector{S}, + function MassActionJump{Nothing, Vector{S}, + Vector{U}, V}(::Nothing, rs_in::Vector{S}, ns::Vector{U}, pmapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool, - rescale_rates_on_update::Bool = scale_rates - ) where { - S <: AbstractVector, - U <: AbstractVector, V, - } + rescale_rates_on_update::Bool = scale_rates) where {S <: AbstractVector, + U <: AbstractVector, V} rs = nocopy ? rs_in : copy(rs_in) for i in eachindex(rs) if useiszero && (length(rs[i]) == 1) && iszero(rs[i][1][1]) rs[i] = typeof(rs[i])() end end - return new(nothing, rs, ns, pmapper, rescale_rates_on_update) + new(nothing, rs, ns, pmapper, rescale_rates_on_update) end - function MassActionJump{T, S, U, V}( - rate::T, rs_in::S, ns::U, pmapper::V, + function MassActionJump{T, S, U, V}(rate::T, rs_in::S, ns::U, pmapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool, - rescale_rates_on_update::Bool = scale_rates - ) where {T <: Number, S, U, V} + rescale_rates_on_update::Bool = scale_rates) where {T <: Number, S, U, V} rs = rs_in if useiszero && (length(rs) == 1) && iszero(rs[1][1]) rs = typeof(rs)() end sr = scale_rates ? scalerate(rate, rs) : rate - return new(sr, rs, ns, pmapper, rescale_rates_on_update) + new(sr, rs, ns, pmapper, rescale_rates_on_update) end - function MassActionJump{Nothing, S, U, V}( - ::Nothing, rs_in::S, ns::U, pmapper::V, + function MassActionJump{Nothing, S, U, V}(::Nothing, rs_in::S, ns::U, pmapper::V, scale_rates::Bool, useiszero::Bool, nocopy::Bool, - rescale_rates_on_update::Bool = scale_rates - ) where {S, U, V} + rescale_rates_on_update::Bool = scale_rates) where {S, U, V} rs = rs_in if useiszero && (length(rs) == 1) && iszero(rs[1][1]) rs = typeof(rs)() end - return new(nothing, rs, ns, pmapper, rescale_rates_on_update) + new(nothing, rs, ns, pmapper, rescale_rates_on_update) end end -function MassActionJump( - usr::T, rs::S, ns::U, pmapper::V; scale_rates = true, +function MassActionJump(usr::T, rs::S, ns::U, pmapper::V; scale_rates = true, useiszero = true, nocopy = false, - rescale_rates_on_update = scale_rates - ) where {T, S, U, V} - return MassActionJump{T, S, U, V}( - usr, rs, ns, pmapper, scale_rates, useiszero, nocopy, - rescale_rates_on_update - ) -end -function MassActionJump( - usr::T, rs, ns; scale_rates = true, useiszero = true, - nocopy = false, rescale_rates_on_update = scale_rates - ) where {T <: AbstractVector} - return MassActionJump( - usr, rs, ns, nothing; scale_rates, useiszero, nocopy, - rescale_rates_on_update - ) -end -function MassActionJump( - usr::T, rs, ns; scale_rates = true, useiszero = true, - nocopy = false, rescale_rates_on_update = scale_rates - ) where {T <: Number} - return MassActionJump( - usr, rs, ns, nothing; scale_rates, useiszero, nocopy, - rescale_rates_on_update - ) + rescale_rates_on_update = scale_rates) where {T, S, U, V} + MassActionJump{T, S, U, V}(usr, rs, ns, pmapper, scale_rates, useiszero, nocopy, + rescale_rates_on_update) +end +function MassActionJump(usr::T, rs, ns; scale_rates = true, useiszero = true, + nocopy = false, rescale_rates_on_update = scale_rates) where {T <: AbstractVector} + MassActionJump(usr, rs, ns, nothing; scale_rates, useiszero, nocopy, + rescale_rates_on_update) +end +function MassActionJump(usr::T, rs, ns; scale_rates = true, useiszero = true, + nocopy = false, rescale_rates_on_update = scale_rates) where {T <: Number} + MassActionJump(usr, rs, ns, nothing; scale_rates, useiszero, nocopy, + rescale_rates_on_update) end # with parameter indices or mapping, multiple jump case -function MassActionJump( - rs, ns; param_idxs = nothing, param_mapper = nothing, +function MassActionJump(rs, ns; param_idxs = nothing, param_mapper = nothing, scale_rates = true, useiszero = true, nocopy = false, - rescale_rates_on_update = scale_rates - ) + rescale_rates_on_update = scale_rates) if param_mapper === nothing (param_idxs === nothing) && error("If no parameter indices are given via param_idxs, an explicit parameter mapping must be passed in via param_mapper.") @@ -452,10 +422,8 @@ function MassActionJump( pmapper = param_mapper end - return MassActionJump( - nothing, nocopy ? rs : copy(rs), ns, pmapper; scale_rates, useiszero, - nocopy = true, rescale_rates_on_update - ) + MassActionJump(nothing, nocopy ? rs : copy(rs), ns, pmapper; scale_rates, useiszero, + nocopy = true, rescale_rates_on_update) end using_params(maj::MassActionJump{T, S, U, Nothing}) where {T, S, U} = false @@ -495,53 +463,43 @@ end # create the initial parameter vector for use in a MassActionJump # Note these are unscaled function (ratemap::MassActionJumpParamMapper{U})(params) where {U <: AbstractArray} - return [params[pidx] for pidx in ratemap.param_idxs] + [params[pidx] for pidx in ratemap.param_idxs] end # Note this is unscaled function (ratemap::MassActionJumpParamMapper{U})(params) where {U <: Int} - return params[ratemap.param_idxs] + params[ratemap.param_idxs] end # update a maj with parameter vectors -function (ratemap::MassActionJumpParamMapper{U})( - maj::MassActionJump, newparams; +function (ratemap::MassActionJumpParamMapper{U})(maj::MassActionJump, newparams; scale_rates, - kwargs... - ) where {U <: AbstractArray} + kwargs...) where {U <: AbstractArray} for i in 1:get_num_majumps(maj) maj.scaled_rates[i] = newparams[ratemap.param_idxs[i]] end scale_rates && scalerates!(maj.scaled_rates, maj.reactant_stoch) - return nothing + nothing end function to_collection(ratemap::MassActionJumpParamMapper{Int}) - return MassActionJumpParamMapper([ratemap.param_idxs]) + MassActionJumpParamMapper([ratemap.param_idxs]) end -function Base.merge!( - pmap1::MassActionJumpParamMapper{U}, - pmap2::MassActionJumpParamMapper{U} - ) where {U <: AbstractVector} - return append!(pmap1.param_idxs, pmap2.param_idxs) +function Base.merge!(pmap1::MassActionJumpParamMapper{U}, + pmap2::MassActionJumpParamMapper{U}) where {U <: AbstractVector} + append!(pmap1.param_idxs, pmap2.param_idxs) end -function Base.merge!( - pmap1::MassActionJumpParamMapper{U}, - pmap2::MassActionJumpParamMapper{V} - ) where { - U <: AbstractVector, - V <: Int, - } - return push!(pmap1.param_idxs, pmap2.param_idxs) +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} - ) - return MassActionJumpParamMapper([pmap1.param_idxs, pmap2.param_idxs]) +function Base.merge(pmap1::MassActionJumpParamMapper{Int}, + pmap2::MassActionJumpParamMapper{Int}) + MassActionJumpParamMapper([pmap1.param_idxs, pmap2.param_idxs]) end """ @@ -560,7 +518,7 @@ Notes: 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.") - return maj.param_mapper(maj, newparams; scale_rates, kwargs) + maj.param_mapper(maj, newparams; scale_rates, kwargs) end """ @@ -608,24 +566,22 @@ struct JumpSet{T1, T2, T3, T4} <: AbstractJump massaction_jump::T4 end function JumpSet(vj, cj, rj, maj::MassActionJump{S, T, U, V}) where {S <: Number, T, U, V} - return JumpSet(vj, cj, rj, check_majump_type(maj)) + JumpSet(vj, cj, rj, check_majump_type(maj)) end JumpSet(jump::ConstantRateJump) = JumpSet((), (jump,), nothing, nothing) JumpSet(jump::VariableRateJump) = JumpSet((jump,), (), nothing, nothing) JumpSet(jump::RegularJump) = JumpSet((), (), jump, nothing) JumpSet(jump::AbstractMassActionJump) = JumpSet((), (), nothing, jump) -function JumpSet(; - variable_jumps = (), constant_jumps = (), - regular_jumps = nothing, massaction_jumps = nothing - ) - return JumpSet(variable_jumps, constant_jumps, regular_jumps, massaction_jumps) +function JumpSet(; variable_jumps = (), constant_jumps = (), + regular_jumps = nothing, massaction_jumps = nothing) + JumpSet(variable_jumps, constant_jumps, regular_jumps, massaction_jumps) end JumpSet(jb::Nothing) = JumpSet() # For Varargs, use recursion to make it type-stable function JumpSet(jumps::AbstractJump...) - return JumpSet(split_jumps((), (), nothing, nothing, jumps...)...) + JumpSet(split_jumps((), (), nothing, nothing, jumps...)...) end # handle vector of mass action jumps @@ -638,34 +594,32 @@ function JumpSet(vjs, cjs, rj, majv::Vector{T}) where {T <: MassActionJump} if !all(m -> m.rescale_rates_on_update == sr_val, majv) error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") end - maj = setup_majump_to_merge( - majv[1].scaled_rates, majv[1].reactant_stoch, - majv[1].net_stoch, majv[1].param_mapper, sr_val - ) + maj = setup_majump_to_merge(majv[1].scaled_rates, majv[1].reactant_stoch, + majv[1].net_stoch, majv[1].param_mapper, sr_val) for i in 2:length(majv) massaction_jump_combine(maj, majv[i]) end - return JumpSet(vjs, cjs, rj, maj) + JumpSet(vjs, cjs, rj, maj) end @inline get_num_majumps(jset::JumpSet) = get_num_majumps(jset.massaction_jump) @inline num_majumps(jset::JumpSet) = get_num_majumps(jset) @inline function num_crjs(jset::JumpSet) - return (jset.constant_jumps !== nothing) ? length(jset.constant_jumps) : 0 + (jset.constant_jumps !== nothing) ? length(jset.constant_jumps) : 0 end @inline function num_vrjs(jset::JumpSet) - return (jset.variable_jumps !== nothing) ? length(jset.variable_jumps) : 0 + (jset.variable_jumps !== nothing) ? length(jset.variable_jumps) : 0 end @inline function num_bndvrjs(jset::JumpSet) - return (jset.variable_jumps !== nothing) ? count(isbounded, jset.variable_jumps) : 0 + (jset.variable_jumps !== nothing) ? count(isbounded, jset.variable_jumps) : 0 end @inline function num_continvrjs(jset::JumpSet) - return (jset.variable_jumps !== nothing) ? count(!isbounded, jset.variable_jumps) : 0 + (jset.variable_jumps !== nothing) ? count(!isbounded, jset.variable_jumps) : 0 end num_jumps(jset::JumpSet) = num_majumps(jset) + num_crjs(jset) + num_vrjs(jset) @@ -674,24 +628,22 @@ num_cdiscretejumps(jset::JumpSet) = num_majumps(jset) + num_crjs(jset) @inline split_jumps(vj, cj, rj, maj) = vj, cj, rj, maj @inline function split_jumps(vj, cj, rj, maj, v::VariableRateJump, args...) - return split_jumps((vj..., v), cj, rj, maj, args...) + split_jumps((vj..., v), cj, rj, maj, args...) end @inline function split_jumps(vj, cj, rj, maj, c::ConstantRateJump, args...) - return split_jumps(vj, (cj..., c), rj, maj, args...) + split_jumps(vj, (cj..., c), rj, maj, args...) end @inline function split_jumps(vj, cj, rj, maj, c::RegularJump, args...) - return split_jumps(vj, cj, regular_jump_combine(rj, c), maj, args...) + split_jumps(vj, cj, regular_jump_combine(rj, c), maj, args...) end @inline function split_jumps(vj, cj, rj, maj, c::MassActionJump, args...) - return split_jumps(vj, cj, rj, massaction_jump_combine(maj, c), args...) + split_jumps(vj, cj, rj, massaction_jump_combine(maj, c), args...) end @inline function split_jumps(vj, cj, rj, maj, j::JumpSet, args...) - return split_jumps( - (vj..., j.variable_jumps...), + split_jumps((vj..., j.variable_jumps...), (cj..., j.constant_jumps...), regular_jump_combine(rj, j.regular_jump), - massaction_jump_combine(maj, j.massaction_jump), args... - ) + massaction_jump_combine(maj, j.massaction_jump), args...) end regular_jump_combine(rj1::RegularJump, rj2::Nothing) = rj1 @@ -703,65 +655,43 @@ end # functionality to merge two mass action jumps together function check_majump_type(maj::MassActionJump{S, T, U, V}) where {S <: Number, T, U, V} - return setup_majump_to_merge( - maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, - maj.param_mapper, maj.rescale_rates_on_update - ) + setup_majump_to_merge(maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, + maj.param_mapper, maj.rescale_rates_on_update) end function check_majump_type(maj::MassActionJump{Nothing, T, U, V}) where {T, U, V} - return setup_majump_to_merge( - maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, - maj.param_mapper, maj.rescale_rates_on_update - ) + setup_majump_to_merge(maj.scaled_rates, maj.reactant_stoch, maj.net_stoch, + maj.param_mapper, maj.rescale_rates_on_update) end # if given containers of rates and stoichiometry directly create a jump -function setup_majump_to_merge( - sr::T, rs::AbstractVector{S}, ns::AbstractVector{U}, pmapper, - rescale_rates_on_update::Bool - ) where { - T <: AbstractVector, S <: AbstractArray, - U <: AbstractArray, - } - return MassActionJump(sr, rs, ns, pmapper; scale_rates = false, rescale_rates_on_update) +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) end # if just given the data for one jump (and not in a container) wrap in a vector -function setup_majump_to_merge( - sr::S, rs::T, ns::U, pmapper, - rescale_rates_on_update::Bool - ) where { - S <: Number, T <: AbstractArray, - U <: AbstractArray, - } - return MassActionJump( - [sr], [rs], [ns], +function setup_majump_to_merge(sr::S, rs::T, ns::U, pmapper, + rescale_rates_on_update::Bool) where {S <: Number, T <: AbstractArray, + U <: AbstractArray} + MassActionJump([sr], [rs], [ns], (pmapper === nothing) ? pmapper : to_collection(pmapper); - scale_rates = false, rescale_rates_on_update - ) + scale_rates = false, rescale_rates_on_update) end # if no rate field setup yet -function setup_majump_to_merge( - ::Nothing, rs::T, ns::U, pmapper, - rescale_rates_on_update::Bool - ) where {T <: AbstractArray, U <: AbstractArray} - return MassActionJump( - nothing, [rs], [ns], +function setup_majump_to_merge(::Nothing, rs::T, ns::U, pmapper, + rescale_rates_on_update::Bool) where {T <: AbstractArray, U <: AbstractArray} + MassActionJump(nothing, [rs], [ns], (pmapper === nothing) ? pmapper : to_collection(pmapper); - scale_rates = false, rescale_rates_on_update - ) + scale_rates = false, rescale_rates_on_update) end # when given a collection of reactions to add to maj -function majump_merge!( - maj::MassActionJump{U, <:AbstractVector{V}, <:AbstractVector{W}, X}, +function majump_merge!(maj::MassActionJump{U, <:AbstractVector{V}, <:AbstractVector{W}, X}, sr::U, rs::AbstractVector{V}, ns::AbstractVector{W}, - param_mapper - ) where { - U <: Union{AbstractVector, Nothing}, - V <: AbstractVector, W <: AbstractVector, X, - } + param_mapper) where {U <: Union{AbstractVector, Nothing}, + V <: AbstractVector, W <: AbstractVector, X} (U <: AbstractVector) && append!(maj.scaled_rates, sr) append!(maj.reactant_stoch, rs) append!(maj.net_stoch, ns) @@ -771,20 +701,16 @@ function majump_merge!( else merge!(maj.param_mapper, param_mapper) end - return maj + maj end # when given a single jump's worth of data to add to maj -function majump_merge!( - maj::MassActionJump{U, V, W, X}, sr::T, rs::S1, ns::S2, - param_mapper - ) where { - T <: Union{Number, Nothing}, +function majump_merge!(maj::MassActionJump{U, V, W, X}, sr::T, rs::S1, ns::S2, + param_mapper) where {T <: Union{Number, Nothing}, S1 <: AbstractArray, S2 <: AbstractArray, U <: Union{AbstractVector{T}, Nothing}, V <: AbstractVector{S1}, - W <: AbstractVector{S2}, X, - } + W <: AbstractVector{S2}, X} (T <: Number) && push!(maj.scaled_rates, sr) push!(maj.reactant_stoch, rs) push!(maj.net_stoch, ns) @@ -795,34 +721,26 @@ function majump_merge!( merge!(maj.param_mapper, param_mapper) end - return maj + maj end # when maj only stores a single jump's worth of data (and not in a collection) # create a new jump with the merged data stored in vectors -function majump_merge!( - maj::MassActionJump{T, S, U, V}, sr::T, rs::S, ns::U, - param_mapper::V - ) where { - T <: Union{Number, Nothing}, +function majump_merge!(maj::MassActionJump{T, S, U, V}, sr::T, rs::S, ns::U, + param_mapper::V) where {T <: Union{Number, Nothing}, S <: AbstractArray{<:Pair}, - U <: AbstractArray{<:Pair}, V, - } + U <: AbstractArray{<:Pair}, V} rates = (T <: Nothing) ? nothing : [maj.scaled_rates, sr] if maj.param_mapper === nothing (param_mapper === nothing) || error("Error, trying to merge a MassActionJump with a parameter mapping to one without a parameter mapping.") - return MassActionJump( - rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], + return MassActionJump(rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], param_mapper; scale_rates = false, - rescale_rates_on_update = maj.rescale_rates_on_update - ) + rescale_rates_on_update = maj.rescale_rates_on_update) else - return MassActionJump( - rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], + return MassActionJump(rates, [maj.reactant_stoch, rs], [maj.net_stoch, ns], merge(maj.param_mapper, param_mapper); scale_rates = false, - rescale_rates_on_update = maj.rescale_rates_on_update - ) + rescale_rates_on_update = maj.rescale_rates_on_update) end end @@ -832,10 +750,8 @@ massaction_jump_combine(maj1::Nothing, maj2::Nothing) = maj1 function massaction_jump_combine(maj1::MassActionJump, maj2::MassActionJump) (maj1.rescale_rates_on_update == maj2.rescale_rates_on_update) || error("Cannot merge MassActionJumps with different rescale_rates_on_update settings.") - return majump_merge!( - maj1, maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, - maj2.param_mapper - ) + majump_merge!(maj1, maj2.scaled_rates, maj2.reactant_stoch, maj2.net_stoch, + maj2.param_mapper) end ##### helper methods for unpacking rates and affects! from constant jumps ##### @@ -848,14 +764,12 @@ function get_jump_info_tuples(jumps) affects! = () end - return rates, affects! + rates, affects! end function get_jump_info_fwrappers(u, p, t, jumps) - RateWrapper = FunctionWrappers.FunctionWrapper{ - typeof(t), - Tuple{typeof(u), typeof(p), typeof(t)}, - } + RateWrapper = FunctionWrappers.FunctionWrapper{typeof(t), + Tuple{typeof(u), typeof(p), typeof(t)}} if (jumps !== nothing) && !isempty(jumps) rates = [RateWrapper(c.rate) for c in jumps] @@ -865,5 +779,5 @@ function get_jump_info_fwrappers(u, p, t, jumps) affects! = Any[] end - return rates, affects! + rates, affects! end diff --git a/src/simple_regular_solve.jl b/src/simple_regular_solve.jl index 4e08378c2..db852eb67 100644 --- a/src/simple_regular_solve.jl +++ b/src/simple_regular_solve.jl @@ -71,7 +71,7 @@ sol = solve(jprob, SimpleExplicitTauLeaping()) ``` """ struct SimpleExplicitTauLeaping{T <: AbstractFloat} <: SciMLBase.AbstractDEAlgorithm - epsilon::T + epsilon::T # Error control parameter end SimpleExplicitTauLeaping(; epsilon = 0.05) = SimpleExplicitTauLeaping(epsilon) @@ -82,7 +82,7 @@ function validate_pure_leaping_inputs(jump_prob::JumpProblem, alg) JumpProblem, i.e. call JumpProblem(::DiscreteProblem, PureLeaping(),...). \ Passing $(jump_prob.aggregator) is deprecated and will be removed in the next breaking release." end - return isempty(jump_prob.jump_callback.continuous_callbacks) && + isempty(jump_prob.jump_callback.continuous_callbacks) && isempty(jump_prob.jump_callback.discrete_callbacks) && isempty(jump_prob.constant_jumps) && isempty(jump_prob.variable_jumps) && @@ -96,7 +96,7 @@ function validate_pure_leaping_inputs(jump_prob::JumpProblem, alg::SimpleExplici JumpProblem, i.e. call JumpProblem(::DiscreteProblem, PureLeaping(),...). \ Passing $(jump_prob.aggregator) is deprecated and will be removed in the next breaking release." end - return isempty(jump_prob.jump_callback.continuous_callbacks) && + isempty(jump_prob.jump_callback.continuous_callbacks) && isempty(jump_prob.jump_callback.discrete_callbacks) && isempty(jump_prob.constant_jumps) && isempty(jump_prob.variable_jumps) && @@ -122,7 +122,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 @@ -139,11 +139,9 @@ function _process_saveat(saveat, tspan, save_start, save_end) return saveat_vec, _save_start, _save_end end -function DiffEqBase.solve( - jump_prob::JumpProblem, alg::SimpleTauLeaping; +function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleTauLeaping; seed = nothing, dt = error("dt is required for SimpleTauLeaping."), - saveat = nothing, save_start = nothing, save_end = nothing - ) + 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.") @@ -214,11 +212,9 @@ function DiffEqBase.solve( push!(tsave, tspan[2]) end - return sol = SciMLBase.build_solution( - prob, alg, tsave, usave, + sol = SciMLBase.build_solution(prob, alg, tsave, usave, calculate_error = false, - interp = SciMLBase.ConstantInterpolation(tsave, usave) - ) + interp = SciMLBase.ConstantInterpolation(tsave, usave)) end # Compute the highest order of reaction (HOR) for each reaction j, as per Cao et al. (2006), Section IV. @@ -229,8 +225,7 @@ function compute_hor(reactant_stoch, numjumps) hor = zeros(stoch_type, numjumps) for j in 1:numjumps order = sum( - stoch for (spec_idx, stoch) in reactant_stoch[j]; init = zero(stoch_type) - ) + stoch for (spec_idx, stoch) in reactant_stoch[j]; init = zero(stoch_type)) if order > 3 error("Reaction $j has order $order, which is not supported (maximum order is 3).") end @@ -286,19 +281,19 @@ function compute_gi(u, max_hor, max_stoich, i, t) return 2 * one_max_hor else # if max_stoich[i] == 2 return u[i] > one_max_hor ? - 2 * one_max_hor + one_max_hor / (u[i] - one_max_hor) : 2 * one_max_hor # Fallback to 2 if x_i <= 1 + 2 * one_max_hor + one_max_hor / (u[i] - one_max_hor) : 2 * one_max_hor # Fallback to 2 if x_i <= 1 end elseif max_hor[i] == 3 if max_stoich[i] == 1 return 3 * one_max_hor elseif max_stoich[i] == 2 return u[i] > one_max_hor ? - (3 * one_max_hor / 2) * - (2 * one_max_hor + one_max_hor / (u[i] - one_max_hor)) : 3 * one_max_hor # Fallback to 3 if x_i <= 1 + (3 * one_max_hor / 2) * + (2 * one_max_hor + one_max_hor / (u[i] - one_max_hor)) : 3 * one_max_hor # Fallback to 3 if x_i <= 1 else # if max_stoich[i] == 3 return u[i] > 2 * one_max_hor ? - 3 * one_max_hor + one_max_hor / (u[i] - one_max_hor) + - 2 * one_max_hor / (u[i] - 2 * one_max_hor) : 3 * one_max_hor # Fallback to 3 if x_i <= 2 + 3 * one_max_hor + one_max_hor / (u[i] - one_max_hor) + + 2 * one_max_hor / (u[i] - 2 * one_max_hor) : 3 * one_max_hor # Fallback to 3 if x_i <= 2 end end return one_max_hor # Default case @@ -310,8 +305,7 @@ end # mu_i(x) = sum_j nu_ij * a_j(x), sigma_i^2(x) = sum_j nu_ij^2 * a_j(x) # I_rs is the set of reactant species (assumed to be all species here, as critical reactions are not specified). function compute_tau( - u, rate_cache, nu, hor, p, t, epsilon, rate, dtmin, max_hor, max_stoich, numjumps - ) + u, rate_cache, nu, hor, p, t, epsilon, rate, dtmin, max_hor, max_stoich, numjumps) rate(rate_cache, u, p, t) if all(<=(0), rate_cache) # Handle case where all rates are zero or negative return dtmin @@ -346,8 +340,7 @@ function simple_explicit_tau_leaping_loop!( 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 - ) + save_end) save_idx = 1 while t_current < t_end @@ -356,13 +349,11 @@ function simple_explicit_tau_leaping_loop!( t_current = t_end break end - tau = compute_tau( - u_current, rate_cache, nu, hor, p, t_current, - epsilon, rate, dtmin, max_hor, max_stoich, numjumps - ) + tau = compute_tau(u_current, rate_cache, nu, hor, p, t_current, + epsilon, rate, dtmin, max_hor, max_stoich, numjumps) tau = min(tau, 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 # Calculate Poisson random numbers only for positive rates @@ -394,7 +385,7 @@ function simple_explicit_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] @@ -407,18 +398,16 @@ function simple_explicit_tau_leaping_loop!( end # Save endpoint if requested and not already saved - return if save_end && (isempty(tsave) || tsave[end] != t_end) + if save_end && (isempty(tsave) || tsave[end] != t_end) push!(usave, copy(u_current)) push!(tsave, t_end) end end -function DiffEqBase.solve( - jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; +function DiffEqBase.solve(jump_prob::JumpProblem, alg::SimpleExplicitTauLeaping; seed = nothing, dtmin = nothing, - saveat = nothing, save_start = nothing, save_end = 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.") @@ -427,7 +416,7 @@ function DiffEqBase.solve( tspan = prob.tspan if dtmin === nothing - dtmin = 1.0e-10 * one(typeof(tspan[2])) + dtmin = 1e-10 * one(typeof(tspan[2])) end (seed !== nothing) && seed!(rng, seed) @@ -472,21 +461,17 @@ function DiffEqBase.solve( reactant_stoch = maj.reactant_stoch hor = compute_hor(reactant_stoch, numjumps) max_hor, max_stoich = precompute_reaction_conditions( - reactant_stoch, hor, length(u0), numjumps - ) + reactant_stoch, hor, length(u0), numjumps) simple_explicit_tau_leaping_loop!( 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 - ) + save_end) - sol = SciMLBase.build_solution( - prob, alg, tsave, usave, + sol = SciMLBase.build_solution(prob, alg, tsave, usave, calculate_error = false, - interp = SciMLBase.ConstantInterpolation(tsave, usave) - ) + interp = SciMLBase.ConstantInterpolation(tsave, usave)) return sol end @@ -525,9 +510,9 @@ struct EnsembleGPUKernel{Backend} <: SciMLBase.EnsembleAlgorithm end function EnsembleGPUKernel(backend) - return EnsembleGPUKernel(backend, 0.0) + EnsembleGPUKernel(backend, 0.0) end function EnsembleGPUKernel() - return EnsembleGPUKernel(nothing, 0.0) + EnsembleGPUKernel(nothing, 0.0) end diff --git a/src/spatial/spatial_massaction_jump.jl b/src/spatial/spatial_massaction_jump.jl index cb1be14de..f75e21260 100644 --- a/src/spatial/spatial_massaction_jump.jl +++ b/src/spatial/spatial_massaction_jump.jl @@ -56,32 +56,31 @@ get_num_majumps(smaj) == 1 ``` """ struct SpatialMassActionJump{A <: AVecOrNothing, B <: AMatOrNothing, S, U, V} <: - AbstractMassActionJump - uniform_rates::A - spatial_rates::B + AbstractMassActionJump + uniform_rates::A # reactions that are uniform in space + spatial_rates::B # reactions whose rate depends on the site reactant_stoch::S net_stoch::U param_mapper::V - function SpatialMassActionJump{A, B, S, U, V}( - uniform_rates::A, spatial_rates::B, + """ + uniform rates go first in ordering + """ + function SpatialMassActionJump{A, B, S, U, V}(uniform_rates::A, spatial_rates::B, reactant_stoch::S, net_stoch::U, param_mapper::V, scale_rates::Bool, useiszero::Bool, - nocopy::Bool - ) where { - A <: AVecOrNothing, + nocopy::Bool) where {A <: AVecOrNothing, B <: AMatOrNothing, - S, U, V, - } + S, U, V} uniform_rates = (nocopy || isnothing(uniform_rates)) ? uniform_rates : - copy(uniform_rates) + copy(uniform_rates) spatial_rates = (nocopy || isnothing(spatial_rates)) ? spatial_rates : - copy(spatial_rates) + copy(spatial_rates) reactant_stoch = nocopy ? reactant_stoch : copy(reactant_stoch) for i in eachindex(reactant_stoch) if useiszero && (length(reactant_stoch[i]) == 1) && - iszero(reactant_stoch[i][1][1]) + iszero(reactant_stoch[i][1][1]) reactant_stoch[i] = typeof(reactant_stoch[i])() end end @@ -92,154 +91,102 @@ struct SpatialMassActionJump{A <: AVecOrNothing, B <: AMatOrNothing, S, U, V} <: if scale_rates && !isnothing(spatial_rates) && !isempty(spatial_rates) scalerates!(spatial_rates, reactant_stoch[(num_unif_rates + 1):end]) end - return new(uniform_rates, spatial_rates, reactant_stoch, net_stoch, param_mapper) + new(uniform_rates, spatial_rates, reactant_stoch, net_stoch, param_mapper) end end ################ Constructors ################## -function SpatialMassActionJump( - urates::A, srates::B, rs::S, ns::U, pmapper::V; +function SpatialMassActionJump(urates::A, srates::B, rs::S, ns::U, pmapper::V; scale_rates = true, useiszero = true, - nocopy = false - ) where { - A <: AVecOrNothing, - B <: AMatOrNothing, S, U, V, - } - return SpatialMassActionJump{A, B, S, U, V}( - urates, srates, rs, ns, pmapper, scale_rates, - useiszero, nocopy - ) -end -function SpatialMassActionJump( - urates::A, srates::B, rs, ns; scale_rates = true, + nocopy = false) where {A <: AVecOrNothing, + B <: AMatOrNothing, S, U, V} + SpatialMassActionJump{A, B, S, U, V}(urates, srates, rs, ns, pmapper, scale_rates, + useiszero, nocopy) +end +function SpatialMassActionJump(urates::A, srates::B, rs, ns; scale_rates = true, useiszero = true, - nocopy = false - ) where { - A <: AVecOrNothing, - B <: AMatOrNothing, - } - return SpatialMassActionJump( - urates, srates, rs, ns, nothing; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy - ) -end - -function SpatialMassActionJump( - srates::B, rs, ns, pmapper; scale_rates = true, + nocopy = false) where {A <: AVecOrNothing, + B <: AMatOrNothing} + SpatialMassActionJump(urates, srates, rs, ns, nothing; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy) +end + +function SpatialMassActionJump(srates::B, rs, ns, pmapper; scale_rates = true, useiszero = true, - nocopy = false - ) where {B <: AMatOrNothing} - return SpatialMassActionJump( - nothing, srates, rs, ns, pmapper; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy - ) -end -function SpatialMassActionJump( - srates::B, rs, ns; scale_rates = true, useiszero = true, - nocopy = false - ) where {B <: AMatOrNothing} - return SpatialMassActionJump( - nothing, srates, rs, ns, nothing; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy - ) -end - -function SpatialMassActionJump( - urates::A, rs, ns, pmapper; scale_rates = true, + nocopy = false) where {B <: AMatOrNothing} + SpatialMassActionJump(nothing, srates, rs, ns, pmapper; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy) +end +function SpatialMassActionJump(srates::B, rs, ns; scale_rates = true, useiszero = true, + nocopy = false) where {B <: AMatOrNothing} + SpatialMassActionJump(nothing, srates, rs, ns, nothing; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy) +end + +function SpatialMassActionJump(urates::A, rs, ns, pmapper; scale_rates = true, useiszero = true, - nocopy = false - ) where {A <: AVecOrNothing} - return SpatialMassActionJump( - urates, nothing, rs, ns, pmapper; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy - ) -end -function SpatialMassActionJump( - urates::A, rs, ns; scale_rates = true, useiszero = true, - nocopy = false - ) where {A <: AVecOrNothing} - return SpatialMassActionJump( - urates, nothing, rs, ns, nothing; scale_rates = scale_rates, - useiszero = useiszero, nocopy = nocopy - ) + nocopy = false) where {A <: AVecOrNothing} + SpatialMassActionJump(urates, nothing, rs, ns, pmapper; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy) +end +function SpatialMassActionJump(urates::A, rs, ns; scale_rates = true, useiszero = true, + nocopy = false) where {A <: AVecOrNothing} + SpatialMassActionJump(urates, nothing, rs, ns, nothing; scale_rates = scale_rates, + useiszero = useiszero, nocopy = nocopy) 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, - useiszero = true, nocopy = false - ) where {T, S, U, V} - return SpatialMassActionJump( - ma_jumps.scaled_rates, ma_jumps.reactant_stoch, +function SpatialMassActionJump(ma_jumps::MassActionJump{T, S, U, V}; scale_rates = false, + useiszero = true, nocopy = false) where {T, S, U, V} + SpatialMassActionJump(ma_jumps.scaled_rates, ma_jumps.reactant_stoch, ma_jumps.net_stoch, ma_jumps.param_mapper; - scale_rates = scale_rates, useiszero = useiszero, nocopy = nocopy - ) + scale_rates = scale_rates, useiszero = useiszero, nocopy = nocopy) end ############################################## -function get_num_majumps( - smaj::SpatialMassActionJump{ - Nothing, Nothing, S, U, V, - } - ) where - {S, U, V} - return 0 -end -function get_num_majumps( - smaj::SpatialMassActionJump{ - Nothing, B, S, U, V, - } - ) where - {B, S, U, V} - return size(smaj.spatial_rates, 1) -end -function get_num_majumps( - smaj::SpatialMassActionJump{ - A, Nothing, S, U, V, - } - ) where - {A, S, U, V} - return length(smaj.uniform_rates) -end -function get_num_majumps( - smaj::SpatialMassActionJump{ - A, B, S, U, V, - } - ) where - {A <: AbstractVector, B <: AbstractMatrix, S, U, V} - return length(smaj.uniform_rates) + size(smaj.spatial_rates, 1) +function get_num_majumps(smaj::SpatialMassActionJump{ + Nothing, Nothing, S, U, V}) where + {S, U, V} + 0 +end +function get_num_majumps(smaj::SpatialMassActionJump{ + Nothing, B, S, U, V}) where + {B, S, U, V} + size(smaj.spatial_rates, 1) +end +function get_num_majumps(smaj::SpatialMassActionJump{ + A, Nothing, S, U, V}) where + {A, S, U, V} + length(smaj.uniform_rates) +end +function get_num_majumps(smaj::SpatialMassActionJump{ + A, B, S, U, V}) where + {A <: AbstractVector, B <: AbstractMatrix, S, U, V} + length(smaj.uniform_rates) + size(smaj.spatial_rates, 1) end using_params(smaj::SpatialMassActionJump) = false -function rate_at_site( - rx, site, - smaj::SpatialMassActionJump{Nothing, B, S, U, V} - ) where {B, S, U, V} - return smaj.spatial_rates[rx, site] -end -function rate_at_site( - rx, site, - smaj::SpatialMassActionJump{A, Nothing, S, U, V} - ) where {A, S, U, V} - return smaj.uniform_rates[rx] -end -function rate_at_site( - rx, site, - smaj::SpatialMassActionJump{A, B, S, U, V} - ) where - {A <: AbstractVector, B <: AbstractMatrix, S, U, V} +function rate_at_site(rx, site, + smaj::SpatialMassActionJump{Nothing, B, S, U, V}) where {B, S, U, V} + smaj.spatial_rates[rx, site] +end +function rate_at_site(rx, site, + smaj::SpatialMassActionJump{A, Nothing, S, U, V}) where {A, S, U, V} + smaj.uniform_rates[rx] +end +function rate_at_site(rx, site, + smaj::SpatialMassActionJump{A, B, S, U, V}) where + {A <: AbstractVector, B <: AbstractMatrix, S, U, V} num_unif_rxs = length(smaj.uniform_rates) - return rx <= num_unif_rxs ? smaj.uniform_rates[rx] : - smaj.spatial_rates[rx - num_unif_rxs, site] + rx <= num_unif_rxs ? smaj.uniform_rates[rx] : + smaj.spatial_rates[rx - num_unif_rxs, site] end -function evalrxrate( - speciesmat::AbstractMatrix{T}, rxidx::S, majump::SpatialMassActionJump, - site::Int - ) where {T, S} +function evalrxrate(speciesmat::AbstractMatrix{T}, rxidx::S, majump::SpatialMassActionJump, + site::Int) where {T, S} val = one(T) @inbounds for specstoch in majump.reactant_stoch[rxidx] specpop = speciesmat[specstoch[1], site] diff --git a/src/spatial/topology.jl b/src/spatial/topology.jl index 3c1d2218d..7641ef54c 100644 --- a/src/spatial/topology.jl +++ b/src/spatial/topology.jl @@ -38,7 +38,7 @@ const offsets_2D = [ CartesianIndex(0, -1), CartesianIndex(-1, 0), CartesianIndex(1, 0), - CartesianIndex(0, 1), + CartesianIndex(0, 1) ] const offsets_3D = [ CartesianIndex(0, 0, -1), @@ -46,7 +46,7 @@ const offsets_3D = [ CartesianIndex(-1, 0, 0), CartesianIndex(1, 0, 0), CartesianIndex(0, 1, 0), - CartesianIndex(0, 0, 1), + CartesianIndex(0, 0, 1) ] """ @@ -107,7 +107,7 @@ function nth_nbr(grid, site, n) CI = grid.CI offsets = grid.offsets @inbounds I = CI[site] - return @inbounds for off in offsets + @inbounds for off in offsets nb = I + off if nb in CI n -= 1 @@ -125,7 +125,7 @@ function neighbors(grid, site) CI = grid.CI LI = grid.LI I = CI[site] - return Iterators.map(off -> LI[off + I], Iterators.filter(off -> off + I in CI, grid.offsets)) + Iterators.map(off -> LI[off + I], Iterators.filter(off -> off + I in CI, grid.offsets)) end """ @@ -143,7 +143,7 @@ function pad_hop_vec!(to_pad::AbstractVector, grid, site, hop_vec::AbstractVecto to_pad[i] = zero(eltype(to_pad)) end end - return to_pad + to_pad end """ @@ -173,8 +173,9 @@ num_sites(grid) == 12 collect(neighbors(grid, 1)) == [2, 4] ``` """ -CartesianGrid(dims) = CartesianGridRej(dims) +CartesianGrid(dims) = CartesianGridRej(dims) # use CartesianGridRej by default +# neighbor sampling is rejection-based """ CartesianGridRej(dims) CartesianGridRej(dimension, linear_size::Int) @@ -234,11 +235,11 @@ function CartesianGridRej(dims::Tuple) LI = LinearIndices(dims) offsets = potential_offsets(dim) nums_neighbors = Int8[count(x -> x + CI[site] in CI, offsets) for site in 1:prod(dims)] - return CartesianGridRej(dims, nums_neighbors, CI, LI, offsets) + CartesianGridRej(dims, nums_neighbors, CI, LI, offsets) end CartesianGridRej(dims) = CartesianGridRej(Tuple(dims)) function CartesianGridRej(dimension, linear_size::Int) - return CartesianGridRej([linear_size for i in 1:dimension]) + CartesianGridRej([linear_size for i in 1:dimension]) end function rand_nbr(rng, grid::CartesianGridRej, site::Int) CI = grid.CI @@ -248,12 +249,9 @@ function rand_nbr(rng, grid::CartesianGridRej, site::Int) @inbounds nb = rand(rng, offsets) + I @inbounds nb in CI && return grid.LI[nb] end - return end -function Base.show( - io::IO, ::MIME"text/plain", - grid::CartesianGridRej - ) - return println(io, "A Cartesian grid with dimensions $(grid.dims)") +function Base.show(io::IO, ::MIME"text/plain", + grid::CartesianGridRej) + println(io, "A Cartesian grid with dimensions $(grid.dims)") end diff --git a/src/variable_rate.jl b/src/variable_rate.jl index 8abf0140a..5f882d4f1 100644 --- a/src/variable_rate.jl +++ b/src/variable_rate.jl @@ -58,10 +58,8 @@ sol = solve(jprob, Tsit5()) """ struct VR_FRM <: VariableRateAggregator end -function configure_jump_problem( - prob, vr_aggregator::VR_FRM, jumps, cvrjs; - rng = DEFAULT_RNG - ) +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) return new_prob, variable_jump_callback @@ -86,7 +84,7 @@ function extend_problem(prob::SciMLBase.AbstractODEProblem, jumps; rng = DEFAULT jump_f = let _f = _f function (du::ExtendedJumpArray, u::ExtendedJumpArray, p, t) _f(du.u, u.u, p, t) - return update_jumps!(du, u, p, t, length(u.u), jumps...) + update_jumps!(du, u, p, t, length(u.u), jumps...) end end else @@ -100,11 +98,9 @@ function extend_problem(prob::SciMLBase.AbstractODEProblem, jumps; rng = DEFAULT end u0 = extend_u0(prob, length(jumps), rng) - f = ODEFunction{isinplace(prob)}( - jump_f; sys = prob.f.sys, - observed = prob.f.observed - ) - return remake(prob; f, u0) + 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) @@ -114,7 +110,7 @@ function extend_problem(prob::SciMLBase.AbstractSDEProblem, jumps; rng = DEFAULT jump_f = let _f = _f function (du::ExtendedJumpArray, u::ExtendedJumpArray, p, t) _f(du.u, u.u, p, t) - return update_jumps!(du, u, p, t, length(u.u), jumps...) + update_jumps!(du, u, p, t, length(u.u), jumps...) end end else @@ -129,20 +125,18 @@ function extend_problem(prob::SciMLBase.AbstractSDEProblem, jumps; rng = DEFAULT if prob.noise_rate_prototype === nothing jump_g = function (du, u, p, t) - return prob.g(du.u, u.u, p, t) + prob.g(du.u, u.u, p, t) end else jump_g = function (du, u, p, t) - return prob.g(du, u.u, p, t) + prob.g(du, u.u, p, t) end end u0 = extend_u0(prob, length(jumps), rng) - f = SDEFunction{isinplace(prob)}( - jump_f, jump_g; sys = prob.f.sys, - observed = prob.f.observed - ) - return remake(prob; f, g = jump_g, u0) + 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) @@ -152,7 +146,7 @@ function extend_problem(prob::SciMLBase.AbstractDDEProblem, jumps; rng = DEFAULT jump_f = let _f = _f function (du::ExtendedJumpArray, u::ExtendedJumpArray, h, p, t) _f(du.u, u.u, h, p, t) - return update_jumps!(du, u, p, t, length(u.u), jumps...) + update_jumps!(du, u, p, t, length(u.u), jumps...) end end else @@ -166,11 +160,9 @@ function extend_problem(prob::SciMLBase.AbstractDDEProblem, jumps; rng = DEFAULT end u0 = extend_u0(prob, length(jumps), rng) - f = DDEFunction{isinplace(prob)}( - jump_f; sys = prob.f.sys, - observed = prob.f.observed - ) - return remake(prob; f, u0) + 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 @@ -181,7 +173,7 @@ function extend_problem(prob::SciMLBase.AbstractDAEProblem, jumps; rng = DEFAULT jump_f = let _f = _f function (out, du::ExtendedJumpArray, u::ExtendedJumpArray, h, p, t) _f(out, du.u, u.u, h, p, t) - return update_jumps!(out, u, p, t, length(u.u), jumps...) + update_jumps!(out, u, p, t, length(u.u), jumps...) end end else @@ -195,11 +187,9 @@ function extend_problem(prob::SciMLBase.AbstractDAEProblem, jumps; rng = DEFAULT end u0 = extend_u0(prob, length(jumps), rng) - f = DAEFunction{isinplace(prob)}( - jump_f, sys = prob.f.sys, - observed = prob.f.observed - ) - return remake(prob; f, u0) + f = DAEFunction{isinplace(prob)}(jump_f, sys = prob.f.sys, + observed = prob.f.observed) + remake(prob; f, u0) end struct VR_FRMEventCallback{F, RNG} @@ -215,50 +205,48 @@ end function (c::VR_FRMEventCallback)(integrator) c.affect!(integrator) integrator.u.jump_u[c.idx] = -randexp(c.rng, typeof(integrator.t)) - return nothing + 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)) derivative_discontinuity!(integrator, true) - return nothing + nothing end function wrap_jump_in_callback(idx, jump; rng = DEFAULT_RNG) cb_functor = VR_FRMEventCallback(idx, jump.affect!, rng) - return ContinuousCallback( - cb_functor, cb_functor; + ContinuousCallback(cb_functor, cb_functor; initialize = cb_functor, idxs = jump.idxs, rootfind = jump.rootfind, interp_points = jump.interp_points, save_positions = jump.save_positions, abstol = jump.abstol, - reltol = jump.reltol - ) + reltol = jump.reltol) end function build_variable_callback(cb, idx, jump, jumps...; rng = DEFAULT_RNG) idx += 1 new_cb = wrap_jump_in_callback(idx, jump; rng) - return build_variable_callback(CallbackSet(cb, new_cb), idx, jumps...; rng = DEFAULT_RNG) + build_variable_callback(CallbackSet(cb, new_cb), idx, jumps...; rng = DEFAULT_RNG) end function build_variable_callback(cb, idx, jump; rng = DEFAULT_RNG) idx += 1 - return CallbackSet(cb, wrap_jump_in_callback(idx, jump; rng)) + CallbackSet(cb, wrap_jump_in_callback(idx, jump; rng)) end @inline function update_jumps!(du, u, p, t, idx, jump) idx += 1 - return du[idx] = jump.rate(u.u, p, t) + du[idx] = jump.rate(u.u, p, t) end @inline function update_jumps!(du, u, p, t, idx, jump, jumps...) idx += 1 du[idx] = jump.rate(u.u, p, t) - return update_jumps!(du, u, p, t, idx, jumps...) + update_jumps!(du, u, p, t, idx, jumps...) end ################################### VR_Direct and VR_DirectFW #################################### @@ -307,7 +295,6 @@ sol = solve(jprob, Tsit5()) - `VR_Direct` and `VR_DirectFW` are expected to generally be more performant than `VR_FRM`. """ struct VR_Direct <: VariableRateAggregator end - """ $(TYPEDEF) @@ -359,8 +346,7 @@ mutable struct VR_DirectEventCache{T, RNG, F1, F2} end function VR_DirectEventCache( - jumps::JumpSet, ::VR_Direct, prob, ::Type{T}; rng = DEFAULT_RNG - ) where {T} + jumps::JumpSet, ::VR_Direct, prob, ::Type{T}; rng = DEFAULT_RNG) where {T} initial_threshold = randexp(rng, T) vjumps = jumps.variable_jumps @@ -369,16 +355,13 @@ function VR_DirectEventCache( cum_rate_sum = Vector{T}(undef, length(vjumps)) - return VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), typeof(affect_funcs)}( - zero(T), + VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), typeof(affect_funcs)}(zero(T), initial_threshold, zero(T), initial_threshold, zero(T), rng, rate_funcs, - affect_funcs, cum_rate_sum - ) + affect_funcs, cum_rate_sum) end function VR_DirectEventCache( - jumps::JumpSet, ::VR_DirectFW, prob, ::Type{T}; rng = DEFAULT_RNG - ) where {T} + jumps::JumpSet, ::VR_DirectFW, prob, ::Type{T}; rng = DEFAULT_RNG) where {T} initial_threshold = randexp(rng, T) vjumps = jumps.variable_jumps @@ -389,11 +372,9 @@ function VR_DirectEventCache( cum_rate_sum = Vector{T}(undef, length(vjumps)) - return VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), Any}( - zero(T), + VR_DirectEventCache{T, typeof(rng), typeof(rate_funcs), Any}(zero(T), initial_threshold, zero(T), initial_threshold, zero(T), rng, rate_funcs, - affect_funcs, cum_rate_sum - ) + affect_funcs, cum_rate_sum) end # Initialization function for VR_DirectEventCache @@ -404,29 +385,23 @@ function initialize_vr_direct_cache!(cache::VR_DirectEventCache, u, t, integrato cache.current_threshold = cache.prev_threshold cache.total_rate = zero(integrator.t) cache.cum_rate_sum .= 0 - return nothing + nothing end -@inline function concretize_vr_direct_affects!( - cache::VR_DirectEventCache, - ::I - ) where {I <: SciMLBase.DEIntegrator} +@inline function concretize_vr_direct_affects!(cache::VR_DirectEventCache, + ::I) where {I <: SciMLBase.DEIntegrator} if (cache.affect_funcs isa Vector) && - !(cache.affect_funcs isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}}) + !(cache.affect_funcs isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}}) AffectWrapper = FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}} - cache.affect_funcs = AffectWrapper[ - makewrapper(AffectWrapper, aff) - for aff in cache.affect_funcs - ] + cache.affect_funcs = AffectWrapper[makewrapper(AffectWrapper, aff) + for aff in cache.affect_funcs] end - return nothing + 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} - return nothing +@inline function concretize_vr_direct_affects!(cache::VR_DirectEventCache{T, RNG, F1, F2}, + ::I) where {T, RNG, F1, F2 <: Tuple, I <: SciMLBase.DEIntegrator} + nothing end # Wrapper for initialize to match ContinuousCallback signature @@ -434,7 +409,7 @@ function initialize_vr_direct_wrapper(cb::ContinuousCallback, u, t, integrator) concretize_vr_direct_affects!(cb.condition, integrator) initialize_vr_direct_cache!(cb.condition, u, t, integrator) derivative_discontinuity!(integrator, false) - return nothing + nothing end # Merge callback parameters across all jumps for VR_Direct @@ -449,10 +424,8 @@ function build_variable_integcallback(cache::VR_DirectEventCache, jumps) reltol = min(reltol, jump.reltol) end - return ContinuousCallback( - cache, cache; initialize = initialize_vr_direct_wrapper, - save_positions, abstol, reltol - ) + return ContinuousCallback(cache, cache; initialize = initialize_vr_direct_wrapper, + save_positions, abstol, reltol) end function configure_jump_problem(prob, ::VR_Direct, jumps, cvrjs; rng = DEFAULT_RNG) @@ -472,7 +445,7 @@ end # recursively evaluate the cumulative sum of the rates for type stability @inline function cumsum_rates!(cum_rate_sum, u, p, t, rates::Tuple) cur_sum = zero(eltype(cum_rate_sum)) - return cumsum_rates!(cum_rate_sum, u, p, t, 1, cur_sum, rates...) + cumsum_rates!(cum_rate_sum, u, p, t, 1, cur_sum, rates...) end # loop-based version for Vector rate_funcs (avoids dynamic splatting) @@ -489,18 +462,16 @@ end new_sum = cur_sum + rate(u, p, t) @inbounds cum_rate_sum[idx] = new_sum idx += 1 - return cumsum_rates!(cum_rate_sum, u, p, t, idx, new_sum, rates...) + cumsum_rates!(cum_rate_sum, u, p, t, idx, new_sum, rates...) end @inline function cumsum_rates!(cum_rate_sum, u, p, t, idx, cur_sum, rate) - return @inbounds cum_rate_sum[idx] = cur_sum + rate(u, p, t) + @inbounds cum_rate_sum[idx] = cur_sum + rate(u, p, t) end function total_variable_rate( cache::VR_DirectEventCache{ - T, RNG, F1, F2, - }, u, p, t - ) where {T, RNG, F1, F2} + T, RNG, F1, F2}, u, p, t) where {T, RNG, F1, F2} (; cum_rate_sum, rate_funcs) = cache sum_rate = cumsum_rates!(cum_rate_sum, u, p, t, rate_funcs) return sum_rate @@ -545,22 +516,18 @@ 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} - return quote +@generated function execute_affect!(cache::VR_DirectEventCache{T, RNG, F1, F2}, + integrator::I, idx) where {T, RNG, 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)) end end -@inline function execute_affect!( - cache::VR_DirectEventCache, - integrator::I, idx - ) where {I <: SciMLBase.DEIntegrator} +@inline function execute_affect!(cache::VR_DirectEventCache, + integrator::I, idx) where {I <: SciMLBase.DEIntegrator} (; affect_funcs) = cache - return if affect_funcs isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}} + if affect_funcs isa Vector{FunctionWrappers.FunctionWrapper{Nothing, Tuple{I}}} @inbounds affect_funcs[idx](integrator) else error("Error, invalid affect_funcs type. Expected a vector of function wrappers and got $(typeof(affect_funcs))") From 2994d82857729da5463eeb8bb1be848fbe7748f7 Mon Sep 17 00:00:00 2001 From: ChrisRackauckas-Claude Date: Sun, 12 Jul 2026 02:59:48 -0400 Subject: [PATCH 5/5] Retrigger documentation CI Co-Authored-By: Chris Rackauckas