Stochastic resetting in SciML

Introduction

During my stay at IMSc I attended the Frontiers in Non-Equilibrium Physics meeting (my first physics meeting in very many years!), where there were several talks on stochastic resetting. This refers to stochastic processes which at random times “reset” to a given state, for instance to its initial position.

To a physicist, stochastic resetting is of interest because even in the simplest example of a diffusive particle undergoing Brownian motion, the system reaches a non-equilibrium stationary state. From a practical perspective, the study of stochastic resetting is often motivated in the context of search processes (see Section 1 of this review article). In biology, an interesting example of stochastic resetting is given in Section 5.5 of the same review article: consider a molecular interaction where an enzyme \(E\) binds to a substrate \(S\), and when bound triggers the production of a product \(P\), say a protein, that is, the reaction scheme

\[ E+S \rightleftharpoons ES \rightarrow E + P \]

Every time the enzyme (reversibly) unbinds and binds, production of \(P\) is reset and restarted.

In this post, I want to illustrate how easy it is to simulate stochastic processes under resetting in SciML, specifically using the DifferentialEquations and JumpProcesses packages.

The unperturbed stochastic process

Stochastic processes in SciML are defined as stochastic differential equation problems (instances of SDEProblem), if the state space is continuous, or discrete problems (instances of DiscreteProblem), if the state space is discrete. Here’s we’ll look at SDE problems.

A (one-dimensional) SDEProblem is defined by the equation

\[ du = f(u,p,t) dt + g(u,p,t) dW \]

where \(f\) is the drift term, \(g\) the diffusion term, \(W\) the Wiener process, and \(p\) are a set of parameters. Let’s consider the simplest example of Brownian motion in one dimension:

using DifferentialEquations
σ² = 1.0
f(u,p,t) = 0.0
g(u,p,t) = σ²
u0 = 0.0
tspan = (0.0, 10.0)
prob = SDEProblem(f, g, u0, tspan)

SDEProblem with uType Float64 and tType Float64. In-place: false
Non-trivial mass matrix: false
timespan: (0.0, 10.0)
u0: 0.0

We can sample a trajectory of the process using the solve function. Here we use the simple Euler-Maruyama algorithm with a time step of 0.01:

using Plots
sol = solve(prob, EM(), dt=0.01)
plot(sol; linewidth=2, legend=false)

Multiple trajectories can be sampled using a parallel ensemble algorithm.

ensemble_prob = EnsembleProblem(prob)
ensemble_sol = solve(ensemble_prob, EM(), dt=0.01, EnsembleThreads(), trajectories = 5)
plot(ensemble_sol; linewidth=2, legend=false)

The resetting mechanism

To model stochastic resetting, we need to augment the stochastic differential equation with a discrete jump process, that is, create a jump-diffusion problem. As explained in the JumpProcesses documentation, we need two functions: a rate function which determines the distribution of waiting times between jumps (resetting events) and a function that determines what happens to the state of the system at the time of a jump. Let’s model resetting with a constant rate, that is, Poissonian resetting, and resetting the process to its initial position:

r = 2.0
reset_rate(u, p, t) = r
function affect!(integrator)
    integrator.u = u0 # reset to the starting position
    nothing
end

affect! (generic function with 1 method)

Now we can build the resetting mechanism:

reset_jump = ConstantRateJump(reset_rate, affect!)

ConstantRateJump{typeof(reset_rate), typeof(affect!)}(reset_rate, affect!)

The stochastic process under resetting

The unperturbed stochastic process and the resetting mechanism are combined in a JumpProblem:

reset_prob = JumpProblem(prob, Direct(), reset_jump)

JumpProblem with problem SDEProblem with aggregator Direct
Number of jumps with discrete aggregation: 1
Number of jumps with continuous aggregation: 0
Number of mass action jumps: 0

A trajectory can be sampled and plotted using any SDE solver:

reset_sol = solve(reset_prob, EM(), dt=0.01)
plot(reset_sol; linewidth=2, legend=false)

The resetting mechanism concentrates the process around the initial position, which can be seen more clearly using an ensemble simulation:

ensemble_reset_prob = EnsembleProblem(reset_prob)
ensemble_reset_sol = solve(ensemble_reset_prob, EM(), dt=0.01, EnsembleThreads(), trajectories = 5)
plot(ensemble_reset_sol; linewidth=2, legend=false)

The stationary state of diffusion with Poissonian resetting

The example of simple diffusion (Brownian motion) with Poissonian (constant rate) resetting reaches a stationary state which can be computed analytically (see equation (2.18) in the review article):

