GPU-Accelerated Stochastic Partial Differential Equations

Let's solve stochastic PDEs in Julia using GPU parallelism. To do this, we will use the type-genericness of the DifferentialEquations.jl library in order to write a code that uses within-method GPU-parallelism on the system of PDEs. The OrdinaryDiffEq.jl solvers of DifferentialEquations.jl, including implicit solvers with GMRES, etc., and the same for SDEs, DAEs, DDEs, etc. are all GPU-compatible with a fast form of broadcast.

Note

The non-native Julia solvers, like Sundials are incompatible with arbitrary input types and thus not compatible with GPUs.

Let's dive into showing how to accelerate ODE solving with GPUs!

Before we start: the two ways to accelerate ODE solvers with GPUs

Before we dive deeper, let us remark that there are two very different ways that one can accelerate an ODE solution with GPUs. There is one case where u is very big and f is very expensive, but very structured, and you use GPUs to accelerate the computation of said f. The other use case is where u is very small but you want to solve the ODE f over many different initial conditions (u0) or parameters p. In that case, you can use GPUs to parallelize over different parameters and initial conditions. In other words:

Type of Problem	SciML Solution
Accelerate a big ODE	Use CUDA.jl's CuArray as `u0`
Solve the same ODE with many `u0` and `p`	Use DiffEqGPU.jl's`EnsembleGPUArray` and `EnsembleGPUKernel`

This showcase will focus on the former case. For the latter, see the massively parallel GPU ODE solving showcase.

Our Problem: 2-dimensional Reaction-Diffusion Equations

The reaction-diffusion equation is a PDE commonly handled in systems biology, which is a diffusion equation plus a nonlinear reaction term. The dynamics are defined as:

\[u_t = D \Delta u + f(t,u)\]

But this doesn't need to only have a single “reactant” u: this can be a vector of reactants and the $f$ is then the nonlinear vector equations describing how these different pieces react together. Let's settle on a specific equation to make this easier to explain. Let's use a simple model of a 3-component system where A can diffuse through space to bind with the non-diffusive B to form the complex C (also non-diffusive, assume B is too big and gets stuck in a cell which causes C=A+B to be stuck as well). Other than the binding, we make each of these undergo a simple birth-death process, and we write down the equations which result from mass-action kinetics. If this all is meaningless to you, just understand that it gives the system of PDEs:

\[\begin{align} A_t &= D \Delta A + \alpha_A(x) - \beta_A A - r_1 A B + r_2 C\\ B_t &= \alpha_B - \beta_B B - r_1 A B + r_2 C\\ C_t &= \alpha_C - \beta_C C + r_1 A B - r_2 C \end{align}\]

One addition that was made to the model is that we let $\alpha_A(x)$ be the production of $A$, and we let that be a function of space so that way it only is produced on one side of our equation. Let's make it a constant when x>80, and 0 otherwise, and let our spatial domain be $x \in [0,100]$ and $y \in [0,100]$.

This model is spatial: each reactant $u(t,x,y)$ is defined at each point in space, and all of the reactions are local, meaning that $f$ at spatial point $(x,y)$ only uses $u_i(t,x,y)$. This is an important fact which will come up later for parallelization.

Discretizing the PDE into ODEs

In order to solve this via a method of lines (MOL) approach, we need to discretize the PDE into a system of ODEs. Let's do a simple uniformly-spaced grid finite difference discretization. Choose $dx = 1$ and $dy = 1$ so that we have 100*100=10000 points for each reactant. Notice how fast that grows! Put the reactants in a matrix such that A[i,j] = A(x_j,y_i), i.e. the columns of the matrix are the $x$ values and the rows are the $y$ values (this way looking at the matrix is essentially like looking at the discretized space).

