Revamp MassActionJump initialization

.github/workflows/Downgrade.yml

@@ -3,6 +3,7 @@ on:
   pull_request:
     branches:
       - master
+      - v10_working
     paths-ignore:
       - 'docs/**'
   push:

.github/workflows/Tests.yml

@@ -8,6 +8,7 @@ on:
   pull_request:
     branches:
       - master
+      - v10_working
     paths-ignore:
       - 'docs/**'
 

.github/workflows/ThreadSafety.yml

@@ -9,6 +9,7 @@ on:
   pull_request:
     branches:
       - master
+      - v10_working
     paths-ignore:
       - 'docs/**'
 

HISTORY.md

@@ -21,6 +21,31 @@
     solver pathways (SSAStepper, ODE, SDE, tau-leaping).
   - `SSAIntegrator` now supports the `SciMLBase` RNG interface (`has_rng`,
     `get_rng`, `set_rng!`).
+  - **Breaking**: The `scale_rates` and `useiszero` keyword arguments have been
+    removed from `JumpProblem`. Set them on the `MassActionJump` directly:
+    ```julia
+    # Before (no longer works):
+    jprob = JumpProblem(dprob, Direct(), maj; scale_rates = false)
+
+    # After:
+    maj = MassActionJump(rates, reactant_stoch, net_stoch; scale_rates = false)
+    jprob = JumpProblem(dprob, Direct(), maj)
+    ```
+  - **Breaking**: Parameterized `MassActionJump`s (those constructed with
+    `param_idxs` or a custom `param_mapper`) are now immutable — rates are
+    computed from parameters at aggregator initialization rather than being
+    materialized into the jump at `JumpProblem` construction time. This means:
+      - `update_parameters!` has been removed. Mass action rates are now
+        automatically recomputed from the current parameter values whenever the
+        aggregator reinitializes. After modifying parameters (e.g. in a
+        callback), call `reset_aggregated_jumps!(integrator)` to trigger
+        reinitialization with the updated parameter values.
+      - Custom parameter mappers (e.g. ModelingToolkitBase's
+        `JumpSysMajParamMapper`) must implement the 3-arg callable API:
+        `(mapper)(dest::AbstractVector, maj::MassActionJump, params)`.
+        See [`MassActionJumpParamMapper`](@ref) for details.
+  - Scalar `param_idxs` (e.g. `param_idxs = 1`) is now internally converted to
+    a one-element vector. The scalar form continues to work as before.
 
 ## 9.14
 

Project.toml

@@ -1,7 +1,7 @@
 name = "JumpProcesses"
 uuid = "ccbc3e58-028d-4f4c-8cd5-9ae44345cda5"
 authors = ["Chris Rackauckas <accounts@chrisrackauckas.com>"]
-version = "9.23.1"
+version = "9.23.3"
 
 [deps]
 ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
@@ -24,9 +24,11 @@ SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
+OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
 
 [extensions]
 JumpProcessesKernelAbstractionsExt = ["Adapt", "KernelAbstractions"]
+JumpProcessesOrdinaryDiffEqCoreExt = "OrdinaryDiffEqCore"
 
 [compat]
 ADTypes = "1"
@@ -45,7 +47,7 @@ KernelAbstractions = "0.9"
 LinearAlgebra = "1"
 LinearSolve = "3"
 OrdinaryDiffEq = "6"
-OrdinaryDiffEqCore = "3.11"
+OrdinaryDiffEqCore = "3"
 Pkg = "1"
 PoissonRandom = "0.4"
 Random = "1"
@@ -72,7 +74,6 @@ FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
-OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"

docs/Project.toml

@@ -13,7 +13,7 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [compat]
-Catalyst = "14.0, 15"
+Catalyst = "16"
 DensityInterface = "0.4"
 DifferentialEquations = "7.11"
 Distributions = "0.25"

docs/src/assets/Project.toml

@@ -13,7 +13,7 @@ StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"
 
 [compat]
-Catalyst = "14.0, 15"
+Catalyst = "16"
 DensityInterface = "0.4"
 DifferentialEquations = "7.11"
 Distributions = "0.25"

docs/src/tutorials/discrete_stochastic_example.md