\[ p^\ast(x) = \lim_{t\to\infty}p(x,t\mid x_0, 0) = \frac{\alpha_0}{2} e^{-\alpha_0\vert x - x_r\vert} \]

where

\[ \alpha_0 = \left(\frac{r}{D}\right)^{1/2}, \]

\(D=\sigma^2/2\) is the diffusion constant, and \(x_r\) is the position to which the system is reset. In other words, \(p^\ast(x)\) is a Laplace distribution with mean \(x_r\) and scale parameter \(\theta=1/\alpha_0\).

We can verify that the simulation reproduces this distribution. First generate an ensemble with more trajectories:

ensemble_reset_sol = solve(ensemble_reset_prob, EM(), dt=0.01, EnsembleThreads(), trajectories = 1000)

EnsembleSolution Solution of length 1000 with uType:
RODESolution{Float64, 1, Vector{Float64}, Nothing, Nothing, Vector{Float64}, NoiseProcess{Float64, 1, Float64, Float64, Nothing, Nothing, typeof(DiffEqNoiseProcess.WHITE_NOISE_DIST), typeof(DiffEqNoiseProcess.WHITE_NOISE_BRIDGE), Nothing, false, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Nothing}, false}, ResettableStacks.ResettableStack{Tuple{Float64, Float64, Nothing}, false}, RSWM{Float64}, Nothing, RandomNumbers.Xorshifts.Xoroshiro128Plus}, Nothing, SDEProblem{Float64, Tuple{Float64, Float64}, false, SciMLBase.NullParameters, Nothing, SDEFunction{false, SciMLBase.FullSpecialize, typeof(f), typeof(g), LinearAlgebra.UniformScaling{Bool}, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, Nothing, typeof(SciMLBase.DEFAULT_OBSERVED), Nothing, Nothing, Nothing}, typeof(g), Base.Pairs{Symbol, Union{}, Tuple{}, @NamedTuple{}}, Nothing}, EM{true}, StochasticDiffEq.LinearInterpolationData{Vector{Float64}, Vector{Float64}}, SciMLBase.DEStats, Nothing, Nothing}

The DifferentialEquations package contains several convenience functions for analyzing an ensemble experiment. Here we only extract the final values of each trajectory:

using Distributions
using StatsPlots
D = σ²/2
α0 = (r/D)^0.5
dstar = Laplace(u0,1/α0)
histogram(EnsembleAnalysis.componentwise_vectors_timepoint(ensemble_reset_sol,tspan[2]), normalize=:pdf, color=:gray, label="Simulation")
plot!(dstar, linewidth=2, color=:red, label="Analytic")

First-passage resetting

There exist many interesting variations on the above resetting scheme. An interesting one is first-passage resetting, where the particle resets not after a certain waiting time, but whenever it crosses a threshold position \(x_0 \pm x^*\). Because resetting is now defined purely by the state of the system and not by a waiting time, it can be simulated directly using event handling functions in the DifferentialEquations package.

In particular we define a function whose zeros correspond to the event triggering, and then define a ContinuousCallback function:

xstar = 1.0
function reset_condition(u, t, integrator)
    (u[1] - u0 - xstar)*(u[1] - u0 + xstar)
end
reset_cb = ContinuousCallback(reset_condition, affect!)

ContinuousCallback{typeof(reset_condition), typeof(affect!), typeof(affect!), typeof(SciMLBase.INITIALIZE_DEFAULT), typeof(SciMLBase.FINALIZE_DEFAULT), Float64, Int64, Rational{Int64}, Nothing, Nothing, Int64}(reset_condition, affect!, affect!, SciMLBase.INITIALIZE_DEFAULT, SciMLBase.FINALIZE_DEFAULT, nothing, SciMLBase.LeftRootFind, 10, Bool[1, 1], 1, 2.220446049250313e-15, 0, 1//100, nothing)

A trajectory can be sampled by solving the original unperturbed SDE with this callback:

reset_sol_fp = solve(prob, EM(), dt=0.01, callback = reset_cb)
plot(reset_sol_fp; linewidth=2, legend=false)

The same is true for sampling an ensemble of trajectories:

ensemble_reset_sol_fp = solve(ensemble_prob, EM(), dt=0.01, callback = reset_cb, EnsembleThreads(), trajectories = 1000)
histogram(EnsembleAnalysis.componentwise_vectors_timepoint(ensemble_reset_sol_fp,tspan[2]), normalize=:pdf, color=:gray, label="")