Stochasticad constant rate

Project.toml

@@ -8,6 +8,7 @@ ArrayInterface = "4fba245c-0d91-5ea0-9b3e-6abc04ee57a9"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 DiffEqBase = "2b5f629d-d688-5b77-993f-72d75c75574e"
 DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def"
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
 Graphs = "86223c79-3864-5bf0-83f7-82e725a168b6"
@@ -17,6 +18,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 RecursiveArrayTools = "731186ca-8d62-57ce-b412-fbd966d074cd"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5"
 
@@ -36,6 +38,7 @@ ArrayInterface = "7.25.0"
 DataStructures = "0.19.5, 0.19"
 DiffEqBase = "7.5.5, 7"
 DiffEqCallbacks = "4.17.0"
+Distributions = "0.25"
 DocStringExtensions = "0.9.5"
 FastBroadcast = "1.3.2, 1"
 FunctionWrappers = "1.1.3"
@@ -53,6 +56,7 @@ RecursiveArrayTools = "4.3.0, 4"
 Reexport = "1.2.2"
 SafeTestsets = "0.1"
 SciMLBase = "3.18.0, 3.1"
+SciMLStructures = "1.10.0, 1"
 SciMLTesting = "1.6"
 StableRNGs = "1"
 StaticArrays = "1.9.18"

docs/src/api.md

@@ -9,6 +9,7 @@ CurrentModule = JumpProcesses
 ```@docs
 JumpProblem
 SSAStepper
+BoundedSSA
 reset_aggregated_jumps!
 ```
 

docs/src/jump_solve.md

@@ -62,6 +62,17 @@ algorithms are optimized for pure jump problems.
     [`SSAStepper`](@ref) generated-solution uses piecewise constant
     interpolation, and can therefore exactly evaluate the sampled solution
     path at any time when only saving the post-jump state for each jump.
+  - `BoundedSSA`: a uniformization (thinning) SSA for jump-only `ConstantRateJump`
+    `DiscreteProblem`s, called as `solve(jprob, BoundedSSA(; rate_bound))` with a
+    constant bound on the total propensity. Unlike `SSAStepper` it is designed to
+    compose with automatic differentiation: when `StochasticAD.jl` is loaded and a
+    `StochasticTriple` parameter is used, `derivative_estimate` /
+    `stochastic_triple` give correct gradients of expectations over the process
+    (including state-dependent rates); it supports `saveat` for the full path. With
+    ordinary parameters it is just an (unbiased) SSA simulation. Valid whenever the
+    total propensity is bounded over the trajectory (e.g. population-bounded
+    systems); currently `ConstantRateJump`s with additive affects only (no
+    `MassActionJump`/`VariableRateJump` yet). See [`BoundedSSA`](@ref).
 
 ## RegularJump Compatible Methods
 

src/JumpProcesses.jl

@@ -11,6 +11,8 @@ using Random: Random, randexp, seed!
 using DocStringExtensions: DocStringExtensions, FIELDS, TYPEDEF
 using DataStructures: DataStructures, MutableBinaryMinHeap, sizehint!, top_with_handle
 using PoissonRandom: PoissonRandom, pois_rand
+using Distributions: Distributions, Bernoulli
+using SciMLStructures: SciMLStructures
 using ArrayInterface: ArrayInterface
 using FunctionWrappers: FunctionWrappers
 using Graphs: Graphs, AbstractGraph, dst, grid, src
@@ -133,6 +135,10 @@ export SSAStepper
 include("simple_regular_solve.jl")
 export SimpleTauLeaping, SimpleExplicitTauLeaping, EnsembleGPUKernel
 
+# BoundedSSA: uniformization SSA solver (differentiable when StochasticAD is loaded)
+include("bounded_ssa.jl")
+export BoundedSSA
+
 # spatial:
 include("spatial/spatial_massaction_jump.jl")
 export SpatialMassActionJump

src/bounded_ssa.jl

