Continuous-Time Jump Processes and Gillespie Methods

In this tutorial we will describe how to define and simulate continuous-time jump processes, also known in biological fields as stochastic chemical kinetics (i.e. Gillespie) models. It is not necessary to have read the first tutorial. We will illustrate

  • The different types of jumps that can be represented in DiffEqJump and their use cases.
  • How to speed up pure-jump simulations with only ConstantRateJumps and MassActionJumps by using the SSAStepper time stepper.
  • How to define and use MassActionJumps, a more specialized type of ConstantRateJump that offers improved computational performance.
  • How to use saving controls to reduce memory use per simulation.
  • How to use VariableRateJumps and when they should be preferred over ConstantRateJumps and MassActionJumps.
  • How to create hybrid problems mixing the various jump types with ODEs or SDEs.
  • How to use RegularJumps to enable faster, but approximate, time stepping via τ-leaping methods.

This tutorial assumes you have read the Ordinary Differential Equations tutorial in DifferentialEquations.jl.

We begin by demonstrating how to build jump processes using DiffEqJump.jl's different jump types, which encode the rate functions (i.e. transition rates, intensities, or propensities) and state changes when a given jump occurs.

Note, the SIR model considered here is a type of stochastic chemical kinetics jump process model, and as such the biological modeling functionality of Catalyst.jl can be used to easily specify the model and automatically calculate inputs needed for DiffEqJump's optimized simulation algorithms. We summarize this alternative approach at the beginning for users who may be interested in modeling chemical systems, but note this tutorial is intended to explain the general jump process formulation of DiffEqJump for all users. However, for those users constructing models that can be represented as a collection of chemical reactions we strongly recommend using Catalyst, which should ensure optimal jump types are selected to represent each reaction, and necessary data structures for the simulation algorithms, such as dependency graphs, are automatically calculated.

We'll make use of the DifferentialEquations.jl meta package, which includes DiffEqJump and ODE/SDE solvers, Plots.jl, and (optionally) Catalyst.jl in this tutorial. If not already installed they can be added as follows

