Tutorial on using spatial SSAs in DiffEqJump

Vasily Ilin"

This tutorial shows how to use spatial solvers added to DiffEqJump in summer 2021. See the documentation for a tutorial on getting started with DiffEqJump.

Installing DiffEqJump

Once in REPL, do ] add DiffEqJump. After the installation finishes, you will be able to use all the functionality described below.

Reversible binding model on a grid

A 5 by 5 Cartesian grid:

....B
.....
.....
.....
A....

Suppose we have a reversible binding system described by $A+B \to C$ at rate $k_1$ and $C \to A+B$ at rate $k_2$. Further suppose that all $A$ molecules start in the lower left corner, while all $B$ molecules start in the upper right corner of a 5 by 5 grid. There are no $C$ molecules at the start.

We first create the grid:

using DiffEqJump
dims = (5,5)
num_nodes = prod(dims) # number of sites
grid = CartesianGrid(dims) # or use LightGraphs.grid(dims)
A Cartesian grid with dimensions (5, 5)

Now we set the initial state of the simulation. It has to be a matrix with entry $(s,i)$ being the number of species $s$ at site $i$ (with the standard column-major ordering of the grid).

num_species = 3
starting_state = zeros(Int, num_species, num_nodes)
starting_state[1,1] = 25
starting_state[2,end] = 25
starting_state
3×25 Matrix{Int64}:
 25  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0   
0
  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  2
5
  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0   
0

We now set the time-span of the simulation and the reaction rates. These can be chosen arbitrarily.

tspan = (0.0, 3.0)
rates = [6.0, 0.05] # k_1 = rates[1], k_2 = rates[2]
2-element Vector{Float64}:
 6.0
 0.05

Now we can create the DiscreteProblem:

prob = DiscreteProblem(starting_state, tspan, rates)
DiscreteProblem with uType Matrix{Int64} and tType Float64. In-place: true
timespan: (0.0, 3.0)
u0: 3×25 Matrix{Int64}:
 25  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0   
0
  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  2
5
  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0  0   
0

Since both reactions are massaction reactions, we put them together in a MassActionJump. In order to do that we create two stoichiometry vectors. The net stoichiometry vector describes which molecules change in number and how much after each reaction; for example, [1 => -1] is the first molecule disappearing. The reaction stoichiometry vector describes what the reactants of each reaction are; for example, [1 => 1, 2 => 1] would mean that the reactants are one molecule of type 1 and one molecule of type 2.

netstoch = [[1 => -1, 2 => -1, 3 => 1],[1 => 1, 2 => 1, 3 => -1]]
reactstoch = [[1 => 1, 2 => 1],[3 => 1]]
majumps = MassActionJump(rates, reactstoch, netstoch)
DiffEqJump.MassActionJump{Vector{Float64}, Vector{Vector{Pair{Int64, Int64}
}}, Vector{Vector{Pair{Int64, Int64}}}, Nothing}([6.0, 0.05], [[1 => 1, 2 =
> 1], [3 => 1]], [[1 => -1, 2 => -1, 3 => 1], [1 => 1, 2 => 1, 3 => -1]], n
othing)

The last thing to set up is the hopping constants – the probability per time of an individual molecule of each species hopping from one site to another site. In practice this parameter, as well as reaction rates, are obtained empirically. Suppose that molecule $C$ cannot diffuse, while molecules $A$ and $B$ diffuse at probability per time 1 (i.e. the time of the diffusive hop is exponentially distributed with mean 1). Entry $(s,i)$ of hopping_constants is the hopping rate of species $s$ at site $i$ to any of its neighboring sites (diagonal hops are not allowed).

hopping_constants = ones(num_species, num_nodes)
hopping_constants[3, :] .= 0.0
hopping_constants
3×25 Matrix{Float64}:
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0  …  1.0  1.0  1.0  1.0  1.0  1.0  1
.0
 1.0  1.0  1.0  1.0  1.0  1.0  1.0  1.0     1.0  1.0  1.0  1.0  1.0  1.0  1
.0
 0.0  0.0  0.0  0.0  0.0  0.0  0.0  0.0     0.0  0.0  0.0  0.0  0.0  0.0  0
.0

We are now ready to set up the JumpProblem with the Next Subvolume Method.

alg = NSM()
jump_prob = JumpProblem(prob, alg, majumps, hopping_constants=hopping_constants, spatial_system = grid, save_positions=(true, false))
Number of constant rate jumps: 0
Number of variable rate jumps: 0
Have a mass action jump

The save_positions keyword tells the solver to save the positions just before the jumps. To solve the jump problem do