@@ -190,36 +190,25 @@ sir_model = @reaction_network begin
 end
 ```
 
-To build a pure jump process model of the reaction system, where the state is
-constant between jumps, we will use a
-[`DiscreteProblem`](https://docs.sciml.ai/DiffEqDocs/stable/types/discrete_types/).
-This encodes that the state only changes at the jump times. We do this by giving
-the constructor `u₀`, the initial condition, and `tspan`, the timespan. Here, we
-will start with ``990`` susceptible people, ``10`` infected person, and `0` recovered
-people, and solve the problem from `t=0.0` to `t=250.0`. We use the parameters
-`β = 0.1/1000` and `ν = 0.01`. Thus, we build the problem via:
+To build a pure jump process model of the reaction system we construct a
+[`JumpProblem`](@ref) directly from the Catalyst `ReactionSystem`. We specify
+`u₀`, the initial condition, `tspan`, the timespan, and `p`, the parameters.
+Here, we will start with ``990`` susceptible people, ``10`` infected person, and
+`0` recovered people, and solve the problem from `t=0.0` to `t=250.0`. We use
+the parameters `β = 0.1/1000` and `ν = 0.01`.
 
 ```@example tut2
 p = (:β => 0.1 / 1000, :ν => 0.01)
 u₀ = [:S => 990, :I => 10, :R => 0]
 tspan = (0.0, 250.0)
