Add optional StochasticAD extension for differentiating jump processes

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"]

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

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.

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

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