@@ -0,0 +1,295 @@
+
+"""
+    BoundedSSA(; rate_bound)
+
+A StochasticAD-compatible SSA algorithm for **jump-only** `ConstantRateJump` /
+`MassActionJump` `DiscreteProblem`s, giving correct gradients via StochasticAD's
+`derivative_estimate`/`stochastic_triple` — with `saveat` support, so the whole
+sampled path is differentiable, not only the terminal state.
+
+The stock `SSAStepper` cannot be differentiated with StochasticAD: it advances
+time with a `while integrator.t < integrator.tstop < end_time` loop, i.e. a
+boolean predicate on (triple-valued) time, which StochasticAD forbids by design —
+so the event-count derivative is dropped (a state-dependent rate yields a gradient
+of `0`). `BoundedSSA` instead uses **uniformization (thinning)** against a fixed
+total-propensity bound `Λ = rate_bound`:
+
+  - candidate event times form a homogeneous Poisson process of rate `Λ` on the
+    time span — these are **parameter-free**, so the loop never branches on a
+    triple and the times stay `Float64`;
+  - at each candidate the current total propensity `a(u)` is recomputed and the
+    event is *accepted* with a tracked `Bernoulli(a(u)/Λ)` (otherwise it is a
+    **null event** absorbing the slack `Λ - a(u)`);
+  - the firing channel is chosen by stick-breaking `Bernoulli`s.
+
+All parameter dependence flows through the accept / channel `Bernoulli`s, so the
+gradient is captured. This is **unbiased** (no step cap) whenever `Λ` is a valid
+bound, and `saveat` is exact because the candidate times are fixed `Float64`.
+
+With ordinary `Float64` parameters `solve(jprob, BoundedSSA(; rate_bound))` is an
+ordinary (uniformization) SSA simulation; with StochasticAD triples it
+differentiates.
+
+# Keyword arguments
+
+  - `rate_bound` (required): a constant `Λ` upper-bounding the **total** propensity
+    `Σₖ rateₖ(u, p, t)` over the whole trajectory (and over the parameter
+    perturbation). Valid for systems with rigorously bounded populations; a looser
+    bound only costs efficiency (more null events), not accuracy. If `Λ` is
+    violated the accept probability exceeds 1 and sampling errors — pick it with
+    margin.
+
+# `solve` options
+
+  - `saveat`: times (a vector, or a `Number` step) at which to return the solution,
+    with `save_start`/`save_end` controlling the endpoints (same conventions as
+    `SimpleTauLeaping`, via `_process_saveat`); defaults to `[t0, tf]`. `sol.u[i]` is
+    the differentiable state at `sol.t[i]`, and `sol(t)` interpolates (piecewise
+    constant, as with `SSAStepper`).
+
+# Scope / limitations
+
+  - `ConstantRateJump`s and `MassActionJump`s (state-dependent / mass-action rates
+    supported); jump-only, no continuous drift, no `VariableRateJump`.
+  - Additive affects only. A `ConstantRateJump`'s net change is inferred from `affect!`
+    and checked; a `MassActionJump`'s `net_stoch` is inherently additive. For a
+    `MassActionJump` rate constant to be *differentiated* it must flow from `p` via a
+    `param_idxs`/`param_mapper` jump (the combinatoric scaling is matched); a jump with
+    fixed numeric `scaled_rates` still simulates, but those constants carry no derivative.
+    Note: MTK/Catalyst-generated mass-action jumps *simulate* under `BoundedSSA` but do
+    not yet *differentiate* their rate constants — MTK's mass-action parameter mapper
+    coerces rates to `Float64`, which drops the `StochasticTriple` (a documented follow-up).
+  - The differentiation parameter `prob.p` may be a plain numeric collection (e.g. a
+    `Vector`) or a SciMLStructures parameter object (MTK/Catalyst `MTKParameters`); in
+    the latter case the differentiable **tunable** portion is the target (extracted via
+    `SciMLStructures.canonicalize`), matching how the rest of JumpProcesses treats `p`.
+
+Internally this wraps `JumpProcesses.bounded_ssa_path`, the (unexported)
+differentiable core; `solve(jprob, BoundedSSA(; rate_bound))` is the public entry.
+"""
+struct BoundedSSA{B} <: SciMLBase.AbstractDEAlgorithm
+    rate_bound::B
+end
+function BoundedSSA(; rate_bound = nothing)
+    rate_bound === nothing && error("BoundedSSA requires the keyword argument " *
+        "`rate_bound` (a constant upper bound on the total propensity).")
+    BoundedSSA{typeof(rate_bound)}(rate_bound)
+end
+
+mutable struct BoundedSSAShim{U, P, T}
+    u::U
+    p::P
+    t::T
+end
+
+function _bssa_net_change(affect!, ubase, p, t0)
+    u = collect(ubase)
+    affect!(BoundedSSAShim(u, p, t0))
+    return u .- ubase
+end
+
+# infer a jump's additive net state change, verifying it is state-independent.
+function _bssa_additive_change(jump, u0, p, t0)
+    base = float.(collect(u0))
+    Δ = _bssa_net_change(jump.affect!, base, p, t0)
+    Δ2 = _bssa_net_change(jump.affect!, base .+ one(eltype(base)), p, t0)
+    isapprox(Δ, Δ2) || error(
+        "BoundedSSA supports only additive affects (a constant net state change), " *
+        "but a jump's affect! gave a state-dependent change ($Δ vs $Δ2 from a " *
+        "shifted state).")
+    return Δ
+end
+
+function _bssa_check_supported(jprob)
+    jprob.prob isa DiscreteProblem || error(
+        "BoundedSSA only supports JumpProblems over DiscreteProblems (pure jumps).")
+    vj = jprob.variable_jumps
+    (vj === nothing || isempty(vj)) || error(
+        "BoundedSSA supports jump-only problems only (no VariableRateJumps).")
+    cj = jprob.constant_jumps
+    nc = (cj === nothing) ? 0 : length(cj)
+    nm = get_num_majumps(jprob.massaction_jump)
+    (nc + nm >= 1) || error(
+        "BoundedSSA requires at least one ConstantRateJump or MassActionJump.")
+    nothing
+end
+
+# Differentiable parameters used to choose the state's element type. Arrays are used
+# directly; SciMLStructures objects use their tunable values, so injected AD types
+# promote the state. Scalars and NullParameters are returned unchanged.
+function _bssa_tunables(p)
+    SciMLStructures.isscimlstructure(p) ?
+    SciMLStructures.canonicalize(SciMLStructures.Tunable(), p)[1] : p
+end
+
+# Dense additive net state change of MassActionJump reaction `r` (from its net_stoch),
+# over `n` species. MassAction affects are inherently additive, so -- unlike a
+# ConstantRateJump -- no affect! probing / verification is needed.
+function _bssa_ma_delta(net_stoch_r, n)
+    Δ = zeros(Float64, n)
+    for (spec, change) in net_stoch_r
+        Δ[spec] += change
+    end
+    return Δ
+end
+
+# Propensity of MassActionJump reaction `r` at state `u`. `unscaled` is the per-reaction
+# unscaled rate constants from `param_mapper(p)` (so a StochasticTriple parameter flows
+# through), or `nothing` to fall back to the jump's stored (constant) scaled_rates. The
+# mass-action factor is the falling factorial ∏_j (u[spec] - j): it is naturally zero
+# when a reactant is insufficient (integer populations), so no boolean on the (triple)
+# population is needed, and BoundedSSA never calls `evalrxrate` (avoiding its `::R`
+# return assertion). Combinatoric scaling matches the stock aggregator via `scalerate`.
+function _bssa_ma_rate(u, maj, unscaled, r)
+    sr = unscaled === nothing ? maj.scaled_rates[r] :
+         (maj.rescale_rates_on_update ? scalerate(unscaled[r], maj.reactant_stoch[r]) :
+          unscaled[r])
+    val = sr
+    for (spec, s) in maj.reactant_stoch[r]
+        pop = u[spec]
+        for j in 0:(s - 1)
+            val = val * (pop - j)
+        end
+    end
+    return val
+end
+
+# Internal driver: returns `(tsave, usave)` at the resolved save schedule. Uses
+# `_process_saveat` (shared with SimpleTauLeaping) for saveat/save_start/save_end.
+function _bounded_ssa(jprob, p, Λ, tspan, saveat, save_start, save_end)
+    _bssa_check_supported(jprob)
+    u0 = jprob.prob.u0
+    cjumps = jprob.constant_jumps === nothing ? () : jprob.constant_jumps
+    maj = jprob.massaction_jump
+    t0, tf = first(tspan), last(tspan)
+    ΔT = tf - t0
+    Kc = length(cjumps)
+    nrx = get_num_majumps(maj)
+    K = Kc + nrx
+    n = length(u0)
+
+    saveat_times, ss, se = _process_saveat(saveat, (t0, tf), save_start, save_end)
+
+    # additive net change per channel: ConstantRateJumps (net change inferred from
+    # affect! and verified additive) first, then MassActionJump reactions (net_stoch).
+    Δ = Vector{Vector{Float64}}(undef, K)
+    for k in 1:Kc
+        Δ[k] = _bssa_additive_change(cjumps[k], u0, p, t0)
+    end
+    for r in 1:nrx
+        Δ[Kc + r] = _bssa_ma_delta(maj.net_stoch[r], n)
+    end
+
+    # `0 * sum(...)` promotes the state to the parameter's element type (giving a triple
+    # zero when a StochasticTriple flows in). We seed from the **tunable** numeric portion
+    # of `p` (see `_bssa_tunables`), so both plain `Vector` parameters and SciMLStructures
+    # parameter objects (MTK/Catalyst tunables) are supported; a `Vector` seeds from itself.
+    z = 0 * sum(_bssa_tunables(p))
+    u = [float(u0[i]) + z for i in 1:n]
+
+    tsave = typeof(t0)[]
+    usave = typeof(u)[]
+    if ss
+        push!(tsave, t0)
+        push!(usave, copy(u))
+    end
+
+    # candidate events ~ homogeneous Poisson(Λ) on [t0, tf]. PARAMETER-FREE (Λ is a
+    # constant), so the count and times carry no derivative and never branch on a
+    # triple. Uses PoissonRandom's `pois_rand`, as elsewhere in JumpProcesses.
+    M = pois_rand(Random.default_rng(), Λ * ΔT)
+    ctimes = sort!(t0 .+ ΔT .* rand(M))
+
+    save_idx = 1
+    for m in 1:M
+        tm = @inbounds ctimes[m]
+        while save_idx <= length(saveat_times) && @inbounds(saveat_times[save_idx]) < tm
+            push!(tsave, @inbounds saveat_times[save_idx])
+            push!(usave, copy(u))
+            save_idx += 1
+        end
+
+        # per-channel propensities at the current state: ConstantRateJumps then
+        # MassActionJump reactions. MA rate constants come from `param_mapper(p)` (so a
+        # StochasticTriple parameter flows through), else the stored scaled_rates.
+        maunscaled = (nrx > 0 && using_params(maj)) ? maj.param_mapper(p) : nothing
+        rates = [k <= Kc ? cjumps[k].rate(u, p, tm) :
+                 _bssa_ma_rate(u, maj, maunscaled, k - Kc) for k in 1:K]
+        total = sum(rates)
+        # thinning: real vs null event. `rand(Bernoulli(p))` handles both the primal
+        # draw and — when a StochasticTriple `p` flows in with StochasticAD loaded —
+        # the differentiable decision (StochasticAD's own `rand(::Bernoulli)` rule).
+        accept = rand(Bernoulli(total / Λ))
+
+        # which channel: stick-breaking conditional Bernoullis (last deterministic)
+        notchosen = 1 + z
+        sel = [z for _ in 1:n]
+        for k in 1:K
+            chose = k < K ?
+                    rand(Bernoulli(rates[k] / (sum(rates[j] for j in k:K) + 1e-300))) :
+                    (1 + z)
+            take = notchosen * chose
+            sel = [sel[i] + take * Δ[k][i] for i in 1:n]
+            notchosen = notchosen * (1 - chose)
+        end
+
+        u = [u[i] + accept * sel[i] for i in 1:n]        # apply only on a real event
+    end
+    while save_idx <= length(saveat_times)
+        push!(tsave, @inbounds saveat_times[save_idx])
+        push!(usave, copy(u))
+        save_idx += 1
+    end
+    if se
+        push!(tsave, tf)
+        push!(usave, copy(u))
+    end
+    return tsave, usave
+end
+
+"""
+    bounded_ssa_path(jprob, p; rate_bound, saveat = tf, save_start = nothing,
+                     save_end = nothing, tspan = jprob.prob.tspan)
+
+Differentiable core behind [`BoundedSSA`](@ref): simulate the jump-only
+`ConstantRateJump` / `MassActionJump` process by uniformization against the constant
+total-propensity bound `rate_bound`, and return the state at each save time as a `Vector` of state
+vectors. When a `StochasticTriple` parameter flows in (StochasticAD loaded) the result
+is differentiable, so this can be wrapped in `derivative_estimate`:
+
+```julia
+derivative_estimate(p0[k]) do pk
+    pv = [j == k ? pk : oftype(pk, p0[j]) for j in eachindex(p0)]
+    bounded_ssa_path(jprob, pv; rate_bound = Λ, saveat = [tf])[end][1]
+end
+```
+
+`saveat`/`save_start`/`save_end` follow the usual JumpProcesses conventions (via
+`_process_saveat`, as `SimpleTauLeaping`). `p` may be a plain numeric collection (e.g. a
+`Vector`) or a SciMLStructures parameter object (MTK/Catalyst), whose tunable portion is
+the differentiation target. See [`BoundedSSA`](@ref) for the method and the
+meaning/validity of `rate_bound`.
+"""
+function bounded_ssa_path(jprob, p; rate_bound, saveat = last(jprob.prob.tspan),
+        save_start = nothing, save_end = nothing, tspan = jprob.prob.tspan)
+    _, usave = _bounded_ssa(jprob, p, rate_bound, tspan, saveat, save_start, save_end)
+    return usave
+end
+
+# solve(jprob, BoundedSSA(; rate_bound); saveat, save_start, save_end). Defined as
+# `solve` (like SimpleTauLeaping) since BoundedSSA is self-contained and does not use
+# the integrator/init machinery. `sol(t)` works via piecewise-constant interpolation.
+function DiffEqBase.solve(jump_prob::JumpProblem, alg::BoundedSSA;
+        seed = nothing, saveat = nothing, save_start = nothing, save_end = nothing,
+        tspan = jump_prob.prob.tspan, kwargs...)
+    seed === nothing || Random.seed!(seed)
+    prob = jump_prob.prob
+    ts, us = _bounded_ssa(jump_prob, prob.p, alg.rate_bound, tspan, saveat,
+        save_start, save_end)
+    SciMLBase.build_solution(prob, alg, ts, us;
+        dense = true,
+        interp = SciMLBase.ConstantInterpolation(ts, us),
+        calculate_error = false,
+        stats = DiffEqBase.Stats(0),
+        retcode = ReturnCode.Success)
+end

test/qa.jl

@@ -31,6 +31,8 @@ run_qa(
                 :AbstractQ,
                 # FunctionWrappers non-public
                 :FunctionWrapper,
+                # SciMLStructures non-public
+                :Tunable, :canonicalize, :isscimlstructure,
             ),
         ),
         all_explicit_imports_are_public = (;