-prob = DiscreteProblem(sir_model, u₀, tspan, p)
+jump_prob = JumpProblem(sir_model, u₀, tspan, p)
 ```
 
 *Notice, the initial populations are integers, since we want the exact number of
 people in the different states.*
 
-The Catalyst reaction network can be converted into various
-DifferentialEquations.jl problem types, including `JumpProblem`s, `ODEProblem`s,
-or `SDEProblem`s. To turn it into a [`JumpProblem`](@ref) representing the SIR jump
-process model, we simply write
-
-```@example tut2
-jump_prob = JumpProblem(sir_model, prob, Direct())
-```
-
-Here `Direct()` indicates that we will determine the random times and types of
-reactions using [Gillespie's Direct stochastic simulation algorithm
+Here `Direct()` is the default aggregator, which determines the random times and
+types of reactions using [Gillespie's Direct stochastic simulation algorithm
 (SSA)](https://doi.org/10.1016/0021-9991(76)90041-3), also known as Doob's
 method or Kinetic Monte Carlo. See [Jump Aggregators for Exact Simulation](@ref) for
 other supported SSAs.

ext/JumpProcessesOrdinaryDiffEqCoreExt.jl

@@ -0,0 +1,26 @@
+module JumpProcessesOrdinaryDiffEqCoreExt
+
+using JumpProcesses
+import DiffEqBase
+import OrdinaryDiffEqCore: OrdinaryDiffEqAlgorithm, DAEAlgorithm,
+    StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm
+
+# Ambiguity fix: OrdinaryDiffEqCore defines
+#   __init(::Union{..., AbstractJumpProblem}, ::Union{OrdinaryDiffEqAlgorithm,
+#          DAEAlgorithm, StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm})
+# which is ambiguous with JumpProcesses'
+#   __init(::AbstractJumpProblem{P}, ::DEAlgorithm)
+#
+# This method resolves the ambiguity by being more specific in the problem type
+# (AbstractJumpProblem vs Union{..., AbstractJumpProblem, ...}) while matching
+# the exact algorithm union from OrdinaryDiffEqCore.
+function DiffEqBase.__init(
+        _jump_prob::DiffEqBase.AbstractJumpProblem{P},
+        alg::Union{OrdinaryDiffEqAlgorithm, DAEAlgorithm,
+            StochasticDiffEqAlgorithm, StochasticDiffEqRODEAlgorithm};
+        merge_callbacks = true, kwargs...) where {P}
+    kwargs = DiffEqBase.merge_problem_kwargs(_jump_prob; merge_callbacks, kwargs...)
+    JumpProcesses.__jump_init(_jump_prob, alg; kwargs...)
+end
+
+end

src/JumpProcesses.jl

@@ -5,7 +5,7 @@ using Reexport: Reexport, @reexport
 
 # Explicit imports from standard libraries
 using LinearAlgebra: LinearAlgebra, mul!
-using Random: Random, randexp, seed!
+using Random: Random, randexp
 
 # Explicit imports from external packages
 using DocStringExtensions: DocStringExtensions, FIELDS, TYPEDEF

src/aggregators/aggregated_api.jl

@@ -1,55 +1,41 @@
 """
-    reset_aggregated_jumps!(integrator, uprev = nothing; update_jump_params=true)
+    reset_aggregated_jumps!(integrator, uprev = nothing)
 
 Reset the state of jump processes and associated solvers following a change
-in parameters or such.
-
-Notes
-
-  - `update_jump_params=true` will recalculate the rates stored within any
-    MassActionJump that was built from the parameter vector. If the parameter
-    vector is unchanged, this can safely be set to false to improve performance.
+in parameters or such. Rate updates are handled automatically by `initialize!`
+via `fill_scaled_rates!`.
 """
-function reset_aggregated_jumps!(integrator, uprev = nothing; update_jump_params = true,
-        kwargs...)
-    reset_aggregated_jumps!(integrator, uprev, integrator.opts.callback,
-        update_jump_params = update_jump_params, kwargs...)
+function reset_aggregated_jumps!(integrator, uprev = nothing; kwargs...)
+    if haskey(kwargs, :update_jump_params)
+        throw(ArgumentError("`update_jump_params` keyword argument has been removed. " *
+                            "Rate updates are now handled automatically by `initialize!` " *
+                            "via `fill_scaled_rates!`."))
+    end
+    reset_aggregated_jumps!(integrator, uprev, integrator.opts.callback)
     nothing
 end
 
-function reset_aggregated_jumps!(integrator, uprev, callback::Nothing;
-        update_jump_params = true, kwargs...)
+function reset_aggregated_jumps!(integrator, uprev, callback::Nothing)
     nothing
 end
 
-function reset_aggregated_jumps!(integrator, uprev, callback::CallbackSet;
-        update_jump_params = true, kwargs...)
+function reset_aggregated_jumps!(integrator, uprev, callback::CallbackSet)
     if !isempty(callback.discrete_callbacks)
-        reset_aggregated_jumps!(integrator, uprev, callback.discrete_callbacks...,
-            update_jump_params = update_jump_params, kwargs...)
+        reset_aggregated_jumps!(integrator, uprev, callback.discrete_callbacks...)
     end
     nothing
 end
 
-function reset_aggregated_jumps!(integrator, uprev, cb::DiscreteCallback, cbs...;
-        update_jump_params = true, kwargs...)
+function reset_aggregated_jumps!(integrator, uprev, cb::DiscreteCallback, cbs...)
     if cb.condition isa AbstractSSAJumpAggregator
-        maj = cb.condition.ma_jumps
-        update_jump_params && using_params(maj) &&
-            update_parameters!(cb.condition.ma_jumps, integrator.p; kwargs...)
         cb.condition(cb, integrator.u, integrator.t, integrator)
     end
-    reset_aggregated_jumps!(integrator, uprev, cbs...;
-        update_jump_params = update_jump_params, kwargs...)
+    reset_aggregated_jumps!(integrator, uprev, cbs...)
     nothing
 end
 
-function reset_aggregated_jumps!(integrator, uprev, cb::DiscreteCallback;
-        update_jump_params = true, kwargs...)
+function reset_aggregated_jumps!(integrator, uprev, cb::DiscreteCallback)
     if cb.condition isa AbstractSSAJumpAggregator
-        maj = cb.condition.ma_jumps
-        update_jump_params && using_params(maj) &&
-            update_parameters!(cb.condition.ma_jumps, integrator.p; kwargs...)
         cb.condition(cb, integrator.u, integrator.t, integrator)
     end
     nothing

src/aggregators/bracketing.jl

@@ -47,8 +47,8 @@ end
 @inline get_spec_brackets(bd, i, u::AbstractVector) = get_spec_brackets(bd, i, u[i])
 
 # Get propensity brackets of massaction jump k.
-@inline function get_majump_brackets(ulow, uhigh, k, majumps)
-    evalrxrate(ulow, k, majumps), evalrxrate(uhigh, k, majumps)
+@inline function get_majump_brackets(ulow, uhigh, k, majumps, maj_rates)
+    evalrxrate(ulow, k, majumps, maj_rates), evalrxrate(uhigh, k, majumps, maj_rates)
 end
 
 # for constant rate jumps we must check the ordering of the bracket values
@@ -66,7 +66,7 @@ get brackets for the rate of reaction rx by first checking if the reaction is a
     ma_jumps = p.ma_jumps
     num_majumps = get_num_majumps(ma_jumps)
     if rx <= num_majumps
-        return get_majump_brackets(p.ulow, p.uhigh, rx, ma_jumps)
+        return get_majump_brackets(p.ulow, p.uhigh, rx, ma_jumps, p.maj_rates)
     else
         @inbounds return get_cjump_brackets(p.ulow, p.uhigh, p.rates[rx - num_majumps],
             params, t)
@@ -108,8 +108,9 @@ function set_bracketing!(p::AbstractSSAJumpAggregator, u, params, t)
     majumps = p.ma_jumps
     crlow = p.cur_rate_low
     crhigh = p.cur_rate_high
+    maj_rates = p.maj_rates
     @inbounds for k in 1:get_num_majumps(majumps)
-        crlow[k], crhigh[k] = get_majump_brackets(p.ulow, p.uhigh, k, majumps)
+        crlow[k], crhigh[k] = get_majump_brackets(p.ulow, p.uhigh, k, majumps, maj_rates)
         sum_rate += crhigh[k]
     end
 

src/aggregators/ccnrm.jl

@@ -12,6 +12,7 @@ mutable struct CCNRMJumpAggregation{T, S, F1, F2, DEPGR, PT} <:
     cur_rates::Vector{T}
     sum_rate::T
     ma_jumps::S
+    maj_rates::Vector{T}
     rates::F1
     affects!::F2
     save_positions::Tuple{Bool, Bool}
@@ -21,6 +22,7 @@ end
 
 function CCNRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T,
         maj::S, rs::F1, affs!::F2, sps::Tuple{Bool, Bool};
+        maj_rates = Vector{T}(undef, get_num_majumps(maj)),
         num_specs, dep_graph = nothing,
         kwargs...) where {T, S, F1, F2}
 
@@ -46,7 +48,7 @@ function CCNRMJumpAggregation(nj::Int, njt::T, et::T, crs::Vector{T}, sr::T,
     affecttype = F2 <: Tuple ? F2 : Any
     CCNRMJumpAggregation{T, S, F1, affecttype, typeof(dg), typeof(ptt)}(
         nj, nj, njt, et,
-        crs, sr, maj,
+        crs, sr, maj, maj_rates,
         rs, affs!, sps,
         dg, ptt)
 end
@@ -67,6 +69,7 @@ end
 # set up a new simulation and calculate the first jump / jump time
 function initialize!(p::CCNRMJumpAggregation, integrator, u, params, t)
     p.end_time = integrator.sol.prob.tspan[2]
+    fill_scaled_rates!(p.maj_rates, p.ma_jumps, params)
     rng = get_rng(integrator)
     initialize_rates_and_times!(p, u, params, t, rng)
     generate_jumps!(p, integrator, u, params, t)
@@ -115,7 +118,7 @@ function update_dependent_rates!(p::CCNRMJumpAggregation, u, params, t, rng)
 
         # update the jump rate
         @inbounds cur_rates[rx] = calculate_jump_rate(ma_jumps, num_majumps, rates, u,
-            params, t, rx)
+            params, t, rx, p.maj_rates)
 
         # Calculate new jump times for dependent jumps
         if rx != p.next_jump && oldrate > zero(oldrate)
@@ -141,8 +144,9 @@ function initialize_rates_and_times!(p::CCNRMJumpAggregation, u, params, t, rng)
     majumps = p.ma_jumps
     cur_rates = p.cur_rates
     pttdata = Vector{typeof(t)}(undef, length(cur_rates))
+    maj_rates = p.maj_rates
     @inbounds for i in 1:get_num_majumps(majumps)
-        cur_rates[i] = evalrxrate(u, i, majumps)
+        cur_rates[i] = evalrxrate(u, i, majumps, maj_rates)
         pttdata[i] = t + randexp(rng) / cur_rates[i]
     end