solution = solve(jump_prob, SSAStepper())
retcode: Default
Interpolation: Piecewise constant interpolation
t: 375-element Vector{Float64}:
 0.0
 0.0002479400984079533
 0.007854536056330234
 0.011631890520484809
 0.0147987233042991
 0.0238552775200484
 0.03886671070478702
 0.040525766567334046
 0.042175898036310625
 0.0586862177154805
 ⋮
 2.9457016489837624
 2.9556766640912757
 2.9586474023648273
 2.9612726151313833
 2.968040076511825
 2.9713845083251478
 2.979701571327976
 2.9956787923813795
 3.0
u: 375-element Vector{Matrix{Int64}}:
 [25 0 … 0 0; 0 0 … 0 25; 0 0 … 0 0]
 [24 1 … 0 0; 0 0 … 0 25; 0 0 … 0 0]
 [23 1 … 0 0; 0 0 … 0 25; 0 0 … 0 0]
 [22 2 … 0 0; 0 0 … 0 25; 0 0 … 0 0]
 [22 2 … 0 0; 0 0 … 1 24; 0 0 … 0 0]
 [21 2 … 0 0; 0 0 … 1 24; 0 0 … 0 0]
 [20 3 … 0 0; 0 0 … 1 24; 0 0 … 0 0]
 [20 3 … 0 0; 0 0 … 0 24; 0 0 … 0 0]
 [19 3 … 0 0; 0 0 … 0 24; 0 0 … 0 0]
 [19 3 … 0 0; 0 0 … 0 24; 0 0 … 0 0]
 ⋮
 [4 1 … 0 0; 0 0 … 1 4; 0 0 … 0 0]
 [4 1 … 0 0; 0 0 … 1 4; 0 0 … 0 0]
 [4 1 … 0 0; 0 0 … 1 4; 0 0 … 0 0]
 [4 1 … 0 0; 0 0 … 1 4; 0 0 … 0 0]
 [4 1 … 0 0; 0 0 … 1 4; 0 0 … 0 0]
 [4 1 … 0 0; 0 0 … 0 5; 0 0 … 0 0]
 [4 2 … 0 0; 0 0 … 0 5; 0 0 … 0 0]
 [4 2 … 0 0; 0 0 … 0 5; 0 0 … 0 0]
 [4 2 … 0 0; 0 0 … 0 5; 0 0 … 0 0]

Animation

Visualizing solutions of spatial jump problems is best done with animations.

using Plots
is_static(spec) = (spec == 3) # true if spec does not hop
"get frame k"
function get_frame(k, sol, linear_size, labels, title)
    num_species = length(labels)
    h = 1/linear_size
    t = sol.t[k]
    state = sol.u[k]
    xlim=(0,1+3h/2); ylim=(0,1+3h/2);
    plt = plot(xlim=xlim, ylim=ylim, title = "$title, $(round(t, sigdigits=3)) seconds")

    species_seriess_x = [[] for i in 1:num_species]
    species_seriess_y = [[] for i in 1:num_species]
    CI = CartesianIndices((linear_size, linear_size))
    for ci in CartesianIndices(state)
        species, site = Tuple(ci)
        x,y = Tuple(CI[site])
        num_molecules = state[ci]
        sizehint!(species_seriess_x[species], num_molecules)
        sizehint!(species_seriess_y[species], num_molecules)
        if !is_static(species)
            randsx = rand(num_molecules)
            randsy = rand(num_molecules)
        else
            randsx = zeros(num_molecules)
            randsy = zeros(num_molecules)
        end
        for k in 1:num_molecules
            push!(species_seriess_x[species], x*h - h/4 + 0.5h*randsx[k])
            push!(species_seriess_y[species], y*h - h/4 + 0.5h*randsy[k])
        end
    end
    for species in 1:num_species
        scatter!(plt, species_seriess_x[species], species_seriess_y[species], label = labels[species], marker = 6)
    end
    xticks!(plt, range(xlim...,length = linear_size+1))
    yticks!(plt, range(ylim...,length = linear_size+1))
    xgrid!(plt, 1, 0.7)
    ygrid!(plt, 1, 0.7)
    return plt
end

"make an animation of solution sol in 2 dimensions"
function animate_2d(sol, linear_size; species_labels, title, verbose = true)
    num_frames = length(sol.t)
    anim = @animate for k=1:num_frames
        verbose && println("Making frame $k")
        get_frame(k, sol, linear_size, species_labels, title)
    end
    anim
end
# animate
anim=animate_2d(solution, 5, species_labels = ["A", "B", "C"], title = "A + B <--> C", verbose = false)
fps = 5
gif(anim, fps = fps)