diff --git a/Project.toml b/Project.toml index a2ff93e3..02178ff7 100644 --- a/Project.toml +++ b/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" diff --git a/docs/src/api.md b/docs/src/api.md index d8fd944b..00c83088 100644 --- a/docs/src/api.md +++ b/docs/src/api.md @@ -9,6 +9,7 @@ CurrentModule = JumpProcesses ```@docs JumpProblem SSAStepper +BoundedSSA reset_aggregated_jumps! ``` diff --git a/docs/src/jump_solve.md b/docs/src/jump_solve.md index 41ac4412..4e379743 100644 --- a/docs/src/jump_solve.md +++ b/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 diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl index d6ce8bfe..a36c65d9 100644 --- a/src/JumpProcesses.jl +++ b/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 diff --git a/src/bounded_ssa.jl b/src/bounded_ssa.jl new file mode 100644 index 00000000..7a4be32f --- /dev/null +++ b/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 diff --git a/test/qa.jl b/test/qa.jl index 685ea29d..05e7beb8 100644 --- a/test/qa.jl +++ b/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 = (; diff --git a/test/runtests.jl b/test/runtests.jl index 9a2dfc24..50de7445 100644 --- a/test/runtests.jl +++ b/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 diff --git a/test/stochasticad/Project.toml b/test/stochasticad/Project.toml new file mode 100644 index 00000000..00871ca6 --- /dev/null +++ b/test/stochasticad/Project.toml @@ -0,0 +1,7 @@ +[deps] +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SciMLStructures = "53ae85a6-f571-4167-b2af-e1d143709226" +Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +StochasticAD = "e4facb34-4f7e-4bec-b153-e122c37934ac" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" diff --git a/test/stochasticad_tests.jl b/test/stochasticad_tests.jl new file mode 100644 index 00000000..ff37ca55 --- /dev/null +++ b/test/stochasticad_tests.jl @@ -0,0 +1,182 @@ +using JumpProcesses, StochasticAD, Distributions, SciMLStructures +using Statistics, Random, Test + +# Tests for BoundedSSA — a uniformization (thinning) SSA for jump-only +# ConstantRateJump/MassActionJump DiscreteProblems that composes with StochasticAD. +# Everything is exercised through the PUBLIC interface only: +# solve(jprob, BoundedSSA(; rate_bound); saveat = ...) +# To differentiate, the JumpProblem is (re)built with the StochasticTriple parameter +# inside the estimator closure, then solved — no internal helpers are called. +# Runs in an isolated environment (GROUP=StochasticAD): needs only JumpProcesses + +# StochasticAD + Distributions (no ODE solver). + +# per-partial StochasticAD gradient with fixed seeds (reproducible) +function sad_partial(f, p0, k; N) + s = Vector{Float64}(undef, N) + for i in 1:N + Random.seed!(i) + s[i] = derivative_estimate(p0[k]) do pk + p = [j == k ? pk : oftype(pk, p0[j]) for j in eachindex(p0)] + f(p) + end + end + return mean(s), std(s) / sqrt(N) +end + +# primal Monte-Carlo mean + standard error +function pmean(f; N) + s = [f() for _ in 1:N] + return mean(s), std(s) / sqrt(N) +end + +# A minimal SciMLStructures parameter object standing in for an MTK/Catalyst +# `MTKParameters`: its differentiable tunables live in `.tunables`. Used to check that +# BoundedSSA seeds from and differentiates through the tunable portion of a *structured* +# parameter, not only plain `Vector`s. Only the `Tunable` portion is implemented. +struct TunableParams{T} + tunables::T +end +SciMLStructures.isscimlstructure(::TunableParams) = true +function SciMLStructures.canonicalize(::SciMLStructures.Tunable, p::TunableParams) + p.tunables, TunableParams, true +end + +@testset "BoundedSSA StochasticAD (public solve interface)" begin + + # --- Test A: pure death, STATE-DEPENDENT rate (the case that matters) ------ + # rate = μ·u ; E[u(T)] = u0·e^{-μT} ; d/dμ E[u(T)] = -T·u0·e^{-μT} + # A naive `while t p[1] * u[1], integ -> (integ.u[1] -= 1; nothing)) + analytic = -T * u0 * exp(-μ0 * T) + g, se = sad_partial([μ0], 1; N = 4000) do p + jp = JumpProblem(DiscreteProblem([u0], (0.0, T), p), Direct(), death) + solve(jp, BoundedSSA(; rate_bound = Λ); saveat = [T]).u[end][1] + end + @test abs(g - analytic) < 4 * se # event-count derivative captured + @test abs(g) > 1.0 # explicitly NOT the zero a naive SSA gives + end + + # --- Test B: birth-death, multi-channel + state-dependent ------------------ + # birth λ, death μ·u ; E[u(T)] = λ/μ + (u0-λ/μ)e^{-μT} + @testset "birth-death (multi-channel, state-dependent)" begin + T, u0, λ0, μ0, Λ = 1.0, 50, 10.0, 0.3, 60.0 # total = λ + μ·u ≤ ~28 ≪ Λ + birth = ConstantRateJump((u, p, t) -> p[1], integ -> (integ.u[1] += 1; nothing)) + death = ConstantRateJump((u, p, t) -> p[2] * u[1], integ -> (integ.u[1] -= 1; nothing)) + a, b = λ0 / μ0, exp(-μ0 * T) + analytic = [(1 - b) / μ0, -λ0 / μ0^2 * (1 - b) + (u0 - a) * (-T * b)] + for k in 1:2 + g, se = sad_partial([λ0, μ0], k; N = 4000) do p + jp = JumpProblem(DiscreteProblem([u0], (0.0, T), p), Direct(), birth, death) + solve(jp, BoundedSSA(; rate_bound = Λ); saveat = [T]).u[end][1] + end + @test abs(g - analytic[k]) < 4 * se + end + end + + # --- Test C: saveat returns the path at intermediate times (primal) -------- + # pure death: E[u(s)] = u0·e^{-μs} at every save time s. + @testset "saveat path (intermediate times)" begin + T, u0, μ0, Λ = 1.0, 100, 0.5, 60.0 + death = ConstantRateJump((u, p, t) -> p[1] * u[1], integ -> (integ.u[1] -= 1; nothing)) + jprob = JumpProblem(DiscreteProblem([u0], (0.0, T), [μ0]), Direct(), death) + sat = [0.0, 0.5, 1.0] + for (idx, s) in enumerate(sat) + m, se = pmean(N = 4000) do + solve(jprob, BoundedSSA(; rate_bound = Λ); saveat = sat).u[idx][1] + end + # `s = 0` is deterministic (u(0) = u0, se = 0); the `+ 1e-9` keeps the + # zero-variance point from failing `0 < 0`. + @test abs(m - u0 * exp(-μ0 * s)) <= 4 * se + 1e-9 + end + end + + # --- Test D: primal mean matches the stock SSA ----------------------------- + @testset "primal mean matches stock SSAStepper" begin + T, u0, μ0, Λ = 1.0, 100, 0.4, 60.0 + death = ConstantRateJump((u, p, t) -> p[1] * u[1], integ -> (integ.u[1] -= 1; nothing)) + jprob = JumpProblem(DiscreteProblem([u0], (0.0, T), [μ0]), Direct(), death) + mb, seb = pmean(N = 4000) do + solve(jprob, BoundedSSA(; rate_bound = Λ); saveat = [T]).u[end][1] + end + ms, ses = pmean(N = 4000) do + solve(jprob, SSAStepper()).u[end][1] + end + @test abs(mb - u0 * exp(-μ0 * T)) < 4 * seb + @test abs(mb - ms) < 4 * sqrt(seb^2 + ses^2) # agree within combined MC error + end + + # --- Test E: solve interface — full path + piecewise-constant interpolation - + @testset "solve returns a full path + interpolates" begin + T, u0, μ0, Λ = 1.0, 100, 0.5, 60.0 + death = ConstantRateJump((u, p, t) -> p[1] * u[1], integ -> (integ.u[1] -= 1; nothing)) + jprob = JumpProblem(DiscreteProblem([u0], (0.0, T), [μ0]), Direct(), death) + sol = solve(jprob, BoundedSSA(; rate_bound = Λ); saveat = [0.0, 0.5, 1.0]) + @test sol.t == [0.0, 0.5, 1.0] + @test length(sol.u) == 3 + @test sol.u[1][1] == u0 # state at t0 is the initial condition + @test 0 <= sol.u[end][1] <= u0 + @test sol(0.25)[1] == sol.u[1][1] # piecewise-constant interpolation + end + + # --- Test F: SciMLStructures parameter object (MTK/Catalyst tunables) ------- + # Same pure-death gradient as Test A, but the parameter is a structured + # SciMLStructures object; BoundedSSA must differentiate through its tunable portion. + @testset "SciMLStructures tunable parameters" begin + T, u0, μ0, Λ = 1.0, 100, 0.5, 60.0 + death = ConstantRateJump((u, p, t) -> p.tunables[1] * u[1], + integ -> (integ.u[1] -= 1; nothing)) + analytic = -T * u0 * exp(-μ0 * T) + g, se = sad_partial([μ0], 1; N = 4000) do p + jp = JumpProblem(DiscreteProblem([u0], (0.0, T), TunableParams(p)), Direct(), death) + solve(jp, BoundedSSA(; rate_bound = Λ); saveat = [T]).u[end][1] + end + @test abs(g - analytic) < 4 * se + @test abs(g) > 1.0 + end + + # --- Test G: MassActionJump, state-dependent rate, differentiated ----------- + # pure death X --> 0 at rate μ, as a MassActionJump with the rate constant taken + # from p (param_idxs). Same analytic as Test A; exercises the MA propensity + + # param_mapper differentiation path (no evalrxrate, no boolean on the population). + @testset "MassActionJump pure death (differentiated)" begin + T, u0, μ0, Λ = 1.0, 100, 0.5, 60.0 + majump = MassActionJump([[1 => 1]], [[1 => -1]]; param_idxs = [1]) + analytic = -T * u0 * exp(-μ0 * T) + g, se = sad_partial([μ0], 1; N = 4000) do p + jp = JumpProblem(DiscreteProblem([u0], (0.0, T), p), Direct(), majump) + solve(jp, BoundedSSA(; rate_bound = Λ); saveat = [T]).u[end][1] + end + @test abs(g - analytic) < 4 * se + @test abs(g) > 1.0 + end + + # --- Test H: order-2 MassActionJump, primal mean matches the stock SSA ------- + # dimerization 2X --> 0 at rate k; propensity = (k/2!)·u·(u-1). Validates the + # falling-factorial propensity (order > 1) AND the combinatoric rate scaling against + # SSAStepper/evalrxrate. max propensity ≈ (k/2)·u0·(u0-1) = 12.25 < Λ. + @testset "MassActionJump dimerization (order 2, primal vs SSAStepper)" begin + T, u0, k0, Λ = 1.0, 50, 0.01, 20.0 + majump = MassActionJump([[1 => 2]], [[1 => -2]]; param_idxs = [1]) + jprob = JumpProblem(DiscreteProblem([u0], (0.0, T), [k0]), Direct(), majump) + mb, seb = pmean(N = 4000) do + solve(jprob, BoundedSSA(; rate_bound = Λ); saveat = [T]).u[end][1] + end + ms, ses = pmean(N = 4000) do + solve(jprob, SSAStepper()).u[end][1] + end + @test abs(mb - ms) < 4 * sqrt(seb^2 + ses^2) + end + + # --- guards: misuse should error, not silently mislead --------------------- + @testset "guards" begin + T = 1.0 + # missing rate_bound on the algorithm + @test_throws ErrorException BoundedSSA() + # non-additive (state-dependent) affect + weird = ConstantRateJump((u, p, t) -> p[1], integ -> (integ.u[1] *= 2; nothing)) + jp_w = JumpProblem(DiscreteProblem([10], (0.0, T), [0.5]), Direct(), weird) + @test_throws ErrorException solve(jp_w, BoundedSSA(; rate_bound = 10.0)) + end +end diff --git a/test/test_groups.toml b/test/test_groups.toml index 443c5cbb..620d9fdc 100644 --- a/test/test_groups.toml +++ b/test/test_groups.toml @@ -6,3 +6,6 @@ versions = ["lts", "1", "pre"] [QA] versions = ["lts", "1"] + +[StochasticAD] +versions = ["lts", "1"]