So now we have 3 matrices (A, B, and C) for our reactants. How do we discretize the PDE? In this case, the diffusion term simply becomes a tridiagonal matrix $M$ where $[1,-2,1]$ is the central band. You can notice that $MA$ performs diffusion along the columns of $A$, and so this is diffusion along the $y$. Similarly, $AM$ flips the indices and thus does diffusion along the rows of $A$ making this diffusion along $x$. Thus $D(M_yA + AM_x)$ is the discretized Laplacian (we could have separate diffusion constants and $dx \neq dy$ if we want by using different constants on the $M$, but let's not do that for this simple example. We leave that as an exercise for the reader). We enforced a Neumann boundary condition with zero derivative (also known as a no-flux boundary condition) by reflecting the changes over the boundary. Thus the derivative operator is generated as:

using LinearAlgebra

# Define the constants for the PDE
const α₂ = 1.0
const α₃ = 1.0
const β₁ = 1.0
const β₂ = 1.0
const β₃ = 1.0
const r₁ = 1.0
const r₂ = 1.0
const D = 100.0
const γ₁ = 0.1
const γ₂ = 0.1
const γ₃ = 0.1
const N = 100
const X = reshape([i for i in 1:100 for j in 1:100], N, N)
const Y = reshape([j for i in 1:100 for j in 1:100], N, N)
const α₁ = 1.0 .* (X .>= 80)

const Mx = Tridiagonal([1.0 for i in 1:(N - 1)], [-2.0 for i in 1:N],
                       [1.0 for i in 1:(N - 1)])
const My = copy(Mx)
# Do the reflections, different for x and y operators
Mx[2, 1] = 2.0
Mx[end - 1, end] = 2.0
My[1, 2] = 2.0
My[end, end - 1] = 2.0

2.0

Note

We could have also done these discretization steps using DiffEqOperators.jl or MethodOfLines.jl. However, we are going to keep it in this form, so we can show the full code, making it easier to see how to define GPU-ready code!

Since all of the reactions are local, we only have each point in space react separately. Thus this represents itself as element-wise equations on the reactants. Thus we can write it out quite simply. The ODE which then represents the PDE is thus in pseudo Julia code:

DA = D * (Mx * A + A * My)
@. DA + α₁ - β₁ * A - r₁ * A * B + r₂ * C
@. α₂ - β₂ * B - r₁ * A * B + r₂ * C
@. α₃ - β₃ * C + r₁ * A * B - r₂ * C

Note here that I am using α₁ as a matrix (or row-vector, since that will broadcast just fine) where every point in space with x<80 has this zero, and all of the others have it as a constant. The other coefficients are all scalars.

How do we do this with the ODE solver?

Our Representation via Views of 3-Tensors

We can represent our problem with a 3-dimensional tensor, taking each 2-dimensional slice as our (A,B,C). This means that we can define:

u0 = zeros(N, N, 3);

100×100×3 Array{Float64, 3}:
[:, :, 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.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.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.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.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.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.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.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.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.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.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.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.0  0.0  0.0  0.0  0.0  0.0

[:, :, 2] =
 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  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  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  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  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
 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  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  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  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  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
 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  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

[:, :, 3] =
 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  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  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  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  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
 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  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  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  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  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
 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  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

Now we can decompose it like:

A = @view u[:, :, 1]
B = @view u[:, :, 2]
C = @view u[:, :, 3]
dA = @view du[:, :, 1]
dB = @view du[:, :, 2]
dC = @view du[:, :, 3]

These views will not construct new arrays and will instead just be pointers to the (contiguous) memory pieces, so this is a nice and efficient way to handle this. Together, our ODE using this tensor as its container can be written as follows:

function f(du, u, p, t)
    A = @view u[:, :, 1]
    B = @view u[:, :, 2]
    C = @view u[:, :, 3]
    dA = @view du[:, :, 1]
    dB = @view du[:, :, 2]
    dC = @view du[:, :, 3]
    DA = D * (Mx * A + A * My)
    @. dA = DA + α₁ - β₁ * A - r₁ * A * B + r₂ * C
    @. dB = α₂ - β₂ * B - r₁ * A * B + r₂ * C
    @. dC = α₃ - β₃ * C + r₁ * A * B - r₂ * C
end

f (generic function with 1 method)

where this is using @. to do inplace updates on our du to say how the full tensor should update in time. Note that we can make this more efficient by adding some cache variables to the diffusion matrix multiplications and using mul!, but let's ignore that for now.

Together, the ODE which defines our PDE is thus:

using DifferentialEquations

prob = ODEProblem(f, u0, (0.0, 100.0))
@time sol = solve(prob, ROCK2());

retcode: Success
Interpolation: 3rd order Hermite
t: 125-element Vector{Float64}:
   0.0
   2.796123741859731e-5
   6.437859112070983e-5
   0.0001355696558501728
   0.00023376199838581863
   0.0003776022263950997
   0.000565057887749957
   0.0008071655660257587
   0.0011053533031798865
   0.001466685803229837
   ⋮
   7.947518012781672
   9.216398069047099
  11.048709859457816
  13.932713724920005
  18.815455951532602
  27.10742394873848
  43.60845467747508
  74.75755579382653
 100.0
u: 125-element Vector{Array{Float64, 3}}:
 [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; … ; 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]
 [3.906344697357476e-10 3.909112473032509e-10 … 2.796119589268521e-5 2.7883165317332937e-5; 3.909112473032509e-10 3.911880248707541e-10 … 2.8039337179064475e-5 2.796119589268521e-5; … ; 3.909112473032509e-10 3.911880248707541e-10 … 2.8039337179064475e-5 2.796119589268521e-5; 3.906344697357476e-10 3.909112473032509e-10 … 2.796119589268521e-5 2.7883165317332937e-5;;; 2.7961237415829508e-5 2.7961237415829508e-5 … 2.7961237409320334e-5 2.79612374093249e-5; 2.7961237415829508e-5 2.7961237415829508e-5 … 2.796123740931577e-5 2.7961237409320334e-5; … ; 2.7961237415829508e-5 2.7961237415829508e-5 … 2.796123740931577e-5 2.7961237409320334e-5; 2.7961237415829508e-5 2.7961237415829508e-5 … 2.7961237409320334e-5 2.79612374093249e-5;;; 2.7960455593334905e-5 2.7960455593334905e-5 … 2.796045559984408e-5 2.7960455599839515e-5; 2.7960455593334905e-5 2.7960455593334905e-5 … 2.7960455599848643e-5 2.796045559984408e-5; … ; 2.7960455593334905e-5 2.7960455593334905e-5 … 2.7960455599848643e-5 2.796045559984408e-5; 2.7960455593334905e-5 2.7960455593334905e-5 … 2.796045559984408e-5 2.7960455599839515e-5]
 [2.064777915080023e-9 2.07217804373864e-9 … 6.437748464989885e-5 6.396710057888366e-5; 2.07217804373864e-9 2.0796096029871125e-9 … 6.47908131237352e-5 6.437748464989885e-5; … ; 2.07217804373864e-9 2.0796096029871125e-9 … 6.47908131237352e-5 6.437748464989885e-5; 2.064777915080023e-9 2.07217804373864e-9 … 6.437748464989885e-5 6.396710057888366e-5;;; 6.437859104627274e-5 6.437859104627274e-5 … 6.437859095983186e-5 6.437859096018008e-5; 6.437859104627274e-5 6.437859104627273e-5 … 6.437859095948209e-5 6.437859095983186e-5; … ; 6.437859104627274e-5 6.437859104627273e-5 … 6.437859095948209e-5 6.437859095983186e-5; 6.437859104627274e-5 6.437859104627274e-5 … 6.437859095983186e-5 6.437859096018008e-5;;; 6.437444666658912e-5 6.437444666658914e-5 … 6.437444675303001e-5 6.43744467526818e-5; 6.437444666658914e-5 6.437444666658914e-5 … 6.437444675337979e-5 6.437444675303001e-5; … ; 6.437444666658914e-5 6.437444666658914e-5 … 6.437444675337979e-5 6.437444675303001e-5; 6.437444666658912e-5 6.437444666658914e-5 … 6.437444675303001e-5 6.43744467526818e-5]
 [9.115090300876016e-9 9.18816509527637e-9 … 0.00013555878362043303 0.000133761104469766; 9.18816509527637e-9 9.262006848947153e-9 … 0.00013738531104727627 0.00013555878362043303; … ; 9.18816509527637e-9 9.262006848947153e-9 … 0.00013738531104727627 0.00013555878362043303; 9.115090300876016e-9 9.18816509527637e-9 … 0.00013555878362043303 0.000133761104469766;;; 0.00013556965510873605 0.00013556965510873586 … 0.00013556965429362562 0.00013556965430078604; 0.00013556965510873586 0.00013556965510873567 … 0.0001355696542863866 0.00013556965429362562; … ; 0.00013556965510873586 0.00013556965510873567 … 0.0001355696542863866 0.00013556965429362562; 0.00013556965510873605 0.00013556965510873586 … 0.00013556965429362562 0.00013556965430078604;;; 0.0001355512782014618 0.000135551278201462 … 0.00013555127901657225 0.00013555127900941182; 0.000135551278201462 0.0001355512782014622 … 0.00013555127902381128 0.00013555127901657225; … ; 0.000135551278201462 0.0001355512782014622 … 0.00013555127902381128 0.00013555127901657225; 0.0001355512782014618 0.000135551278201462 … 0.00013555127901657225 0.00013555127900941182]
 [2.6926897581315367e-8 2.7313507714770143e-8 … 0.00023370489649268138 0.00022845315379793862; 2.7313507714770143e-8 2.7707718844309493e-8 … 0.0002391075039605557 0.00023370489649268138; … ; 2.7313507714770143e-8 2.7707718844309493e-8 … 0.0002391075039605557 0.00023370489649268138; 2.6926897581315367e-8 2.7313507714770143e-8 … 0.00023370489649268138 0.00022845315379793862;;; 0.00023376199441293072 0.00023376199441292722 … 0.0002337619902052695 0.00023376199027161357; 0.00023376199441292722 0.00023376199441292367 … 0.00023376199013754758 0.0002337619902052695; … ; 0.00023376199441292722 0.00023376199441292367 … 0.00023376199013754758 0.0002337619902052695; 0.00023376199441293072 0.00023376199441292722 … 0.0002337619902052695 0.00023376199027161357;;; 0.00023370736165976154 0.0002337073616597651 … 0.00023370736586742275 0.0002337073658010787; 0.0002337073616597651 0.0002337073616597686 … 0.0002337073659351447 0.00023370736586742275; … ; 0.0002337073616597651 0.0002337073616597686 … 0.0002337073659351447 0.00023370736586742275; 0.00023370736165976154 0.0002337073616597651 … 0.00023370736586742275 0.0002337073658010787]
 [6.961732292649316e-8 7.124591837656885e-8 … 0.00037736420948259527 0.0003640011486945887; 7.124591837656885e-8 7.292787278493182e-8 … 0.0003913522086821116 0.00037736420948259527; … ; 7.124591837656885e-8 7.292787278493182e-8 … 0.0003913522086821116 0.00037736420948259527; 6.961732292649316e-8 7.124591837656885e-8 … 0.00037736420948259527 0.0003640011486945887;;; 0.0003776022093504568 0.0003776022093504163 … 0.00037760219156997705 0.00037760219202269337; 0.0003776022093504163 0.0003776022093503748 … 0.000377602191101566 0.00037760219156997705; … ; 0.0003776022093504163 0.0003776022093503748 … 0.000377602191101566 0.00037760219156997705; 0.0003776022093504568 0.0003776022093504163 … 0.00037760219156997705 0.00037760219202269337;;; 0.00037745967704351697 0.00037745967704355747 … 0.00037745969482399675 0.0003774596943712805; 0.00037745967704355747 0.00037745967704359894 … 0.00037745969529240787 0.00037745969482399675; … ; 0.00037745967704355747 0.00037745967704359894 … 0.00037745969529240787 0.00037745969482399675; 0.00037745967704351697 0.00037745967704355747 … 0.00037745969482399675 0.0003774596943712805]
 [1.540579894733325e-7 1.5945566561169678e-7 … 0.0005642799709081393 0.0005353147081300882; 1.5945566561169678e-7 1.6512254090905263e-7 … 0.0005952708190357286 0.0005642799709081393; … ; 1.5945566561169678e-7 1.6512254090905263e-7 … 0.0005952708190357286 0.0005642799709081393; 1.540579894733325e-7 1.5945566561169678e-7 … 0.0005642799709081393 0.0005353147081300882;;; 0.0005650578298830339 0.0005650578298827208 … 0.000565057770205711 0.0005650577724579198; 0.0005650578298827208 0.0005650578298823957 … 0.0005650577678343024 0.000565057770205711; … ; 0.0005650578298827208 0.0005650578298823957 … 0.0005650577678343024 0.000565057770205711; 0.0005650578298830339 0.0005650578298827208 … 0.000565057770205711 0.0005650577724579198;;; 0.0005647387130703668 0.0005647387130706799 … 0.0005647387727476897 0.0005647387704954811; 0.0005647387130706799 0.0005647387130710051 … 0.0005647387751190984 0.0005647387727476897; … ; 0.0005647387130706799 0.0005647387130710051 … 0.0005647387751190984 0.0005647387727476897; 0.0005647387130703668 0.0005647387130706799 … 0.0005647387727476897 0.0005647387704954811]
 [3.0969738789168484e-7 3.2507995129791364e-7 … 0.0008049880902276838 0.0007482706516369002; 3.2507995129791364e-7 3.415619008018545e-7 … 0.0008673145106382796 0.0008049880902276838; … ; 3.2507995129791364e-7 3.415619008018545e-7 … 0.0008673145106382796 0.0008049880902276838; 3.0969738789168484e-7 3.2507995129791364e-7 … 0.0008049880902276838 0.0007482706516369002;;; 0.0008071653959737827 0.0008071653959719099 … 0.0008071652219566872 0.0008071652311307488; 0.0008071653959719099 0.0008071653959699315 … 0.0008071652120858652 0.0008071652219566872; … ; 0.0008071653959719099 0.0008071653959699315 … 0.0008071652120858652 0.0008071652219566872; 0.0008071653959737827 0.0008071653959719099 … 0.0008071652219566872 0.0008071652311307488;;; 0.0008065143898932863 0.0008065143898951592 … 0.0008065145639103818 0.0008065145547363204; 0.0008065143898951592 0.0008065143898971374 … 0.0008065145737812039 0.0008065145639103818; … ; 0.0008065143898951592 0.0008065143898971374 … 0.0008065145737812039 0.0008065145639103818; 0.0008065143898932863 0.0008065143898951592 … 0.0008065145639103818 0.0008065145547363204]
 [5.705076773905195e-7 6.087951811632106e-7 … 0.0011000514541360741 0.0009988289301903872; 6.087951811632106e-7 6.508085123006421e-7 … 0.0012147497953614846 0.0011000514541360741; … ; 6.087951811632106e-7 6.508085123006421e-7 … 0.0012147497953614846 0.0011000514541360741; 5.705076773905195e-7 6.087951811632106e-7 … 0.0011000514541360741 0.0009988289301903872;;; 0.0011053528638517067 0.001105352863842754 … 0.0011053524171780035 0.0011053524484369427; 0.001105352863842754 0.0011053528638331044 … 0.0011053523826827932 0.0011053524171780035; … ; 0.001105352863842754 0.0011053528638331044 … 0.0011053523826827932 0.0011053524171780035; 0.0011053528638517067 0.001105352863842754 … 0.0011053524171780035 0.0011053524484369427;;; 0.001104132375968241 0.0011041323759771937 … 0.0011041328226419442 0.0011041327913830052; 0.0011041323759771937 0.0011041323759868433 … 0.0011041328571371547 0.0011041328226419442; … ; 0.0011041323759771937 0.0011041323759868433 … 0.0011041328571371547 0.0011041328226419442; 0.001104132375968241 0.0011041323759771937 … 0.0011041328226419442 0.0011041327913830052]
 [9.837734761082027e-7 1.069728314870665e-6 … 0.0014550919403635942 0.001286939535979884; 1.069728314870665e-6 1.1666270528831613e-6 … 0.0016522401555657791 0.0014550919403635942; … ; 1.069728314870665e-6 1.1666270528831613e-6 … 0.0016522401555657791 0.0014550919403635942; 9.837734761082027e-7 1.069728314870665e-6 … 0.0014550919403635942 0.001286939535979884;;; 0.0014666847725153 0.0014666847724792008 … 0.0014666837305651816 0.0014666838235512587; 0.0014666847724792008 0.0014666847724393811 … 0.0014666836249749018 0.0014666837305651816; … ; 0.0014666847724792008 0.0014666847724393811 … 0.0014666836249749018 0.0014666837305651816; 0.0014666847725153 0.0014666847724792008 … 0.0014666837305651816 0.0014666838235512587;;; 0.001464536697606192 0.0014645366976422912 … 0.0014645377395563106 0.0014645376465702335; 0.0014645366976422912 0.001464536697682111 … 0.0014645378451465904 0.0014645377395563106; … ; 0.0014645366976422912 0.001464536697682111 … 0.0014645378451465904 0.0014645377395563106; 0.001464536697606192 0.0014645366976422912 … 0.0014645377395563106 0.0014645376465702335]
 ⋮
 [0.0860980584145309 0.17043647381136437 … 0.5253487807847173 0.2653793414280966; 0.17043632607368323 0.3382712254382821 … 1.0431153649867155 0.5255803265164318; … ; 0.17043632607047352 0.338271225445665 … 1.0431153649877962 0.5255803265126668; 0.08609805841464986 0.1704364738091774 … 0.5253487807877247 0.265379341430695;;; 1.4373252284252627 1.381514056133696 … 1.1874256906452123 1.3236301768644763; 1.3815141490947984 1.282410462486958 … 0.9854451339241748 1.1873168599586519; … ; 1.3815141490947924 1.2824104624869688 … 0.9854451339241813 1.1873168599586483; 1.4373252284252602 1.3815140561336967 … 1.1874256906452232 1.3236301768644851;;; 0.5618219481395736 0.6176331204311375 … 0.8117214859196261 0.6755169997003598; 0.6176330274700377 0.7167367140778765 … 1.0137020426406604 0.811830316606183; … ; 0.6176330274700433 0.7167367140778665 … 1.0137020426406564 0.811830316606187; 0.5618219481395745 0.6176331204311389 … 0.8117214859196145 0.6755169997003496]
 [0.08611716284361551 0.17047468757767417 … 0.5254093393090755 0.265409623784774; 0.17047453796637485 0.33834765744416206 … 1.043236489945349 0.5256408960059097; … ; 0.17047453797624784 0.3383476574329872 … 1.0432364899545734 0.5256408960113826; 0.08611716284187021 0.17047468757644577 … 0.5254093392900094 0.26540962378754035;;; 1.4378006045760103 1.3819336252553542 … 1.1877342572082403 1.3240275973003055; 1.3819337199879664 1.2827436621439456 … 0.9856521249815257 1.1876253748112229; … ; 1.3819337199879878 1.282743662143918 … 0.9856521249815371 1.187625374811248; 1.4378006045760054 1.3819336252553445 … 1.1877342572081835 1.3240275973003033;;; 0.5618903141179099 0.6177572934385639 … 0.8119566614856801 0.6756633213936142; 0.6177571987059528 0.716947256549974 … 1.014038793712393 0.8120655438826948; … ; 0.6177571987059314 0.7169472565500006 … 1.014038793712381 0.812065543882672; 0.5618903141179137 0.6177572934385737 … 0.8119566614857363 0.675663321393615]
 [0.08612424653409514 0.1704888541714203 … 0.5254313531199226 0.26542063415503336; 0.17048870388305637 0.3383759883686652 … 1.0432805111692454 0.5256629133739961; … ; 0.1704887038883321 0.33837598835841554 … 1.0432805111290577 0.5256629134018219; 0.08612424653691521 0.17048885417575446 … 0.5254313531223933 0.2654206341449212;;; 1.4379761857636695 1.3820885642569618 … 1.187848575254823 1.3241746260446108; 1.3820886597402704 1.2828666720897481 … 0.9857290285709916 1.1877396740993025; … ; 1.382088659740295 1.2828666720897104 … 0.985729028570942 1.1877396740993493; 1.4379761857636821 1.3820885642569747 … 1.1878485752547856 1.3241746260445781;;; 0.5619155910916934 0.6178032125983992 … 0.81204320160054 0.675717150810749; 0.6178031171150915 0.7170251047656166 … 1.014162748284371 0.8121521027560599; … ; 0.6178031171150687 0.7170251047656533 … 1.01416274828442 0.8121521027560125; 0.5619155910916772 0.6178032125983878 … 0.8120432016005785 0.6757171508107845]
 [0.08612614504602102 0.17049265171457642 … 0.525437120686977 0.26542351832198846; 0.1704925012778496 0.33838358347625414 … 1.0432920459311534 0.5256686816812286; … ; 0.17049250125561458 0.33838358352274106 … 1.0432920462707862 0.5256686816281663; 0.08612614502946425 0.17049265174051192 … 0.5254371206311342 0.26542351827734434;;; 1.4380217482813722 1.3821288252436692 … 1.1878784162080114 1.3242128480702946; 1.382128920937962 1.282898770102663 … 0.9857493588365255 1.187769510236514; … ; 1.3821289209378955 1.2828987701027577 … 0.9857493588371241 1.1877695102364358; 1.4380217482813036 1.382128825243725 … 1.1878784162079201 1.3242128480701114;;; 0.5619220947768395 0.6178150178145437 … 0.81206542685021 0.6757309949879174; 0.6178149221202488 0.7170450729555525 … 1.0141944842216821 0.8121743328217007; … ; 0.6178149221203245 0.7170450729554643 … 1.0141944842210908 0.8121743328217819; 0.5619220947769019 0.6178150178144902 … 0.8120654268503005 0.6757309949881005]
 [0.0861263484550237 0.1704930588114406 … 0.5254376433059956 0.26542377925479826; 0.1704929083575903 0.3383843985532119 … 1.043293094769307 0.5256692044097837; … ; 0.1704929083179532 0.3383843985504839 … 1.0432930941125902 0.5256692047684617; 0.08612634841214753 0.17049305897868994 … 0.5254376437522809 0.2654237789478525;;; 1.4380264806957288 1.3821330324293595 … 1.1878815936547746 1.324216899594982; 1.3821331281547102 1.2829020891792937 … 0.9857513300943188 1.1877726872009684; … ; 1.3821331281546714 1.2829020891793919 … 0.9857513300935629 1.1877726872015324; 1.4380264806956726 1.382133032429838 … 1.1878815936554665 1.3242168995943688;;; 0.5619227864032379 0.6178162346696093 … 0.812067673444226 0.6757323675039573; 0.6178161389442847 0.7170471779196865 … 1.0141979370046275 0.8121765798980022; … ; 0.6178161389443182 0.7170471779195954 … 1.0141979370054217 0.8121765798974202; 0.561922786403297 0.61781623466914 … 0.8120676734434983 0.6757323675045981]
 [0.08612696916103961 0.1704942996211746 … 0.5254400816807748 0.26542500060545304; 0.17049414903877258 0.3383868795724996 … 1.043297975800099 0.5256716441583383; … ; 0.17049414919857298 0.33838687869515427 … 1.0432979721279094 0.5256716450098741; 0.08612696914432466 0.17049429987008752 … 0.525440082168004 0.26542500091475424;;; 1.4380418735129417 1.3821466102175435 … 1.1878911197329391 1.324229463323031; 1.3821467059253385 1.2829128644483128 … 0.9857575482619754 1.1877822109612353; … ; 1.3821467059252304 1.282912864446957 … 0.9857575482563989 1.1877822109616085; 1.4380418735133622 1.3821466102178275 … 1.1878911197332265 1.3242294633241454;;; 0.5619250181562303 0.6178202814514968 … 0.8120757719362344 0.675737428346068; 0.6178201857436643 0.7170540272208108 … 1.014209343407195 0.8121846807079487; … ; 0.6178201857438892 0.7170540272223528 … 1.014209343412782 0.8121846807075852; 0.5619250181559755 0.6178202814512223 … 0.8120757719358355 0.6757374283449257]
 [0.08612922377484898 0.17049880746230267 … 0.5254468148723012 0.26542836358957134; 0.17049865534065103 0.33839590092128835 … 1.0433114145796933 0.5256783805383086; … ; 0.17049865860770871 0.3383958958161563 … 1.0433114239576757 0.5256783777167134; 0.08612922322293583 0.17049880814212776 … 0.525446811840995 0.2654283642611259;;; 1.4380984980593468 1.382196601013655 … 1.1879282918393537 1.3242770694216377; 1.382196697047234 1.2829525933107075 … 0.9857827134840025 1.1878193774417978; … ; 1.3821966970533412 1.2829525932993169 … 0.9857827134916414 1.1878193774431431; 1.43809849805989 1.3821966010112683 … 1.1879282918325718 1.3242770694231056;;; 0.5619331295482851 0.6178350265953393 … 0.8121033357688703 0.6757545581872626; 0.6178349305618676 0.7170790342975233 … 1.0142489141247812 0.8122122501663108; … ; 0.6178349305548604 0.7170790343096268 … 1.014248914116363 0.812212250165756; 0.5619331295484676 0.6178350265969836 … 0.8121033357761148 0.675754558186087]
 [0.08612916330502197 0.17049869317371466 … 0.5254465304448525 0.2654282163618434; 0.17049854499628475 0.33839566059686466 … 1.0433108429459548 0.5256780726382269; … ; 0.17049854184749363 0.3383956617550718 … 1.0433108543112954 0.5256780602391141; 0.08612916171450978 0.17049869578421842 … 0.5254465182386399 0.2654282327286097;;; 1.438096899651587 1.3821951874950726 … 1.1879273352067259 1.324275791897684; 1.3821952836875342 1.2829514631453598 … 0.985782108932073 1.1878184210379235; … ; 1.3821952836664313 1.2829514631602468 … 0.9857821089456186 1.187818421030176; 1.4380968996444772 1.3821951874942873 … 1.187927335176687 1.3242757919422374;;; 0.5619328997439027 0.6178346119017598 … 0.8121024641884341 0.6757540074995592; 0.6178345157098714 0.7170783362480253 … 1.0142476904650342 0.8122113783598383; … ; 0.6178345157287075 0.7170783362385516 … 1.0142476904497297 0.8122113783638633; 0.5619328997500804 0.6178346119013959 … 0.8121024642139002 0.6757540074487606]
 [0.08612858502693178 0.17049753011854185 … 0.5254450995508656 0.2654275119357499; 0.17049737654539124 0.33839335967131645 … 1.0433080246587925 0.5256766560874475; … ; 0.17049739061254476 0.33839332923662874 … 1.0433080297881574 0.5256766653154703; 0.08612858461742168 0.17049753677333118 … 0.5254450841778983 0.2654275090734674;;; 1.4380834767959698 1.382183387135683 … 1.1879183611000284 1.3242643710114153; 1.3821834831479929 1.2829421711018256 … 0.9857759838047694 1.1878094482253942; … ; 1.3821834831841513 1.282942171049491 … 0.9857759838212017 1.1878094482498152; 1.4380834768058741 1.3821833871480351 … 1.18791836103523 1.324264370990561;;; 0.5619309119379512 0.6178310015967424 … 0.8120960276328201 0.6757500177235799; 0.617830905584656 0.7170722176292615 … 1.014238404929981 0.8122049405078227; … ; 0.6178309055491185 0.7170722176846236 … 1.0142384049122126 0.8122049404803533; 0.5619309119268662 0.6178310015838927 … 0.8120960276957662 0.675750017738761]

@time sol = solve(prob, ROCK2());

retcode: Success
Interpolation: 3rd order Hermite
t: 125-element Vector{Float64}:
   0.0
   2.796123741859731e-5
   6.437859112070983e-5
   0.0001355696558501728
   0.00023376199838581863
   0.0003776022263950997
   0.000565057887749957
   0.0008071655660257587
   0.0011053533031798865
   0.001466685803229837
   ⋮
   7.947518012781672
   9.216398069047099
  11.048709859457816
  13.932713724920005
  18.815455951532602
  27.10742394873848
  43.60845467747508
  74.75755579382653
 100.0
u: 125-element Vector{Array{Float64, 3}}:
 [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; … ; 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]
 [3.906344697357476e-10 3.909112473032509e-10 … 2.796119589268521e-5 2.7883165317332937e-5; 3.909112473032509e-10 3.911880248707541e-10 … 2.8039337179064475e-5 2.796119589268521e-5; … ; 3.909112473032509e-10 3.911880248707541e-10 … 2.8039337179064475e-5 2.796119589268521e-5; 3.906344697357476e-10 3.909112473032509e-10 … 2.796119589268521e-5 2.7883165317332937e-5;;; 2.7961237415829508e-5 2.7961237415829508e-5 … 2.7961237409320334e-5 2.79612374093249e-5; 2.7961237415829508e-5 2.7961237415829508e-5 … 2.796123740931577e-5 2.7961237409320334e-5; … ; 2.7961237415829508e-5 2.7961237415829508e-5 … 2.796123740931577e-5 2.7961237409320334e-5; 2.7961237415829508e-5 2.7961237415829508e-5 … 2.7961237409320334e-5 2.79612374093249e-5;;; 2.7960455593334905e-5 2.7960455593334905e-5 … 2.796045559984408e-5 2.7960455599839515e-5; 2.7960455593334905e-5 2.7960455593334905e-5 … 2.7960455599848643e-5 2.796045559984408e-5; … ; 2.7960455593334905e-5 2.7960455593334905e-5 … 2.7960455599848643e-5 2.796045559984408e-5; 2.7960455593334905e-5 2.7960455593334905e-5 … 2.796045559984408e-5 2.7960455599839515e-5]
 [2.064777915080023e-9 2.07217804373864e-9 … 6.437748464989885e-5 6.396710057888366e-5; 2.07217804373864e-9 2.0796096029871125e-9 … 6.47908131237352e-5 6.437748464989885e-5; … ; 2.07217804373864e-9 2.0796096029871125e-9 … 6.47908131237352e-5 6.437748464989885e-5; 2.064777915080023e-9 2.07217804373864e-9 … 6.437748464989885e-5 6.396710057888366e-5;;; 6.437859104627274e-5 6.437859104627274e-5 … 6.437859095983186e-5 6.437859096018008e-5; 6.437859104627274e-5 6.437859104627273e-5 … 6.437859095948209e-5 6.437859095983186e-5; … ; 6.437859104627274e-5 6.437859104627273e-5 … 6.437859095948209e-5 6.437859095983186e-5; 6.437859104627274e-5 6.437859104627274e-5 … 6.437859095983186e-5 6.437859096018008e-5;;; 6.437444666658912e-5 6.437444666658914e-5 … 6.437444675303001e-5 6.43744467526818e-5; 6.437444666658914e-5 6.437444666658914e-5 … 6.437444675337979e-5 6.437444675303001e-5; … ; 6.437444666658914e-5 6.437444666658914e-5 … 6.437444675337979e-5 6.437444675303001e-5; 6.437444666658912e-5 6.437444666658914e-5 … 6.437444675303001e-5 6.43744467526818e-5]
 [9.115090300876016e-9 9.18816509527637e-9 … 0.00013555878362043303 0.000133761104469766; 9.18816509527637e-9 9.262006848947153e-9 … 0.00013738531104727627 0.00013555878362043303; … ; 9.18816509527637e-9 9.262006848947153e-9 … 0.00013738531104727627 0.00013555878362043303; 9.115090300876016e-9 9.18816509527637e-9 … 0.00013555878362043303 0.000133761104469766;;; 0.00013556965510873605 0.00013556965510873586 … 0.00013556965429362562 0.00013556965430078604; 0.00013556965510873586 0.00013556965510873567 … 0.0001355696542863866 0.00013556965429362562; … ; 0.00013556965510873586 0.00013556965510873567 … 0.0001355696542863866 0.00013556965429362562; 0.00013556965510873605 0.00013556965510873586 … 0.00013556965429362562 0.00013556965430078604;;; 0.0001355512782014618 0.000135551278201462 … 0.00013555127901657225 0.00013555127900941182; 0.000135551278201462 0.0001355512782014622 … 0.00013555127902381128 0.00013555127901657225; … ; 0.000135551278201462 0.0001355512782014622 … 0.00013555127902381128 0.00013555127901657225; 0.0001355512782014618 0.000135551278201462 … 0.00013555127901657225 0.00013555127900941182]
 [2.6926897581315367e-8 2.7313507714770143e-8 … 0.00023370489649268138 0.00022845315379793862; 2.7313507714770143e-8 2.7707718844309493e-8 … 0.0002391075039605557 0.00023370489649268138; … ; 2.7313507714770143e-8 2.7707718844309493e-8 … 0.0002391075039605557 0.00023370489649268138; 2.6926897581315367e-8 2.7313507714770143e-8 … 0.00023370489649268138 0.00022845315379793862;;; 0.00023376199441293072 0.00023376199441292722 … 0.0002337619902052695 0.00023376199027161357; 0.00023376199441292722 0.00023376199441292367 … 0.00023376199013754758 0.0002337619902052695; … ; 0.00023376199441292722 0.00023376199441292367 … 0.00023376199013754758 0.0002337619902052695; 0.00023376199441293072 0.00023376199441292722 … 0.0002337619902052695 0.00023376199027161357;;; 0.00023370736165976154 0.0002337073616597651 … 0.00023370736586742275 0.0002337073658010787; 0.0002337073616597651 0.0002337073616597686 … 0.0002337073659351447 0.00023370736586742275; … ; 0.0002337073616597651 0.0002337073616597686 … 0.0002337073659351447 0.00023370736586742275; 0.00023370736165976154 0.0002337073616597651 … 0.00023370736586742275 0.0002337073658010787]
 [6.961732292649316e-8 7.124591837656885e-8 … 0.00037736420948259527 0.0003640011486945887; 7.124591837656885e-8 7.292787278493182e-8 … 0.0003913522086821116 0.00037736420948259527; … ; 7.124591837656885e-8 7.292787278493182e-8 … 0.0003913522086821116 0.00037736420948259527; 6.961732292649316e-8 7.124591837656885e-8 … 0.00037736420948259527 0.0003640011486945887;;; 0.0003776022093504568 0.0003776022093504163 … 0.00037760219156997705 0.00037760219202269337; 0.0003776022093504163 0.0003776022093503748 … 0.000377602191101566 0.00037760219156997705; … ; 0.0003776022093504163 0.0003776022093503748 … 0.000377602191101566 0.00037760219156997705; 0.0003776022093504568 0.0003776022093504163 … 0.00037760219156997705 0.00037760219202269337;;; 0.00037745967704351697 0.00037745967704355747 … 0.00037745969482399675 0.0003774596943712805; 0.00037745967704355747 0.00037745967704359894 … 0.00037745969529240787 0.00037745969482399675; … ; 0.00037745967704355747 0.00037745967704359894 … 0.00037745969529240787 0.00037745969482399675; 0.00037745967704351697 0.00037745967704355747 … 0.00037745969482399675 0.0003774596943712805]
 [1.540579894733325e-7 1.5945566561169678e-7 … 0.0005642799709081393 0.0005353147081300882; 1.5945566561169678e-7 1.6512254090905263e-7 … 0.0005952708190357286 0.0005642799709081393; … ; 1.5945566561169678e-7 1.6512254090905263e-7 … 0.0005952708190357286 0.0005642799709081393; 1.540579894733325e-7 1.5945566561169678e-7 … 0.0005642799709081393 0.0005353147081300882;;; 0.0005650578298830339 0.0005650578298827208 … 0.000565057770205711 0.0005650577724579198; 0.0005650578298827208 0.0005650578298823957 … 0.0005650577678343024 0.000565057770205711; … ; 0.0005650578298827208 0.0005650578298823957 … 0.0005650577678343024 0.000565057770205711; 0.0005650578298830339 0.0005650578298827208 … 0.000565057770205711 0.0005650577724579198;;; 0.0005647387130703668 0.0005647387130706799 … 0.0005647387727476897 0.0005647387704954811; 0.0005647387130706799 0.0005647387130710051 … 0.0005647387751190984 0.0005647387727476897; … ; 0.0005647387130706799 0.0005647387130710051 … 0.0005647387751190984 0.0005647387727476897; 0.0005647387130703668 0.0005647387130706799 … 0.0005647387727476897 0.0005647387704954811]
 [3.0969738789168484e-7 3.2507995129791364e-7 … 0.0008049880902276838 0.0007482706516369002; 3.2507995129791364e-7 3.415619008018545e-7 … 0.0008673145106382796 0.0008049880902276838; … ; 3.2507995129791364e-7 3.415619008018545e-7 … 0.0008673145106382796 0.0008049880902276838; 3.0969738789168484e-7 3.2507995129791364e-7 … 0.0008049880902276838 0.0007482706516369002;;; 0.0008071653959737827 0.0008071653959719099 … 0.0008071652219566872 0.0008071652311307488; 0.0008071653959719099 0.0008071653959699315 … 0.0008071652120858652 0.0008071652219566872; … ; 0.0008071653959719099 0.0008071653959699315 … 0.0008071652120858652 0.0008071652219566872; 0.0008071653959737827 0.0008071653959719099 … 0.0008071652219566872 0.0008071652311307488;;; 0.0008065143898932863 0.0008065143898951592 … 0.0008065145639103818 0.0008065145547363204; 0.0008065143898951592 0.0008065143898971374 … 0.0008065145737812039 0.0008065145639103818; … ; 0.0008065143898951592 0.0008065143898971374 … 0.0008065145737812039 0.0008065145639103818; 0.0008065143898932863 0.0008065143898951592 … 0.0008065145639103818 0.0008065145547363204]
 [5.705076773905195e-7 6.087951811632106e-7 … 0.0011000514541360741 0.0009988289301903872; 6.087951811632106e-7 6.508085123006421e-7 … 0.0012147497953614846 0.0011000514541360741; … ; 6.087951811632106e-7 6.508085123006421e-7 … 0.0012147497953614846 0.0011000514541360741; 5.705076773905195e-7 6.087951811632106e-7 … 0.0011000514541360741 0.0009988289301903872;;; 0.0011053528638517067 0.001105352863842754 … 0.0011053524171780035 0.0011053524484369427; 0.001105352863842754 0.0011053528638331044 … 0.0011053523826827932 0.0011053524171780035; … ; 0.001105352863842754 0.0011053528638331044 … 0.0011053523826827932 0.0011053524171780035; 0.0011053528638517067 0.001105352863842754 … 0.0011053524171780035 0.0011053524484369427;;; 0.001104132375968241 0.0011041323759771937 … 0.0011041328226419442 0.0011041327913830052; 0.0011041323759771937 0.0011041323759868433 … 0.0011041328571371547 0.0011041328226419442; … ; 0.0011041323759771937 0.0011041323759868433 … 0.0011041328571371547 0.0011041328226419442; 0.001104132375968241 0.0011041323759771937 … 0.0011041328226419442 0.0011041327913830052]
 [9.837734761082027e-7 1.069728314870665e-6 … 0.0014550919403635942 0.001286939535979884; 1.069728314870665e-6 1.1666270528831613e-6 … 0.0016522401555657791 0.0014550919403635942; … ; 1.069728314870665e-6 1.1666270528831613e-6 … 0.0016522401555657791 0.0014550919403635942; 9.837734761082027e-7 1.069728314870665e-6 … 0.0014550919403635942 0.001286939535979884;;; 0.0014666847725153 0.0014666847724792008 … 0.0014666837305651816 0.0014666838235512587; 0.0014666847724792008 0.0014666847724393811 … 0.0014666836249749018 0.0014666837305651816; … ; 0.0014666847724792008 0.0014666847724393811 … 0.0014666836249749018 0.0014666837305651816; 0.0014666847725153 0.0014666847724792008 … 0.0014666837305651816 0.0014666838235512587;;; 0.001464536697606192 0.0014645366976422912 … 0.0014645377395563106 0.0014645376465702335; 0.0014645366976422912 0.001464536697682111 … 0.0014645378451465904 0.0014645377395563106; … ; 0.0014645366976422912 0.001464536697682111 … 0.0014645378451465904 0.0014645377395563106; 0.001464536697606192 0.0014645366976422912 … 0.0014645377395563106 0.0014645376465702335]
 ⋮
 [0.0860980584145309 0.17043647381136437 … 0.5253487807847173 0.2653793414280966; 0.17043632607368323 0.3382712254382821 … 1.0431153649867155 0.5255803265164318; … ; 0.17043632607047352 0.338271225445665 … 1.0431153649877962 0.5255803265126668; 0.08609805841464986 0.1704364738091774 … 0.5253487807877247 0.265379341430695;;; 1.4373252284252627 1.381514056133696 … 1.1874256906452123 1.3236301768644763; 1.3815141490947984 1.282410462486958 … 0.9854451339241748 1.1873168599586519; … ; 1.3815141490947924 1.2824104624869688 … 0.9854451339241813 1.1873168599586483; 1.4373252284252602 1.3815140561336967 … 1.1874256906452232 1.3236301768644851;;; 0.5618219481395736 0.6176331204311375 … 0.8117214859196261 0.6755169997003598; 0.6176330274700377 0.7167367140778765 … 1.0137020426406604 0.811830316606183; … ; 0.6176330274700433 0.7167367140778665 … 1.0137020426406564 0.811830316606187; 0.5618219481395745 0.6176331204311389 … 0.8117214859196145 0.6755169997003496]
 [0.08611716284361551 0.17047468757767417 … 0.5254093393090755 0.265409623784774; 0.17047453796637485 0.33834765744416206 … 1.043236489945349 0.5256408960059097; … ; 0.17047453797624784 0.3383476574329872 … 1.0432364899545734 0.5256408960113826; 0.08611716284187021 0.17047468757644577 … 0.5254093392900094 0.26540962378754035;;; 1.4378006045760103 1.3819336252553542 … 1.1877342572082403 1.3240275973003055; 1.3819337199879664 1.2827436621439456 … 0.9856521249815257 1.1876253748112229; … ; 1.3819337199879878 1.282743662143918 … 0.9856521249815371 1.187625374811248; 1.4378006045760054 1.3819336252553445 … 1.1877342572081835 1.3240275973003033;;; 0.5618903141179099 0.6177572934385639 … 0.8119566614856801 0.6756633213936142; 0.6177571987059528 0.716947256549974 … 1.014038793712393 0.8120655438826948; … ; 0.6177571987059314 0.7169472565500006 … 1.014038793712381 0.812065543882672; 0.5618903141179137 0.6177572934385737 … 0.8119566614857363 0.675663321393615]
 [0.08612424653409514 0.1704888541714203 … 0.5254313531199226 0.26542063415503336; 0.17048870388305637 0.3383759883686652 … 1.0432805111692454 0.5256629133739961; … ; 0.1704887038883321 0.33837598835841554 … 1.0432805111290577 0.5256629134018219; 0.08612424653691521 0.17048885417575446 … 0.5254313531223933 0.2654206341449212;;; 1.4379761857636695 1.3820885642569618 … 1.187848575254823 1.3241746260446108; 1.3820886597402704 1.2828666720897481 … 0.9857290285709916 1.1877396740993025; … ; 1.382088659740295 1.2828666720897104 … 0.985729028570942 1.1877396740993493; 1.4379761857636821 1.3820885642569747 … 1.1878485752547856 1.3241746260445781;;; 0.5619155910916934 0.6178032125983992 … 0.81204320160054 0.675717150810749; 0.6178031171150915 0.7170251047656166 … 1.014162748284371 0.8121521027560599; … ; 0.6178031171150687 0.7170251047656533 … 1.01416274828442 0.8121521027560125; 0.5619155910916772 0.6178032125983878 … 0.8120432016005785 0.6757171508107845]
 [0.08612614504602102 0.17049265171457642 … 0.525437120686977 0.26542351832198846; 0.1704925012778496 0.33838358347625414 … 1.0432920459311534 0.5256686816812286; … ; 0.17049250125561458 0.33838358352274106 … 1.0432920462707862 0.5256686816281663; 0.08612614502946425 0.17049265174051192 … 0.5254371206311342 0.26542351827734434;;; 1.4380217482813722 1.3821288252436692 … 1.1878784162080114 1.3242128480702946; 1.382128920937962 1.282898770102663 … 0.9857493588365255 1.187769510236514; … ; 1.3821289209378955 1.2828987701027577 … 0.9857493588371241 1.1877695102364358; 1.4380217482813036 1.382128825243725 … 1.1878784162079201 1.3242128480701114;;; 0.5619220947768395 0.6178150178145437 … 0.81206542685021 0.6757309949879174; 0.6178149221202488 0.7170450729555525 … 1.0141944842216821 0.8121743328217007; … ; 0.6178149221203245 0.7170450729554643 … 1.0141944842210908 0.8121743328217819; 0.5619220947769019 0.6178150178144902 … 0.8120654268503005 0.6757309949881005]
 [0.0861263484550237 0.1704930588114406 … 0.5254376433059956 0.26542377925479826; 0.1704929083575903 0.3383843985532119 … 1.043293094769307 0.5256692044097837; … ; 0.1704929083179532 0.3383843985504839 … 1.0432930941125902 0.5256692047684617; 0.08612634841214753 0.17049305897868994 … 0.5254376437522809 0.2654237789478525;;; 1.4380264806957288 1.3821330324293595 … 1.1878815936547746 1.324216899594982; 1.3821331281547102 1.2829020891792937 … 0.9857513300943188 1.1877726872009684; … ; 1.3821331281546714 1.2829020891793919 … 0.9857513300935629 1.1877726872015324; 1.4380264806956726 1.382133032429838 … 1.1878815936554665 1.3242168995943688;;; 0.5619227864032379 0.6178162346696093 … 0.812067673444226 0.6757323675039573; 0.6178161389442847 0.7170471779196865 … 1.0141979370046275 0.8121765798980022; … ; 0.6178161389443182 0.7170471779195954 … 1.0141979370054217 0.8121765798974202; 0.561922786403297 0.61781623466914 … 0.8120676734434983 0.6757323675045981]
 [0.08612696916103961 0.1704942996211746 … 0.5254400816807748 0.26542500060545304; 0.17049414903877258 0.3383868795724996 … 1.043297975800099 0.5256716441583383; … ; 0.17049414919857298 0.33838687869515427 … 1.0432979721279094 0.5256716450098741; 0.08612696914432466 0.17049429987008752 … 0.525440082168004 0.26542500091475424;;; 1.4380418735129417 1.3821466102175435 … 1.1878911197329391 1.324229463323031; 1.3821467059253385 1.2829128644483128 … 0.9857575482619754 1.1877822109612353; … ; 1.3821467059252304 1.282912864446957 … 0.9857575482563989 1.1877822109616085; 1.4380418735133622 1.3821466102178275 … 1.1878911197332265 1.3242294633241454;;; 0.5619250181562303 0.6178202814514968 … 0.8120757719362344 0.675737428346068; 0.6178201857436643 0.7170540272208108 … 1.014209343407195 0.8121846807079487; … ; 0.6178201857438892 0.7170540272223528 … 1.014209343412782 0.8121846807075852; 0.5619250181559755 0.6178202814512223 … 0.8120757719358355 0.6757374283449257]
 [0.08612922377484898 0.17049880746230267 … 0.5254468148723012 0.26542836358957134; 0.17049865534065103 0.33839590092128835 … 1.0433114145796933 0.5256783805383086; … ; 0.17049865860770871 0.3383958958161563 … 1.0433114239576757 0.5256783777167134; 0.08612922322293583 0.17049880814212776 … 0.525446811840995 0.2654283642611259;;; 1.4380984980593468 1.382196601013655 … 1.1879282918393537 1.3242770694216377; 1.382196697047234 1.2829525933107075 … 0.9857827134840025 1.1878193774417978; … ; 1.3821966970533412 1.2829525932993169 … 0.9857827134916414 1.1878193774431431; 1.43809849805989 1.3821966010112683 … 1.1879282918325718 1.3242770694231056;;; 0.5619331295482851 0.6178350265953393 … 0.8121033357688703 0.6757545581872626; 0.6178349305618676 0.7170790342975233 … 1.0142489141247812 0.8122122501663108; … ; 0.6178349305548604 0.7170790343096268 … 1.014248914116363 0.812212250165756; 0.5619331295484676 0.6178350265969836 … 0.8121033357761148 0.675754558186087]
 [0.08612916330502197 0.17049869317371466 … 0.5254465304448525 0.2654282163618434; 0.17049854499628475 0.33839566059686466 … 1.0433108429459548 0.5256780726382269; … ; 0.17049854184749363 0.3383956617550718 … 1.0433108543112954 0.5256780602391141; 0.08612916171450978 0.17049869578421842 … 0.5254465182386399 0.2654282327286097;;; 1.438096899651587 1.3821951874950726 … 1.1879273352067259 1.324275791897684; 1.3821952836875342 1.2829514631453598 … 0.985782108932073 1.1878184210379235; … ; 1.3821952836664313 1.2829514631602468 … 0.9857821089456186 1.187818421030176; 1.4380968996444772 1.3821951874942873 … 1.187927335176687 1.3242757919422374;;; 0.5619328997439027 0.6178346119017598 … 0.8121024641884341 0.6757540074995592; 0.6178345157098714 0.7170783362480253 … 1.0142476904650342 0.8122113783598383; … ; 0.6178345157287075 0.7170783362385516 … 1.0142476904497297 0.8122113783638633; 0.5619328997500804 0.6178346119013959 … 0.8121024642139002 0.6757540074487606]
 [0.08612858502693178 0.17049753011854185 … 0.5254450995508656 0.2654275119357499; 0.17049737654539124 0.33839335967131645 … 1.0433080246587925 0.5256766560874475; … ; 0.17049739061254476 0.33839332923662874 … 1.0433080297881574 0.5256766653154703; 0.08612858461742168 0.17049753677333118 … 0.5254450841778983 0.2654275090734674;;; 1.4380834767959698 1.382183387135683 … 1.1879183611000284 1.3242643710114153; 1.3821834831479929 1.2829421711018256 … 0.9857759838047694 1.1878094482253942; … ; 1.3821834831841513 1.282942171049491 … 0.9857759838212017 1.1878094482498152; 1.4380834768058741 1.3821833871480351 … 1.18791836103523 1.324264370990561;;; 0.5619309119379512 0.6178310015967424 … 0.8120960276328201 0.6757500177235799; 0.617830905584656 0.7170722176292615 … 1.014238404929981 0.8122049405078227; … ; 0.6178309055491185 0.7170722176846236 … 1.0142384049122126 0.8122049404803533; 0.5619309119268662 0.6178310015838927 … 0.8120960276957662 0.675750017738761]

if I want to solve it on $t \in [0,100]$. Done! The solution gives back our tensors (and interpolates to create new ones if you use sol(t)). We can plot it in Plots.jl:

using Plots
p1 = surface(X, Y, sol[end][:, :, 1], title = "[A]")
p2 = surface(X, Y, sol[end][:, :, 2], title = "[B]")
p3 = surface(X, Y, sol[end][:, :, 3], title = "[C]")
plot(p1, p2, p3, layout = grid(3, 1))

and see the pretty gradients. Using this 2nd order ROCK method we solve this equation in about 2 seconds. That's okay.

Some Optimizations

There are some optimizations that can still be done. When we do A*B as matrix multiplication, we create another temporary matrix. These allocations can bog down the system. Instead, we can pre-allocate the outputs and use the inplace functions mul! to make better use of memory. The easiest way to store these cache arrays are constant globals, but you can use closures (anonymous functions which capture data, i.e. (x)->f(x,y)) or call-overloaded types to do it without globals. The globals way (the easy way) is simply:

const MyA = zeros(N, N)
const AMx = zeros(N, N)
const DA = zeros(N, N)
function f(du, u, p, t)
    A = @view u[:, :, 1]
    B = @view u[:, :, 2]
    C = @view u[:, :, 3]
    dA = @view du[:, :, 1]
    dB = @view du[:, :, 2]
    dC = @view du[:, :, 3]
    mul!(MyA, My, A)
    mul!(AMx, A, Mx)
    @. DA = D * (MyA + AMx)
    @. dA = DA + α₁ - β₁ * A - r₁ * A * B + r₂ * C
    @. dB = α₂ - β₂ * B - r₁ * A * B + r₂ * C
    @. dC = α₃ - β₃ * C + r₁ * A * B - r₂ * C
end

For reference, closures looks like:

MyA = zeros(N, N)
AMx = zeros(N, N)
DA = zeros(N, N)
function f_full(du, u, p, t, MyA, AMx, DA)
    A = @view u[:, :, 1]
    B = @view u[:, :, 2]
    C = @view u[:, :, 3]
    dA = @view du[:, :, 1]
    dB = @view du[:, :, 2]
    dC = @view du[:, :, 3]
    mul!(MyA, My, A)
    mul!(AMx, A, Mx)
    @. DA = D * (MyA + AMx)
    @. dA = DA + α₁ - β₁ * A - r₁ * A * B + r₂ * C
    @. dB = α₂ - β₂ * B - r₁ * A * B + r₂ * C
    @. dC = α₃ - β₃ * C + r₁ * A * B - r₂ * C
end
f(du, u, p, t) = f_full(du, u, p, t, MyA, AMx, DA)

and a call overloaded type looks like:

struct MyFunction{T} <: Function
    MyA::T
    AMx::T
    DA::T
end

# Now define the overload
function (ff::MyFunction)(du, u, p, t)
    # This is a function which references itself via ff
    A = @view u[:, :, 1]
    B = @view u[:, :, 2]
    C = @view u[:, :, 3]
    dA = @view du[:, :, 1]
    dB = @view du[:, :, 2]
    dC = @view du[:, :, 3]
    mul!(ff.MyA, My, A)
    mul!(ff.AMx, A, Mx)
    @. ff.DA = D * (ff.MyA + ff.AMx)
    @. dA = f.DA + α₁ - β₁ * A - r₁ * A * B + r₂ * C
    @. dB = α₂ - β₂ * B - r₁ * A * B + r₂ * C
    @. dC = α₃ - β₃ * C + r₁ * A * B - r₂ * C
end

MyA = zeros(N, N)
AMx = zeros(N, N)
DA = zeros(N, N)

f = MyFunction(MyA, AMx, DA)
# Now f(du,u,p,t) is our function!

These last two ways enclose the pointer to our cache arrays locally but still present a function f(du,u,p,t) to the ODE solver.

Now, since PDEs are large, many times we don't care about getting the whole timeseries. Using the output controls from DifferentialEquations.jl, we can make it only output the final timepoint.

prob = ODEProblem(f, u0, (0.0, 100.0))
@time sol = solve(prob, ROCK2(), progress = true, save_everystep = false,
                  save_start = false);
@time sol = solve(prob, ROCK2(), progress = true, save_everystep = false,
                  save_start = false);

Around 0.4 seconds. Much better. Also, if you're using VS Code, this'll give you a nice progress bar, so you can track how it's going.

Quick Note About Performance

Note

We are using the ROCK2 method here because it's a method for stiff equations with eigenvalues that are real-dominated (as opposed to dominated by the imaginary parts). If we wanted to use a more conventional implicit ODE solver, we would need to make use of the sparsity pattern. This is covered in the advanced ODE tutorial It turns out that ROCK2 is more efficient anyway (and doesn't require sparsity handling), so we will keep this setup.

Quick Summary: full PDE ODE Code

As a summary, here's a full PDE code:

using OrdinaryDiffEq, LinearAlgebra

# Define the constants for the PDE
const α₂ = 1.0
const α₃ = 1.0
const β₁ = 1.0
const β₂ = 1.0
const β₃ = 1.0
const r₁ = 1.0
const r₂ = 1.0
const D = 100.0
const γ₁ = 0.1
const γ₂ = 0.1
const γ₃ = 0.1
const N = 100
const X = reshape([i for i in 1:100 for j in 1:100], N, N)
const Y = reshape([j for i in 1:100 for j in 1:100], N, N)
const α₁ = 1.0 .* (X .>= 80)

const Mx = Array(Tridiagonal([1.0 for i in 1:(N - 1)], [-2.0 for i in 1:N],
                             [1.0 for i in 1:(N - 1)]))
const My = copy(Mx)
Mx[2, 1] = 2.0
Mx[end - 1, end] = 2.0
My[1, 2] = 2.0
My[end, end - 1] = 2.0

# Define the initial condition as normal arrays
u0 = zeros(N, N, 3)

const MyA = zeros(N, N);
const AMx = zeros(N, N);
const DA = zeros(N, N)
# Define the discretized PDE as an ODE function
function f(du, u, p, t)
    A = @view u[:, :, 1]
    B = @view u[:, :, 2]
    C = @view u[:, :, 3]
    dA = @view du[:, :, 1]
    dB = @view du[:, :, 2]
    dC = @view du[:, :, 3]
    mul!(MyA, My, A)
    mul!(AMx, A, Mx)
    @. DA = D * (MyA + AMx)
    @. dA = DA + α₁ - β₁ * A - r₁ * A * B + r₂ * C
    @. dB = α₂ - β₂ * B - r₁ * A * B + r₂ * C
    @. dC = α₃ - β₃ * C + r₁ * A * B - r₂ * C
end

# Solve the ODE
prob = ODEProblem(f, u0, (0.0, 100.0))
sol = solve(prob, ROCK2(), progress = true, save_everystep = false, save_start = false)

using Plots;
gr();
p1 = surface(X, Y, sol[end][:, :, 1], title = "[A]")
p2 = surface(X, Y, sol[end][:, :, 2], title = "[B]")
p3 = surface(X, Y, sol[end][:, :, 3], title = "[C]")
plot(p1, p2, p3, layout = grid(3, 1))

Making Use of GPU Parallelism

That was all using the CPU. How do we turn on GPU parallelism with DifferentialEquations.jl? Well, you don't. DifferentialEquations.jl "doesn't have GPU bits". So, wait... can we not do GPU parallelism? No, this is the glory of type-genericness, especially in broadcasted operations. To make things use the GPU, we simply use a CuArray from CUDA.jl. If instead of zeros(N,M) we used CuArray(zeros(N,M)), then the array lives on the GPU. CuArray naturally overrides broadcast such that dotted operations are performed on the GPU. DifferentialEquations.jl uses broadcast internally, and thus just by putting the array as a CuArray, the array-type will take over how all internal updates are performed and turn this algorithm into a fully GPU-parallelized algorithm that doesn't require copying to the CPU. Wasn't that simple?

From that you can probably also see how to multithread everything, or how to set everything up with distributed parallelism. You can make the ODE solvers do whatever you want by defining an array type where the broadcast does whatever special behavior you want.

So to recap, the entire difference from above is changing to:

using CUDA
const gMx = CuArray(Float32.(Mx))
const gMy = CuArray(Float32.(My))
const gα₁ = CuArray(Float32.(α₁))
gu0 = CuArray(Float32.(u0))

const gMyA = CuArray(zeros(Float32, N, N))
const AgMx = CuArray(zeros(Float32, N, N))
const gDA = CuArray(zeros(Float32, N, N))
function gf(du, u, p, t)
    A = @view u[:, :, 1]
    B = @view u[:, :, 2]
    C = @view u[:, :, 3]
    dA = @view du[:, :, 1]
    dB = @view du[:, :, 2]
    dC = @view du[:, :, 3]
    mul!(gMyA, gMy, A)
    mul!(AgMx, A, gMx)
    @. gDA = D * (gMyA + AgMx)
    @. dA = gDA + gα₁ - β₁ * A - r₁ * A * B + r₂ * C
    @. dB = α₂ - β₂ * B - r₁ * A * B + r₂ * C
    @. dC = α₃ - β₃ * C + r₁ * A * B - r₂ * C
end

prob2 = ODEProblem(gf, gu0, (0.0, 100.0))
CUDA.allowscalar(false) # makes sure none of the slow fallbacks are used
@time sol = solve(prob2, ROCK2(), progress = true, dt = 0.003, save_everystep = false,
                  save_start = false);

retcode: Success
Interpolation: 1st order linear
t: 1-element Vector{Float64}:
 100.0
u: 1-element Vector{CUDA.CuArray{Float32, 3, CUDA.Mem.DeviceBuffer}}:
 [0.30325767 0.30234623 … 0.9336535 0.9336841; 0.30325767 0.30234623 … 0.9336535 0.9336841; … ; 0.30325767 0.30234623 … 0.9336535 0.9336841; 0.30325767 0.30234623 … 0.9336535 0.9336841;;; 1.3027548 1.302752 … 1.0227438 1.0225888; 1.3027548 1.302752 … 1.0227438 1.0225888; … ; 1.3027548 1.302752 … 1.0227438 1.0225888; 1.3027548 1.302752 … 1.0227438 1.0225888;;; 0.69724476 0.697248 … 0.977256 0.97741115; 0.69724476 0.697248 … 0.977256 0.97741115; … ; 0.69724476 0.697248 … 0.977256 0.97741115; 0.69724476 0.697248 … 0.977256 0.97741115]

@time sol = solve(prob2, ROCK2(), progress = true, dt = 0.003, save_everystep = false,
                  save_start = false);

retcode: Success
Interpolation: 1st order linear
t: 1-element Vector{Float64}:
 100.0
u: 1-element Vector{CUDA.CuArray{Float32, 3, CUDA.Mem.DeviceBuffer}}:
 [0.30325767 0.30234623 … 0.9336535 0.9336841; 0.30325767 0.30234623 … 0.9336535 0.9336841; … ; 0.30325767 0.30234623 … 0.9336535 0.9336841; 0.30325767 0.30234623 … 0.9336535 0.9336841;;; 1.3027548 1.302752 … 1.0227438 1.0225888; 1.3027548 1.302752 … 1.0227438 1.0225888; … ; 1.3027548 1.302752 … 1.0227438 1.0225888; 1.3027548 1.302752 … 1.0227438 1.0225888;;; 0.69724476 0.697248 … 0.977256 0.97741115; 0.69724476 0.697248 … 0.977256 0.97741115; … ; 0.69724476 0.697248 … 0.977256 0.97741115; 0.69724476 0.697248 … 0.977256 0.97741115]

Go have fun.

And Stochastic PDEs?

Why not make it an SPDE? All that we need to do is extend each of the PDE equations to have a noise function. In this case, let's use multiplicative noise on each reactant. This means that our noise update equation is:

function g(du, u, p, t)
    A = @view u[:, :, 1]
    B = @view u[:, :, 2]
    C = @view u[:, :, 3]
    dA = @view du[:, :, 1]
    dB = @view du[:, :, 2]
    dC = @view du[:, :, 3]
    @. dA = γ₁ * A
    @. dB = γ₂ * A
    @. dC = γ₃ * A
end

g (generic function with 1 method)

Now we just define and solve the system of SDEs:

prob = SDEProblem(f, g, u0, (0.0, 100.0))
@time sol = solve(prob, SRIW1());

retcode: Success
Interpolation: 1st order linear
t: 81467-element Vector{Float64}:
   0.0
   9.999999999999999e-5
   0.0002125
   0.00033906249999999995
   0.0004814453125
   0.0006416259765625
   0.0008218292236328125
   0.0010245578765869141
   0.0012526276111602785
   0.0015092060625553133
   ⋮
  99.979329791279
  99.98115946335041
  99.98321784443074
  99.98553352314612
  99.9881386617009
  99.99106944257504
  99.99436657105845
  99.99807584060228
 100.0
u: 81467-element Vector{Array{Float64, 3}}:
 [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; … ; 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]
 [4.999999999999999e-9 4.999999999999999e-9 … 0.00010003354785815078 9.902717571726942e-5; 4.999999999999999e-9 4.999999999999999e-9 … 0.00010098271793934425 0.00010008456417499692; … ; 4.999999999999999e-9 4.999999999999999e-9 … 0.00010092693987644246 0.0001001298889611478; 4.999999999999999e-9 4.999999999999999e-9 … 0.00010017732665794973 9.897464715118086e-5;;; 9.999999999999999e-5 9.999999999999999e-5 … 9.990824717771995e-5 0.00010003719990013448; 9.999999999999999e-5 9.999999999999999e-5 … 9.991108399118227e-5 9.997925082521627e-5; … ; 9.999999999999999e-5 9.999999999999999e-5 … 0.00010000950788402055 0.00010002466975868248; 9.999999999999999e-5 9.999999999999999e-5 … 0.00010013644932389828 9.991843120171309e-5;;; 9.999e-5 9.999e-5 … 9.990927729023278e-5 9.996807434862905e-5; 9.999e-5 9.999e-5 … 0.00010003501928127977 0.00010001912869613252; … ; 9.999e-5 9.999e-5 … 9.999027398449927e-5 9.99810009536162e-5; 9.999e-5 9.999e-5 … 9.997142034528109e-5 9.990508815891823e-5]
 [2.232770332402062e-8 2.2575808134809752e-8 … 0.0002127593484354011 0.00020823922434927468; 2.2574767437676622e-8 2.281004149827771e-8 … 0.00021701526580228205 0.0002125734116370646; … ; 2.2574889597976524e-8 2.2809525229379915e-8 … 0.0002169536454208501 0.00021258387726039586; 2.232837910465318e-8 2.2563423672375118e-8 … 0.00021254917954630008 0.00020776445618653073;;; 0.000212499999153623 0.00021249998563036666 … 0.0002125552579920488 0.00021252065847997807; 0.00021249999802586578 0.00021249997808362922 … 0.00021220657212474812 0.00021250252495450483; … ; 0.00021249999318567089 0.00021249998470923882 … 0.00021255246045740186 0.0002124511481122418; 0.000212500008304568 0.00021250000373598064 … 0.00021272789735389338 0.00021274063741478404;;; 0.00021245486412802702 0.00021245484111985795 … 0.00021245567804994625 0.00021266588855043815; 0.0002124548414005349 0.0002124548489737366 … 0.0002125970159412083 0.0002125157915390588; … ; 0.00021245482753346007 0.00021245485079683096 … 0.00021245160020916414 0.00021243014219598009; 0.00021245486352009202 0.00021245486433182133 … 0.0002124104685896007 0.0002122852780015129]
 [5.6263069521506305e-8 5.7480099710388855e-8 … 0.0003389287400310471 0.00032853204684212624; 5.744926269270841e-8 5.8541136132626664e-8 … 0.0003505493697818147 0.0003388633189404726; … ; 5.74367973368229e-8 5.852500165562683e-8 … 0.00035019006145487174 0.00033888030754576626; 5.6370586748353135e-8 5.7485704562910854e-8 … 0.0003388608466110118 0.00032735588205481394;;; 0.000339062522178849 0.0003390624264913829 … 0.0003389694104183819 0.00033955504981340303; 0.0003390625284912161 0.0003390624836573401 … 0.00033811593171063827 0.000339163525603146; … ; 0.00033906240912412683 0.00033906244800592256 … 0.0003391751115305677 0.00033934737076385407; 0.00033906248700364786 0.0003390625254647619 … 0.0003389292250618276 0.0003395068185606192;;; 0.000338947591252421 0.0003389475724953584 … 0.00033840313043931487 0.0003388853623227377; 0.0003389474998212709 0.0003389475440777583 … 0.00033924964638918867 0.0003392372498889701; … ; 0.00033894754472576724 0.0003389475321321736 … 0.0003390341704440883 0.00033941253380550466; 0.00033894762733118774 0.0003389475815667519 … 0.0003391931687876087 0.00033821285095143253]
 [1.1262410455633155e-7 1.1582810464167765e-7 … 0.00048114329725648544 0.00046008334356349023; 1.1574103207669379e-7 1.1932448145228782e-7 … 0.0005039474187455218 0.000480219288403279; … ; 1.1576989527007816e-7 1.188781928984951e-7 … 0.0005037711515317602 0.00048127459553267093; 1.1249324014276081e-7 1.1600439099734273e-7 … 0.00047981850640432626 0.00045916719012083625;;; 0.0004814450703436235 0.00048144507681945466 … 0.0004814833970729153 0.0004820288414723127; 0.0004814453353869058 0.00048144550905593244 … 0.00048016308607861047 0.00048188100396471254; … ; 0.00048144530703097855 0.00048144521747663777 … 0.00048135555569016813 0.0004817744795084102; 0.0004814453438827381 0.00048144523405590264 … 0.0004815285174859866 0.0004815701416165949;;; 0.0004812137481275946 0.00048121371961576794 … 0.00048009501757346817 0.0004818950861971289; 0.00048121361903838325 0.0004812137374485358 … 0.0004807560737266274 0.00048194523813321075; … ; 0.0004812135731404426 0.0004812134501735442 … 0.0004800832639231105 0.00048154541737917; 0.0004812136307297277 0.0004812136918897018 … 0.00048074404114893004 0.0004807862680140094]
 [1.9786849281204892e-7 2.055097567278405e-7 … 0.0006398340430733396 0.0006045260152835097; 2.056352237909797e-7 2.137282286393747e-7 … 0.0006793217318625088 0.0006401035949620767; … ; 2.0526132721085307e-7 2.1356152230937373e-7 … 0.0006794559857370349 0.00064108280336269; 1.9800195766384715e-7 2.0565632922178843e-7 … 0.0006391403697569298 0.0006022138671876513;;; 0.0006416257027085002 0.0006416257040289718 … 0.0006422211679441112 0.0006405186107894082; 0.0006416259561041636 0.0006416263243805523 … 0.0006405714024260535 0.0006420505209904965; … ; 0.0006416259370316521 0.0006416258459231636 … 0.0006425845369222983 0.000642762314266274; 0.0006416262416411281 0.0006416261467327789 … 0.0006404799958434139 0.0006431212945278071;;; 0.0006412146287816358 0.0006412145282440409 … 0.0006391752044909052 0.0006414880679886268; 0.0006412145305963855 0.0006412146671435248 … 0.0006403675109885333 0.0006418546109538855; … ; 0.0006412143421290408 0.0006412142946271488 … 0.0006405675314569891 0.0006426701761245204; 0.0006412145910320194 0.0006412144054982365 … 0.0006409005474742935 0.0006405100894992011]
 [3.2025016584472927e-7 3.369452499501272e-7 … 0.0008198386525649481 0.0007625147722351089; 3.375848852754994e-7 3.5391563757935935e-7 … 0.0008828179046116646 0.0008199819420034498; … ; 3.364406910391081e-7 3.5414216064899767e-7 … 0.0008848663537117636 0.0008206590362794441; 3.202297839035182e-7 3.370383153393458e-7 … 0.0008195951715644601 0.0007602287229629449;;; 0.0008218285470911801 0.0008218288410153909 … 0.0008236965483751411 0.000822544150077032; 0.0008218293273299824 0.0008218294131257776 … 0.0008196320729442091 0.0008218368928479574; … ; 0.0008218288287049129 0.0008218292947357749 … 0.0008245341837134013 0.0008252016161314603; 0.0008218293142473482 0.0008218291088855823 … 0.0008198586034556337 0.0008234596238905908;;; 0.0008211539711640749 0.0008211545807249837 … 0.0008183933491751349 0.0008201918625747052; 0.0008211538251420636 0.0008211549005258915 … 0.0008202765639266234 0.0008233614384362694; … ; 0.0008211541888151829 0.000821154276951224 … 0.0008198076669053832 0.000823314994673141; 0.0008211545124650892 0.0008211545160738619 … 0.0008206413689914282 0.0008192231055456598]
 [4.91669236462011e-7 5.235510431701129e-7 … 0.0010200065336939844 0.0009336167520208591; 5.239877982250564e-7 5.580578519467854e-7 … 0.0011188863946627345 0.0010206351394063335; … ; 5.231497663389453e-7 5.568371626119661e-7 … 0.0011195792718640284 0.0010193107199864436; 4.918260617417772e-7 5.233425521867908e-7 … 0.001022817198817395 0.0009327214304045652;;; 0.0010245583324299042 0.0010245572604919023 … 0.0010253827396595973 0.0010277139503110267; 0.0010245581782423235 0.0010245587611285448 … 0.0010208902031936548 0.0010234008402144348; … ; 0.0010245576533994395 0.0010245579157079094 … 0.0010292870545445863 0.0010285481031061733; 0.0010245579694529944 0.0010245570160247143 … 0.0010240075342661105 0.001024535614899195;;; 0.001023508890658365 0.001023509160219746 … 0.0010204590142400085 0.001021278549101574; 0.0010235086444944858 0.001023509353027609 … 0.0010236095255985649 0.0010259512503348045; … ; 0.0010235080950972942 0.0010235089739193456 … 0.001022495286178245 0.0010245533227794714; 0.0010235076327887063 0.0010235089742623579 … 0.001020939907186191 0.0010203942433524953]
 [7.254049764973987e-7 7.80789968820091e-7 … 0.00124384185084368 0.001121420962719144; 7.818473450345788e-7 8.448963491114685e-7 … 0.0013868293070193403 0.0012471946051881748; … ; 7.815971396854572e-7 8.42087789264485e-7 … 0.0013879249439040073 0.0012449738581438675; 7.253119831849194e-7 7.808245565687058e-7 … 0.0012454471625427943 0.0011214491061158876;;; 0.0012526279041697001 0.001252624700406325 … 0.0012523931799472 0.0012577080594158296; 0.0012526276931773108 0.0012526272115548002 … 0.001247851453489419 0.0012509506526417206; … ; 0.0012526279046805857 0.0012526280632851018 … 0.0012569297929144268 0.0012546935781605568; 0.00125262771635558 0.0012526276703783363 … 0.0012520918303974646 0.0012521830796510297;;; 0.0012510593680085762 0.0012510594094758793 … 0.001246971492361705 0.001248764479038949; 0.0012510600480871324 0.0012510590248065933 … 0.0012474861276281408 0.0012517330310673677; … ; 0.0012510611750547842 0.0012510602007754325 … 0.0012542233280401513 0.0012531694385646253; 0.0012510581341302082 0.001251059526715911 … 0.0012502143070211307 0.001247793553639316]
 [1.0382465413051875e-6 1.1323717634813291e-6 … 0.0014959739459812745 0.0013233291275459892; 1.1333889673484837e-6 1.2439934235469783e-6 … 0.0016965279259684547 0.0014977382687929618; … ; 1.1311273388596255e-6 1.2379915549438189e-6 … 0.0017007602944787393 0.001500378367988204; 1.036749572726524e-6 1.1350923713470047e-6 … 0.0015012646069862774 0.0013224543132201471;;; 0.0015092060033015536 0.001509202922759721 … 0.0015055987203454382 0.0015183859086361619; 0.0015092036276462162 0.001509205528193198 … 0.0015033786975595714 0.0015081776941393457; … ; 0.0015092058557681514 0.0015092038650898563 … 0.0015176185751494629 0.001508530154234868; 0.0015092056491239738 0.0015092059799452263 … 0.0015111889626266176 0.0015057395143603964;;; 0.0015069297165554323 0.0015069282171418254 … 0.0015034413016725133 0.001502927737444058; 0.0015069315856247184 0.0015069283637019323 … 0.001499810615667775 0.0015080903394359272; … ; 0.0015069334776396542 0.0015069296010481275 … 0.0015094781162582561 0.001506667273168418; 0.0015069305861959339 0.0015069285905516935 … 0.0015074960245848952 0.00150521958326599]
 ⋮
 [0.09434541973721874 0.18508621993933 … 0.538146999071017 0.2689893106415949; 0.1870730882475272 0.3696462679464202 … 1.0777596704582395 0.5378566516428678; … ; 0.18537856056444654 0.37069416334496214 … 1.084361582743963 0.5436104881639716; 0.09425240179939072 0.18650292679306968 … 0.5423211832126797 0.2748778517910903;;; 1.3650511424189362 1.332413244138798 … 1.0712287928269535 1.4458568302097843; 1.4195077012714592 1.1729090138093985 … 1.740073467859827 1.6141730400073138; … ; 1.2780752991170647 1.2645887279936716 … 1.0166096443308574 0.9669020861909463; 1.4214860353900092 1.348019382539887 … 1.2869764204108245 1.3756202619085276;;; 0.5351659321465314 0.6051649433117224 … 0.8033213824634423 0.5235515744635356; 0.5778574194769819 0.7314115300151492 … 1.8964856095240448 0.8033957222267737; … ; 0.6571260020607604 1.0293915944890404 … 1.352928138228693 0.7512909634309822; 0.5543834529565181 0.5650678244616844 … 0.7319779128441922 0.7027188764044349]
 [0.09399327532259022 0.18643546238866937 … 0.5397055575654344 0.2688526289613556; 0.18635620619866003 0.368140687322 … 1.080695067417763 0.5395344315572943; … ; 0.18576336417453593 0.3677427041031214 … 1.0866586689735507 0.5447054230714783; 0.09381985446526839 0.1847941270013041 … 0.5466400651202619 0.27510870389744524;;; 1.3648683066181453 1.3314746136282838 … 1.074559493605835 1.4449005051092898; 1.4187835069994374 1.1715895429058296 … 1.7378619331409828 1.6148692975724743; … ; 1.278726303577731 1.266257691751873 … 1.013254473243479 0.9664597210531208; 1.4218653130100531 1.3472931951303597 … 1.2843093972179815 1.3760106176460252;;; 0.535971821065098 0.6068644539748497 … 0.8027289818912191 0.5240690508056133; 0.57883208845581 0.7318158006282568 … 1.8956218391400903 0.8019822719814098; … ; 0.6567150207584515 1.0269859544130304 … 1.3492154280698359 0.7498822663986733; 0.5536921426444729 0.5649967026372423 … 0.7337506692685548 0.7018580266993044]
 [0.09382561104328792 0.18602505658290838 … 0.5389958199255546 0.26993268448816327; 0.18581107881773762 0.3675778005498663 … 1.080756121099945 0.5394570611373763; … ; 0.18551785467142157 0.36577176914375154 … 1.0874809252288782 0.5457898400215652; 0.09396359854045067 0.185107287319693 … 0.5487040894136039 0.27599131170331165;;; 1.3642067178637833 1.330474574406424 … 1.0783055264947685 1.444797947701405; 1.4185946214257266 1.169404416535743 … 1.7353873780573212 1.6154278834626106; … ; 1.2792322179384776 1.2648403376113029 … 1.0146796429395915 0.9642682763416207; 1.4225312595874655 1.3478658840926028 … 1.286321218053824 1.3771581317283197;;; 0.5359646691510308 0.6068884068129783 … 0.8027883675331848 0.5238813694660039; 0.5788199460214966 0.734186883348326 … 1.900604941551974 0.8032260918671097; … ; 0.6549431605379535 1.0286258588936004 … 1.3457664788022747 0.7513345483232781; 0.5537015353186234 0.564708145945671 … 0.7327630309150178 0.702692905354324]
 [0.09408676495154014 0.18553425600493287 … 0.5429992806198932 0.27154128050865073; 0.18633435542914123 0.3671095918231597 … 1.089924735409712 0.5414416899501411; … ; 0.18427407578314822 0.3628329476485171 … 1.078679796058594 0.5452355261988019; 0.09307288500879453 0.1833198779461843 … 0.5458727376198304 0.2779763994833913;;; 1.3640970037606135 1.3318270133937467 … 1.0795136486398937 1.4435922129873395; 1.420246463557604 1.1683594719594506 … 1.7352763609955606 1.61162096791031; … ; 1.2793436735080033 1.2637003967448939 … 1.01556136106967 0.9634730012045186; 1.4220492789148933 1.3469041482965365 … 1.2850520587290897 1.3768574696218516;;; 0.5366712455867839 0.6062834088215784 … 0.8007051436282004 0.5275296497148732; 0.5785418734209137 0.7336349237908307 … 1.8985964768816763 0.80383862497938; … ; 0.6548299052260185 1.0285222729617538 … 1.3428363157612535 0.7530352602495695; 0.5538617777109296 0.5655787176826499 … 0.7364020407268415 0.7037738674633416]
 [0.09384675507594395 0.18557834577368662 … 0.5462255765594364 0.2733457597454012; 0.18577130459785268 0.36702191530082484 … 1.0911415859614613 0.5415335467090744; … ; 0.1829087550743569 0.36280547821273745 … 1.073026736126931 0.5429061058310543; 0.09243767004711435 0.18290789617688627 … 0.5458963325467107 0.2768736823066773;;; 1.364284751970085 1.332614358804723 … 1.0813177544661026 1.4435101788343485; 1.4202645772632512 1.167598235434254 … 1.7393901623900114 1.611699597610917; … ; 1.2798626200043073 1.263046317516736 … 1.0153809280412134 0.9677451349163817; 1.4218341818795488 1.3461221387531055 … 1.285686849473725 1.3740196811914673;;; 0.5366637744631295 0.6056355595014555 … 0.8020843320741354 0.5299114473205434; 0.578732242116293 0.7323001474925743 … 1.901528223397652 0.8008614522468483; … ; 0.6546178991822356 1.0262283818986102 … 1.3338405266395459 0.7500539150249581; 0.5538399980607901 0.566041407311132 … 0.7325813002749884 0.7025788593474429]
 [0.09360695448418628 0.18477462765781005 … 0.5450373008639975 0.27460423813479334; 0.18550644310679437 0.36648936387687886 … 1.0892721319793175 0.5420775600196578; … ; 0.18270772972171395 0.36248684659553737 … 1.076176962850606 0.5432597078940415; 0.09274819389090909 0.1823704129236062 … 0.5451441987774687 0.2757512445088988;;; 1.3635304744699888 1.3327518518764212 … 1.0816021774762408 1.440652975384843; 1.4206233443752894 1.1658586199551113 … 1.7348377027173076 1.6106241540090855; … ; 1.2790367211486664 1.2651617295596738 … 1.0168023891441098 0.9750524402018921; 1.4204069245856061 1.3457948391559393 … 1.2814625593165685 1.3724386431074156;;; 0.5361969127742536 0.6071835547669522 … 0.8013555755045405 0.5315568937802654; 0.5780561757903732 0.7266242141393872 … 1.9067200631461627 0.8049999117127731; … ; 0.6552607609962763 1.022055408105506 … 1.3296411204806549 0.7488297094175835; 0.5540648434254295 0.5669035844319097 … 0.7391608093857108 0.6997691169204152]
 [0.09344691697027357 0.18529874358595352 … 0.5448830755906984 0.27712385272158535; 0.18459443308152634 0.36431432055856194 … 1.097125128916498 0.5402293071814896; … ; 0.1811883327238567 0.3605177310396028 … 1.0735331669920383 0.545218258826362; 0.09183602119020472 0.1811529741821661 … 0.5475181012667136 0.2739280730562613;;; 1.3642179746522098 1.3331244958101187 … 1.0817196593012468 1.4389198127319445; 1.4206400734197173 1.1641169344828592 … 1.724732213013308 1.6135837149566659; … ; 1.2795287313770671 1.2672039106842616 … 1.0185759982631448 0.9758128721761322; 1.4202443693506022 1.3453249909873581 … 1.2798251039984279 1.3693098813916333;;; 0.5365909563733068 0.6061002862842293 … 0.8034164244950217 0.534010377711294; 0.5773518694845631 0.725870465374598 … 1.8970778036506337 0.8049957817768209; … ; 0.6548366674868152 1.0184305558916702 … 1.3366061711772004 0.7481224116498103; 0.5542766606822793 0.5661696162756069 … 0.7368531014688237 0.6981632792106875]
 [0.09337110800803788 0.18487715057857582 … 0.5424545710414578 0.2829143177247511; 0.18479413495539512 0.36388008177472186 … 1.103391857799131 0.5383927699026776; … ; 0.1791367824058736 0.35803227195046405 … 1.0637817823791131 0.553274689270747; 0.0918837120471803 0.17900382382450183 … 0.5559713037066539 0.27198242217725127;;; 1.363730781907621 1.332546045214861 … 1.0861363924372909 1.4417157171026458; 1.420010594744118 1.165853469026919 … 1.7205829105774768 1.6103433165858443; … ; 1.2807585399746952 1.2655468545841657 … 1.0288127266974927 0.9724190515304509; 1.4189786533096338 1.3448363176283287 … 1.2789376744659937 1.3691133028701585;;; 0.5368267628608511 0.6056232775435331 … 0.8067104064662555 0.5358336295988847; 0.5796987719784064 0.7242238602594467 … 1.8859391095644733 0.80851914884237; … ; 0.6550492273484197 1.019652742636251 … 1.332138763343938 0.7519183445569416; 0.5537673682117981 0.5684498425349267 … 0.7307024994290051 0.6988421471512095]
 [0.09333978650454249 0.18467648928912234 … 0.547534889629764 0.28107298314472445; 0.18458627959940221 0.3636043167133434 … 1.1016233543820793 0.5449651487215106; … ; 0.17902087309956696 0.3565734667859469 … 1.0639654820819093 0.548678677744122; 0.09127397549111538 0.17972656303006973 … 0.5511320095995607 0.2712803837370611;;; 1.3638447297704235 1.3319004371681107 … 1.0899768734057667 1.4417626285161624; 1.419295605367196 1.166083817607662 … 1.7213556945289574 1.609181149742114; … ; 1.2803819780788588 1.2665856141904217 … 1.028317394730353 0.9738286949541706; 1.4191066435080382 1.344745080516717 … 1.2758749709833315 1.3691325430232602;;; 0.5375868914440133 0.605792536970359 … 0.8054950704148123 0.5360289338859267; 0.580071341483524 0.724206281689264 … 1.8829495070524709 0.8131555304406916; … ; 0.6545980775433827 1.0177762697363633 … 1.3329776482047602 0.7510110568557846; 0.5536661036305505 0.5684038524489858 … 0.7297293317096695 0.6977722954492733]

using Plots;
gr();

# Use `Array` to transform the result back into a CPU-based `Array` for plotting
p1 = surface(X, Y, Array(sol[end][:, :, 1]), title = "[A]")
p2 = surface(X, Y, Array(sol[end][:, :, 2]), title = "[B]")
p3 = surface(X, Y, Array(sol[end][:, :, 3]), title = "[C]")
plot(p1, p2, p3, layout = grid(3, 1))

We can see the cool effect that diffusion dampens the noise in [A] but is unable to dampen the noise in [B] which results in a very noisy [C]. The stiff SPDE takes much longer to solve even using high order plus adaptivity because stochastic problems are just that much more difficult (current research topic is to make new algorithms for this!). It gets GPU'd just by using CuArray like before. But there we go: solving systems of stochastic PDEs using high order adaptive algorithms with within-method GPU parallelism. That's gotta be a first? The cool thing is that nobody ever had to implement the GPU-parallelism either, it just exists by virtue of the Julia type system.

(Note: We can also use one of the SROCK methods for better performance here, but they will require a choice of dt. This is left to the reader to try.)

Note

This can take a while to solve! An explicit Runge-Kutta algorithm isn't necessarily great here, though to use a stiff solver on a problem of this size requires once again smartly choosing sparse linear solvers. The high order adaptive method is pretty much necessary though, since something like Euler-Maruyama is simply not stable enough to solve this at a reasonable dt. Also, the current algorithms are not so great at handling this problem. Good thing there's a publication coming along with some new stuff...