using Pkg
Pkg.add("Catalyst)                # optional

Let's now load the required packages and set some default plot settings

using DifferentialEquations, Plots, LinearAlgebra
default(; lw = 2)

Illustrative Model: SIR Disease Dynamics

To illustrate the jump process solvers, we will build an SIR model which matches the tutorial from Gillespie.jl. SIR stands for susceptible, infected, and recovered, and is a model of disease spread. When a susceptible person comes in contact with an infected person, the disease has a chance of infecting the susceptible person. This "chance" is determined by the number of susceptible persons and the number of infected persons, since in larger populations there is a greater chance that two people come into contact. Every infected person will in turn have a rate at which they recover. In our model we'll assume there are no births or deaths, and a recovered individual is protected from reinfection.

We'll begin by giving the mathematical equations for the jump processes of the number of susceptible ($S(t)$), number of infected ($I(t)$), and number of recovered ($R(t)$). In the next section we give a more intuitive and biological description of the model for users that are less familiar with jump processes. Let $Y_i(t)$, $i = 1,2$, denote independent unit Poisson processes. Our basic mathematical model for the evolution of $(S(t),I(t),R(t))$, written using Kurtz's time-change representation, is then

\[\begin{aligned} S(t) &= S(0) - Y_1\left( \int_0^t \beta S(s^{-}) I(s^{-}) \, ds\right) \\ I(t) &= I(0) + Y_1\left( \int_0^t \beta S(s^{-}) I(s^{-}) \, ds\right) - Y_2 \left( \int_0^t \nu I(s^-) \, ds \right) \\ R(t) &= R(0) + Y_2 \left( \int_0^t \nu I(s^-) \, ds \right) \end{aligned}\]

Notice, our model involves two jump processes with rate functions, also known as intensities or propensities, given by $\beta S(t) I(t)$ and $\nu I(t)$ respectively. These give the probability per time a new infected individual is created, and the probability per time some infected individual recovers.

For those less-familiar with the time-change representation, we next give a more intuitive explanation of the model as a collection of chemical reactions, and then demonstrate how these reactions can be written in Catalyst.jl and seamlessly converted into a form that can be used with the DiffEqJump.jl solvers. Users interested in how to directly define jumps using the lower-level DiffEqJump interface can skip to Building and Simulating the Jump Process Using the DiffEqJump Low-Level Interface.

Specifying the SIR Model with Chemical Reactions

The SIR model described above involves two basic chemical reactions,

\[\begin{aligned} S + I &\overset{\beta}{\to} 2 I \\ I &\overset{\nu}{\to} R, \end{aligned}\]

where $\beta$ and $\nu$ are the rate constants of the reactions (with units of probability per time). In a jump process (stochastic chemical kinetics) model, we keep track of the non-negative integer number of each species at each time (i.e. $(S(t), I(t), R(t))$ above). Each reaction has an associated rate function (i.e. intensity or propensity) giving the probability per time it can occur when the system is in state $(S(t),I(t),R(t))$:

\[\begin{matrix} \text{Reaction} & \text{Rate Functions} \\ \hline S + I \overset{\beta}{\to} 2 I & \beta S(t) I(t) \\ I \overset{\nu}{\to} R & \nu I(t). \end{matrix}\]

$\beta$ is determined by factors like the type of the disease. It can be interpreted as the probability per time one pair of susceptible and infected people encounter each other, with the susceptible person becoming sick. The overall rate (i.e. probability per time) that some susceptible person gets sick is then given by the rate constant multiplied by the number of possible pairs of susceptible and infected people. This formulation is known as the law of mass action. Similarly, we have that each individual infected person is assumed to recover with probability per time $\nu$, so that the probability per time some infected person becomes recovered is $\nu$ times the number of infected people, i.e. $\nu I(t)$.

Rate functions give the probability per time for each of the two types of jumps to occur, and hence determine when the state of our system changes. To fully specify our model we also need to specify how the state changes when a jump occurs, giving what are called affect! functions in DiffEqJump. For example, when the $S + I \to 2 I$ reaction occurs and some susceptible person becomes infected, the subsequent (instantaneous) state change is that

\[\begin{aligned} S &\to S - 1 & I &\to I + 1. \end{aligned}\]

Likewise, when the $I \to R$ reaction occurs so that some infected person becomes recovered the state change is

\[\begin{aligned} I &\to I - 1 & R \to R + 1. \end{aligned}\]

In summary, our model is described by two chemical reactions, which each in turn correspond to a jump process determined by a rate function specifying how frequently jumps should occur, and an affect! function for how the state should change when that jump type occurs.

Building and Simulating the Jump Processes from Catalyst Models

Using Catalyst.jl we can input our full reaction network in a form that can be easily used with DiffEqJump's solvers

using Catalyst
sir_model = @reaction_network begin
    β, S + I --> 2I
    ν, I --> R
end β ν

\[ \begin{align*} \require{mhchem} \ce{ S + I &->[$\beta$] 2 I}\\ \ce{ I &->[$\nu$] R} \end{align*} \]

To build a pure jump process model of the reaction system, where the state is constant between jumps, we will use a DiscreteProblem. This encodes that the state only changes at the jump times. We do this by giving the constructor u₀, the initial condition, and tspan, the timespan. Here, we will start with 999 susceptible people, 1 infected person, and 0 recovered people, and solve the problem from t=0.0 to t=250.0. We use the parameters β = 0.1/1000 and ν = 0.01. Thus we build the problem via:

p     = (:β => 0.1/1000, :ν => 0.01)
u₀    = [:S => 999, :I => 10, :R => 0]
tspan = (0.0, 250.0)
prob  = DiscreteProblem(sir_model, u₀, tspan, p)
DiscreteProblem with uType Vector{Int64} and tType Float64. In-place: true
timespan: (0.0, 250.0)
u0: 3-element Vector{Int64}:

Notice, the initial populations are integers since we want the exact number of people in the different states.

The Catalyst reaction network can be converted into various DifferentialEquations.jl problem types, including JumpProblems, ODEProblems, or SDEProblems. To turn it into a JumpProblem representing the SIR jump process model, we simply write

jump_prob = JumpProblem(sir_model, prob, Direct())
JumpProblem with problem DiscreteProblem with aggregator Direct
Number of constant rate jumps: 0
Number of variable rate jumps: 0
Number of mass action jumps: 2

Here Direct() indicates that we will determine the random times and types of reactions using Gillespie's Direct stochastic simulation algorithm (SSA), also known as Doob's method or Kinetic Monte Carlo. See Constant Rate Jump Aggregators for other supported SSAs.

We now have a problem that can be evolved in time using the DiffEqJump solvers. Since our model is a pure jump process (no continuously-varying components), we will use SSAStepper() to handle time-stepping the Direct method from jump to jump:

sol = solve(jump_prob, SSAStepper())
retcode: Default
Interpolation: Piecewise constant interpolation
t: 1869-element Vector{Float64}:
u: 1869-element Vector{Vector{Int64}}:
 [999, 10, 0]
 [998, 11, 0]
 [997, 12, 0]
 [996, 13, 0]
 [995, 14, 0]
 [995, 13, 1]
 [994, 14, 1]
 [994, 13, 2]
 [993, 14, 2]
 [992, 15, 2]
 [0, 148, 861]
 [0, 147, 862]
 [0, 146, 863]
 [0, 145, 864]
 [0, 144, 865]
 [0, 143, 866]
 [0, 142, 867]
 [0, 141, 868]
 [0, 141, 868]

This solve command takes the standard commands of the common interface, and the solution object acts just like any other differential equation solution. Thus there exists a plot recipe, which we can plot with: