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..f6423004 --- /dev/null +++ b/src/bounded_ssa.jl @@ -0,0 +1,246 @@ + +""" + 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` 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`. + - 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 + +# 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 + +# 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(...)` 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 + + 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` 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) + 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 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..6d98e17d --- /dev/null +++ b/test/stochasticad_tests.jl @@ -0,0 +1,158 @@ +using JumpProcesses, StochasticAD, Distributions, SciMLStructures +using Statistics, Random, Test + +# Tests for the optional StochasticAD extension: differentiating expectations over +# jump-only ConstantRateJump processes via BoundedSSA (uniformization/thinning). +# Needs only JumpProcesses + StochasticAD + Distributions (no ODE solver), so it +# runs in its own isolated environment (GROUP=StochasticAD). + +# 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 — the +# sole portion BoundedSSA consults (via `_bssa_tunables`). +struct TunableParams{T} + tunables::T +end +SciMLStructures.isscimlstructure(::TunableParams) = true +function SciMLStructures.canonicalize(::SciMLStructures.Tunable, p::TunableParams) + p.tunables, TunableParams, true +end + +@testset "StochasticAD constant-rate jumps (BoundedSSA, uniformization)" 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)) + jprob = JumpProblem(DiscreteProblem([u0], (0.0, T)), Direct(), death) + analytic = -T * u0 * exp(-μ0 * T) + g, se = sad_partial([μ0], 1; N = 10000) do p + JumpProcesses.bounded_ssa_path(jprob, p; rate_bound = Λ, saveat = [T])[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)) + jprob = JumpProblem(DiscreteProblem([u0], (0.0, T)), Direct(), birth, death) + 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 = 10000) do p + JumpProcesses.bounded_ssa_path(jprob, p; rate_bound = Λ, saveat = [T])[end][1] + end + @test abs(g - analytic[k]) < 4 * se + end + end + + # --- Test C: saveat returns the path at intermediate times ----------------- + # 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)), Direct(), death) + sat = [0.0, 0.5, 1.0] + for (idx, s) in enumerate(sat) + m, se = pmean(N = 4000) do + JumpProcesses.bounded_ssa_path(jprob, [μ0]; rate_bound = Λ, saveat = sat)[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 + JumpProcesses.bounded_ssa_path(jprob, [μ0]; rate_bound = Λ, saveat = [T])[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: BoundedSSA solve path (native interface) ---------------------- + @testset "BoundedSSA solve path + saveat" 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)) + # primal: solve returns the full path at saveat + jp0 = JumpProblem(DiscreteProblem([u0], (0.0, T), [μ0]), Direct(), death) + sol = solve(jp0, BoundedSSA(; rate_bound = Λ); saveat = [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 + # differentiate through solve (constructs jprob with triple parameters) + 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 + 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 seed from and differentiate through its + # tunable portion (a plain `Vector` canonicalizes to itself — covered by Test A). + @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)) + jprob = JumpProblem(DiscreteProblem([u0], (0.0, T)), Direct(), death) + analytic = -T * u0 * exp(-μ0 * T) + g, se = sad_partial([μ0], 1; N = 10000) do p + JumpProcesses.bounded_ssa_path(jprob, TunableParams(p); + rate_bound = Λ, saveat = [T])[end][1] + end + @test abs(g - analytic) < 4 * se # tunable-portion derivative captured + @test abs(g) > 1.0 + end + + # --- guards: misuse should error, not silently mislead --------------------- + @testset "guards" begin + T = 1.0 + # MassActionJump not yet supported + majump = MassActionJump([0.5], [[1 => 1]], [[1 => -1]]) + jp_ma = JumpProblem(DiscreteProblem([10], (0.0, T), [0.5]), Direct(), majump) + @test_throws ErrorException JumpProcesses.bounded_ssa_path(jp_ma, [0.5]; rate_bound = 10.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)), Direct(), weird) + @test_throws ErrorException JumpProcesses.bounded_ssa_path(jp_w, [0.5]; 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"]