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"
@@ -37,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"
 ExplicitImports = "1"
 FastBroadcast = "1.3.2, 1"

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,7 @@ 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 ArrayInterface: ArrayInterface
 using FunctionWrappers: FunctionWrappers
 using Graphs: Graphs, AbstractGraph, dst, grid, src
@@ -133,6 +134,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,236 @@
+
+"""
+    BoundedSSA(; rate_bound)
+
+A StochasticAD-compatible SSA algorithm for **jump-only** `ConstantRateJump`
+`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 only (state-dependent rates supported); jump-only, no
+    continuous drift, no `VariableRateJump`. `MassActionJump` is not yet supported.
+  - Additive affects only (the net change is inferred from `affect!` and checked).
+  - The differentiation parameter `prob.p` must be a summable/indexable numeric
+    collection (e.g. a `Vector`). SciMLStructures parameter objects (MTK/Catalyst
+    tunables) are not yet specialized — a documented follow-up.
+  - The solver itself is plain (no StochasticAD dependency): with ordinary parameters
+    `solve(jprob, BoundedSSA(; rate_bound))` is a uniformization SSA simulation. It
+    becomes differentiable when the user loads `StochasticAD` and passes a
+    `StochasticTriple` parameter — StochasticAD's own `rand(::Bernoulli)` rule makes
+    the accept/channel decisions differentiable, with no glue needed from this package.
+
+Internally this wraps `JumpProcesses.bounded_ssa_path`, the (unexported)
+differentiable core; `solve(jprob, BoundedSSA(; rate_bound))` is the public entry.
+"""
+struct BoundedSSA{B} <: DiffEqBase.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).")
+    maj = jprob.massaction_jump
+    (maj === nothing || get_num_majumps(maj) == 0) || error(
+        "BoundedSSA does not yet support MassActionJump; use ConstantRateJumps.")
+    vj = jprob.variable_jumps
+    (vj === nothing || isempty(vj)) || error(
+        "BoundedSSA supports jump-only constant-rate problems only (no VariableRateJumps).")
+    cj = jprob.constant_jumps
+    (cj === nothing || isempty(cj)) &&
+        error("BoundedSSA requires at least one ConstantRateJump.")
+    nothing
+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
+    jumps = jprob.constant_jumps
+    t0, tf = first(tspan), last(tspan)
+    ΔT = tf - t0
+    K = length(jumps)
+    n = length(u0)
+
+    saveat_times, ss, se = _process_saveat(saveat, (t0, tf), save_start, save_end)
+
+    Δ = [_bssa_additive_change(jumps[k], u0, p, t0) for k in 1:K]
+
+    # `0 * sum(p)` promotes the state to the parameter's element type (giving a triple
+    # zero when a StochasticTriple flows in). This assumes `p` is a summable/indexable
+    # numeric collection, e.g. a `Vector` — SciMLStructures parameter objects
+    # (MTK/Catalyst tunables) are not yet specialized. See BoundedSSA docs.
+    z = 0 * sum(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
+
+        rates = [jumps[k].rate(u, p, tm) for k in 1:K]   # recomputed at current state
+        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` 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` must be a summable/indexable numeric
+collection (e.g. a `Vector`); SciMLStructures parameter objects are not yet specialized.
+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)
+    DiffEqBase.build_solution(prob, alg, ts, us;
+        dense = true,
+        interp = DiffEqBase.ConstantInterpolation(ts, us),
+        calculate_error = false,
+        stats = DiffEqBase.Stats(0),
+        retcode = ReturnCode.Success)
+end

test/runtests.jl

@@ -9,6 +9,16 @@ function activate_gpu_env()
     Pkg.instantiate()
 end
 
+# Isolated environment for the StochasticAD extension tests. StochasticAD pins
+# old transitive deps (e.g. ForwardDiff 0.10) that conflict with the modern
+# OrdinaryDiffEq stack, so it is kept out of the main test target and run here
+# in its own project (no ODE solver needed -- the extension never calls `solve`).
+function activate_stochasticad_env()
+    Pkg.activate(joinpath(@__DIR__, "stochasticad"))
+    Pkg.develop(PackageSpec(path = dirname(@__DIR__)))
+    Pkg.instantiate()
+end
+
 @time begin
     if GROUP == "QA"
         @time @safetestset "QA Tests" begin include("qa.jl") end
@@ -63,6 +73,14 @@ end
         @time @safetestset "GPU Tau Leaping test" begin include("gpu/regular_jumps.jl") end
     end
 
+    if GROUP == "StochasticAD"
+        # Set up the isolated `stochasticad` project, then run its tests in a fresh
+        # Julia process so it does not clash with the main test environment. If the
+        # tests fail, `run` makes the overall suite fail too.
+        activate_stochasticad_env()
+        @time run(`$(Base.julia_cmd()) --project=$(joinpath(@__DIR__, "stochasticad")) $(joinpath(@__DIR__, "stochasticad_tests.jl"))`)
+    end
+
     if GROUP == "Correctness"
         activate_gpu_env()
     end

test/stochasticad/Project.toml

@@ -0,0 +1,6 @@
+[deps]
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
+StochasticAD = "e4facb34-4f7e-4bec-b153-e122c37934ac"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"