From fc613a3c36c08f3236c8d24fe24616f05798426d Mon Sep 17 00:00:00 2001 From: Roman <116185879+RomanSahakyan03@users.noreply.github.com> Date: Wed, 3 Jun 2026 17:38:16 +0400 Subject: [PATCH] Add optional StochasticAD extension for differentiating jump processes Adds JumpProcessesStochasticADExt, an optional package extension (loaded only when StochasticAD and Distributions are present) that lets StochasticAD's derivative_estimate compute gradients of expectations over variable-rate jump processes. Differentiating through the solver directly -- derivative_estimate(g -> solve(jprob(g), Tsit5())) -- does not compose: the adaptive OrdinaryDiffEq internals require Number-interface methods that StochasticTriple intentionally omits, and VR_Direct's ContinuousCallback rootfind is a boolean predicate on a triple. The extension instead provides a triple-generic, fixed-grid PDMP simulator (generic RK4 + per-channel Bernoulli with first-fire priority + multiplicative-select update) that composes by design. - src/JumpProcesses.jl: fixedgrid_simulate / fixedgrid_jump_observable stubs and exports (no StochasticAD code in src/) - ext/JumpProcessesStochasticADExt.jl: the simulator and a JumpProblem adapter - Project.toml: StochasticAD and Distributions as weakdeps/extension + compat - test/stochasticad_tests.jl: analytic-gradient and MCWF tests (registered) Validated against an analytic benchmark (0.7 sigma) and the standalone Stage-2 prototype gradient (agreement within ~1 sigma on both partials). --- Project.toml | 9 ++- ext/JumpProcessesStochasticADExt.jl | 111 ++++++++++++++++++++++++++++ src/JumpProcesses.jl | 8 ++ test/runtests.jl | 1 + test/stochasticad_tests.jl | 74 +++++++++++++++++++ 5 files changed, 202 insertions(+), 1 deletion(-) create mode 100644 ext/JumpProcessesStochasticADExt.jl create mode 100644 test/stochasticad_tests.jl diff --git a/Project.toml b/Project.toml index 9b2401234..4f38753c8 100644 --- a/Project.toml +++ b/Project.toml @@ -22,12 +22,15 @@ SymbolicIndexingInterface = "2efcf032-c050-4f8e-a9bb-153293bab1f5" [weakdeps] Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" OrdinaryDiffEqCore = "bbf590c4-e513-4bbe-9b18-05decba2e5d8" +StochasticAD = "e4facb34-4f7e-4bec-b153-e122c37934ac" [extensions] JumpProcessesKernelAbstractionsExt = ["Adapt", "KernelAbstractions"] JumpProcessesOrdinaryDiffEqCoreExt = "OrdinaryDiffEqCore" +JumpProcessesStochasticADExt = ["StochasticAD", "Distributions"] [compat] ADTypes = "1" @@ -37,6 +40,7 @@ ArrayInterface = "7.15" DataStructures = "0.18, 0.19" DiffEqBase = "6.192, 7" DiffEqCallbacks = "4.7" +Distributions = "0.25" DocStringExtensions = "0.9" ExplicitImports = "1" FastBroadcast = "0.3, 1" @@ -58,6 +62,7 @@ SciMLBase = "2.115, 3.1" StableRNGs = "1" StaticArrays = "1.9.8" Statistics = "1" +StochasticAD = "0.1" StochasticDiffEq = "6.82, 7" SymbolicIndexingInterface = "0.3.36" Test = "1" @@ -66,6 +71,7 @@ julia = "1.10" [extras] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" +Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -76,8 +82,9 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f" SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f" StableRNGs = "860ef19b-820b-49d6-a774-d7a799459cd3" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" +StochasticAD = "e4facb34-4f7e-4bec-b153-e122c37934ac" StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["ADTypes", "Aqua", "ExplicitImports", "FastBroadcast", "LinearAlgebra", "LinearSolve", "OrdinaryDiffEq", "OrdinaryDiffEqFunctionMap", "Pkg", "SafeTestsets", "StableRNGs", "Statistics", "StochasticDiffEq", "Test"] +test = ["ADTypes", "Aqua", "Distributions", "ExplicitImports", "FastBroadcast", "LinearAlgebra", "LinearSolve", "OrdinaryDiffEq", "OrdinaryDiffEqFunctionMap", "Pkg", "SafeTestsets", "StableRNGs", "Statistics", "StochasticAD", "StochasticDiffEq", "Test"] diff --git a/ext/JumpProcessesStochasticADExt.jl b/ext/JumpProcessesStochasticADExt.jl new file mode 100644 index 000000000..0250bb765 --- /dev/null +++ b/ext/JumpProcessesStochasticADExt.jl @@ -0,0 +1,111 @@ +module JumpProcessesStochasticADExt + +# Optional StochasticAD support for JumpProcesses. +# +# Why a separate fixed-grid simulator instead of differentiating `solve`? +# `derivative_estimate(γ -> solve(jprob(γ), Tsit5()))` does not compose: (1) the +# adaptive OrdinaryDiffEq solver internals assume the full `Number` interface +# that `ForwardDiff.Dual` implements but `StochasticAD.StochasticTriple` omits +# (so triples can't even enter the ODE state), and (2) `VR_Direct` locates jump +# times with a `ContinuousCallback` rootfind, i.e. a boolean predicate on a +# triple, which StochasticAD forbids by design. Both were established +# empirically. This extension instead provides a triple-generic, fixed-grid +# PDMP simulator that composes *by design* (pure generic arithmetic + +# `rand(Bernoulli)`), which is what "AD on the algorithm" amounts to here. + +using JumpProcesses +using StochasticAD +using Distributions: Bernoulli + +# generic RK4 step over a triple-safe out-of-place drift +@inline function _rk4(drift, u, p, t, dt) + k1 = drift(u, p, t) + k2 = drift([u[i] + 0.5dt * k1[i] for i in eachindex(u)], p, t + 0.5dt) + k3 = drift([u[i] + 0.5dt * k2[i] for i in eachindex(u)], p, t + 0.5dt) + k4 = drift([u[i] + dt * k3[i] for i in eachindex(u)], p, t + dt) + return [u[i] + (dt / 6) * (k1[i] + 2k2[i] + 2k3[i] + k4[i]) for i in eachindex(u)] +end + +""" + fixedgrid_simulate(drift, channels, u0, p, tspan, dt) -> u(tspan[2]) + +Triple-safe, fixed-grid piecewise-deterministic Markov process simulator that +composes with StochasticAD's `derivative_estimate`/`stochastic_triple`. + +On each step of size `dt`: integrate the smooth flow with a generic RK4, then for +each jump channel sample `rand(Bernoulli(rateₖ(u,p,t) * dt))` and apply the +post-jump state of the first channel that fires (first-fire priority), all via a +multiplicative-select blend. Using one `Bernoulli` per channel (rather than one +`Bernoulli` for "any jump" plus a `Categorical` for "which") avoids division, so +there is no `Categorical` 0/0 when the total rate vanishes. + +Arguments: + - `drift(u, p, t) -> du`: triple-safe, out-of-place ODE RHS. + - `channels`: iterable of named tuples `(rate, post)` where + `rate(u, p, t) -> ≥0 scalar` and `post(u, p, t) -> post-jump state vector`. + - `u0`, `p`, `tspan`, `dt`: initial state, parameters, `(t0, tf)`, step. + +Validity: this is an O(dt²)-per-step (τ-leap-style) approximation to the exact +continuous-time PDMP, accurate when `rateₖ · dt ≪ 1`. +""" +function JumpProcesses.fixedgrid_simulate(drift, channels, u0, p, tspan, dt) + nsteps = round(Int, (tspan[2] - tspan[1]) / dt) + ptype = sum(p) # carries the (maybe-triple) type + u = [u0i + 0 * ptype for u0i in u0] # promote state to p's type + n = length(u) + t0 = tspan[1] + + for s in 1:nsteps + t = t0 + (s - 1) * dt + u_smooth = _rk4(drift, u, p, t, dt) + + remaining = 1 + 0 * ptype # triple, value 1 + collapse = [0 * ptype for _ in 1:n] # triple zeros + for ch in channels + r = ch.rate(u, p, t) + fired = rand(Bernoulli(r * dt)) # triple 0/1 + take = fired * remaining # 1 iff first firing channel + pj = ch.post(u, p, t) + collapse = [collapse[i] + take * pj[i] for i in 1:n] + remaining = remaining * (1 - fired) + end + any_fired = 1 - remaining + + u = [any_fired * collapse[i] + (1 - any_fired) * u_smooth[i] for i in 1:n] + end + return u +end + +""" + fixedgrid_jump_observable(jprob, p, post_jumps, observable; dt) -> scalar + +Run [`fixedgrid_simulate`](@ref) over the model bundled in a `JumpProblem`, then +return `observable(u(tf))`. Designed to be wrapped in `derivative_estimate`: + +```julia +g = derivative_estimate(p0[k]) do pk + p = [j == k ? pk : oftype(pk, p0[j]) for j in eachindex(p0)] + fixedgrid_jump_observable(jprob, p, post_jumps, observable; dt) +end +``` + +The drift (`jprob.prob.f`) and the channel rates (`jprob.variable_jumps[k].rate`) +are taken from `jprob`; `post_jumps[k](u,p,t)` supplies channel `k`'s post-jump +state (the `VariableRateJump` `affect!`s are *not* used, since arbitrary mutating +affects are not StochasticAD-safe). All of `drift`, the rates, `post_jumps`, and +`observable` must be triple-safe (generic arithmetic; no `ComplexF64(::triple)`, +`fill!(u, 0.0)`, or boolean branch on state). +""" +function JumpProcesses.fixedgrid_jump_observable(jprob, p, post_jumps, observable; dt) + prob = jprob.prob + f = prob.f + drift = JumpProcesses.isinplace(prob) ? + ((u, pp, t) -> (du = similar(u); f(du, u, pp, t); du)) : + ((u, pp, t) -> f(u, pp, t)) + vj = jprob.variable_jumps + channels = [(rate = vj[k].rate, post = post_jumps[k]) for k in eachindex(vj)] + u = JumpProcesses.fixedgrid_simulate(drift, channels, prob.u0, p, prob.tspan, dt) + return observable(u) +end + +end # module diff --git a/src/JumpProcesses.jl b/src/JumpProcesses.jl index 87c8c2c26..85b78f630 100644 --- a/src/JumpProcesses.jl +++ b/src/JumpProcesses.jl @@ -112,6 +112,14 @@ include("aggregators/aggregated_api.jl") include("variable_rate.jl") export VariableRateAggregator, VR_FRM, VR_Direct, VR_DirectFW +# StochasticAD support. These are stubs; methods are provided by the package +# extension `ext/JumpProcessesStochasticADExt.jl`, which loads only when both +# StochasticAD and Distributions are available. No StochasticAD code lives in +# `src/` -- the differentiable fixed-grid jump simulator is isolated in the ext. +function fixedgrid_simulate end +function fixedgrid_jump_observable end +export fixedgrid_simulate, fixedgrid_jump_observable + """ Aggregator to indicate that individual jumps should also be handled via the leaping algorithm that is passed to solve. diff --git a/test/runtests.jl b/test/runtests.jl index 9a2dfc246..ba327b97d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -36,6 +36,7 @@ end @time @safetestset "Bracketing Tests" begin include("bracketing.jl") end @time @safetestset "Composition-Rejection Table Tests" begin include("table_test.jl") end @time @safetestset "Extinction test" begin include("extinction_test.jl") end + @time @safetestset "StochasticAD Extension Tests" begin include("stochasticad_tests.jl") end end if GROUP == "All" || GROUP == "InterfaceII" @time @safetestset "Saveat Regression test" begin include("saveat_regression.jl") end diff --git a/test/stochasticad_tests.jl b/test/stochasticad_tests.jl new file mode 100644 index 000000000..deb836d69 --- /dev/null +++ b/test/stochasticad_tests.jl @@ -0,0 +1,74 @@ +using JumpProcesses, StochasticAD, Distributions +using Random, Statistics, Test + +# Tests for the optional StochasticAD extension (fixed-grid differentiable jump +# simulator). These need only JumpProcesses + StochasticAD + Distributions -- +# no ODE solver, since the extension never calls `solve`. + +@testset "scalar analytic gradient" begin + # du/dt = -A*u ; one jump of rate p ; on fire u -> u/2 ; observable u(T). + # E[u(T)] = exp(-A*T) * exp(-p*T/2) + # d/dp E[u] = -(T/2) * exp(-A*T) * exp(-p*T/2) + A, T, dt, p0 = 1.0, 2.0, 0.01, 0.5 + analytic = -(T / 2) * exp(-A * T) * exp(-p0 * T / 2) # ≈ -0.08208 + + drift!(du, u, p, t) = (du[1] = -A * u[1]; nothing) + oprob = ODEProblem(drift!, [1.0], (0.0, T), [p0]) + vrj = VariableRateJump((u, p, t) -> p[1], integ -> nothing) # affect! unused + jprob = JumpProblem(oprob, Direct(), vrj; vr_aggregator = VR_Direct()) + + post = (u, p, t) -> [0.5 * u[1]] + obs = u -> u[1] + + N = 2000 + s = Vector{Float64}(undef, N) + for i in 1:N + Random.seed!(i) + s[i] = derivative_estimate(p0) do pk + fixedgrid_jump_observable(jprob, [pk], [post], obs; dt = dt) + end + end + g = mean(s) + @test isapprox(g, analytic; atol = 0.01) +end + +@testset "two-channel MCWF runs and is sensible" begin + # Λ-system MCWF over real components; observable |c3(T)|^2. + # γ2 drives the |2>->|3> channel that feeds the observable, so d/dγ2 > 0. + Ω = [2.0, 2.0] + H = zeros(3, 3); H[1,2]=H[2,1]=Ω[1]; H[2,3]=H[3,2]=Ω[2] + function drift!(du, u, p, t) + γtot = p[1] + p[2] + rc1,ic1,rc2,ic2,rc3,ic3 = u[1],u[2],u[3],u[4],u[5],u[6] + Hr = H*[rc1,rc2,rc3]; Hi = H*[ic1,ic2,ic3] + P2 = rc2^2+ic2^2; g = 0.5*γtot + du[1]= Hi[1]+g*P2*rc1; du[2]=-Hr[1]+g*P2*ic1 + du[3]= Hi[2]-g*(1-P2)*rc2; du[4]=-Hr[2]-g*(1-P2)*ic2 + du[5]= Hi[3]+g*P2*rc3; du[6]=-Hr[3]+g*P2*ic3 + nothing + end + u0 = [1.0,0,0,0,0,0] + γ0 = [0.5, 0.3] + oprob = ODEProblem(drift!, u0, (0.0, 5.0), γ0) + vrj1 = VariableRateJump((u,p,t)->p[1]*(u[3]^2+u[4]^2), integ->nothing) + vrj2 = VariableRateJump((u,p,t)->p[2]*(u[3]^2+u[4]^2), integ->nothing) + jprob = JumpProblem(oprob, Direct(), vrj1, vrj2; vr_aggregator = VR_Direct()) + + posts = [(u,p,t)->[1.0,0,0,0,0,0], (u,p,t)->[0.0,0,0,0,1.0,0]] + obs = u -> u[5]^2 + u[6]^2 + + N = 1500 + K = 2 + S = Matrix{Float64}(undef, K, N) + for i in 1:N, k in 1:K + Random.seed!(1000 + (i-1)*K + k) + S[k, i] = derivative_estimate(γ0[k]) do γk + p = [j == k ? γk : oftype(γk, γ0[j]) for j in 1:K] + fixedgrid_jump_observable(jprob, p, posts, obs; dt = 0.01) + end + end + g = vec(mean(S, dims = 2)) + @test all(isfinite, g) + @test g[2] > 0 # γ2 feeds |3>: positive sensitivity + @test 0.0 < g[1] < g[2] # γ1 smaller positive effect (cf. ~[0.06, 0.11]) +end