Bifurcation diagram 2D reaction diffusion type system on GPU

Hi all,

This is my first time posting, so forgive me if I am not following the proper etiquette.

I am currently working on computing bifurcation diagrams of 2D reaction diffusion type models (e.g. Brusselator model) with zero flux boundary conditions. I can successfully compute diagrams, but when increasing the resolution of the domain the computations become quite slow. I already tried to optimize the code by using the analytical expression for the Jacobian matrix and using a shift invariant eigensolver, but this only gets me so far.

Because of this I am interested in performing these computations on the GPU. I tried to adapt my CPU code based on this tutorial 🟠 2d Swift-Hohenberg equation (non-local) on the GPU · Bifurcation Analysis in Julia. However, this is (in my opinion ;)) not a very simple example to start with.

Hence my question: are there any more resources available for computing bifurcation diagrams on the GPU? Any other tips, tricks or general advice is also greatly appreciated.

Thank you in advance!

Sorry I missed your post. Is it still on?

Yes I am still working on this, but currently on lower resolutions on the CPU. I would still be interested to know how to do these computations on the GPU, so I am still interested if you have any recommendations. To give some more specifics: my model is similar to the Brusselator system like this: 🟠 1d Brusselator · Bifurcation Analysis in Julia but then in 2D and with zero flux boundary conditions.

You want to compute periodic orbits or equilibria?

In any case, the only bottleneck is “inverting” the jacobian. I know that the Laplacian precondtionner is quite good. So either you use sparse matrices on the GPU (you probably have to modify my CGLeq model) or you use FFTs.

I want to compute equilibria. I am not super familiar with using FFTs, so my preference would be to just use finite differences and sparse matrices on the GPU. Given the simple example below (2D Ginzburg-Landau from the tutorials) could you maybe give me some concrete pointers what you would change here to make it run on the GPU?

using Revise, DifferentialEquations
using BifurcationKit, LinearAlgebra, Plots, SparseArrays
const BK = BifurcationKit

function Laplacian2D(Nx, Ny, lx, ly)
    hx = 2lx/Nx
    hy = 2ly/Ny
    D2x = spdiagm(0 => -2ones(Nx), 1 => ones(Nx-1), -1 => ones(Nx-1) ) / hx^2
    D2y = spdiagm(0 => -2ones(Ny), 1 => ones(Ny-1), -1 => ones(Ny-1) ) / hy^2

    D2x[1,1] = -2/hx^2
    D2x[end,end] = -2/hx^2

    D2y[1,1] = -2/hy^2
    D2y[end,end] = -2/hy^2

    D2xsp = sparse(D2x)
    D2ysp = sparse(D2y)
    A = kron(sparse(I, Ny, Ny), D2xsp) + kron(D2ysp, sparse(I, Nx, Nx))
    return A, D2x
end


function NL!(f, u, p, t = 0.)
	(;r, μ, ν, c3, c5) = p
	n = div(length(u), 2)
	u1 = @view u[1:n]
	u2 = @view u[n+1:2n]

	ua = u1.^2 .+ u2.^2

	f1 = @view f[1:n]
	f2 = @view f[n+1:2n]

	@. f1 .= r * u1 - ν * u2 - ua * (c3 * u1 - μ * u2) - c5 * ua^2 * u1
	@. f2 .= r * u2 + ν * u1 - ua * (c3 * u2 + μ * u1) - c5 * ua^2 * u2

	return f
end

function Fcgl!(f, u, p, t = 0.)
	mul!(f, p.Δ, u)
	f .= f .+ NL(u, p)
end

NL(u, p) = NL!(similar(u), u, p)
Fcgl(u, p, t = 0.) = Fcgl!(similar(u), u, p, t)

function Jcgl(u, p, t = 0.)
	(;r, μ, ν, c3, c5, Δ) = p

	n = div(length(u), 2)
	u1 = @view u[1:n]
	u2 = @view u[n+1:2n]

	ua = u1.^2 .+ u2.^2

	f1u = zero(u1)
	f2u = zero(u1)
	f1v = zero(u1)
	f2v = zero(u1)

	@. f1u =  r - 2 * u1 * (c3 * u1 - μ * u2) - c3 * ua - 4 * c5 * ua * u1^2 - c5 * ua^2
	@. f1v = -ν - 2 * u2 * (c3 * u1 - μ * u2)  + μ * ua - 4 * c5 * ua * u1 * u2
	@. f2u =  ν - 2 * u1 * (c3 * u2 + μ * u1)  - μ * ua - 4 * c5 * ua * u1 * u2
	@. f2v =  r - 2 * u2 * (c3 * u2 + μ * u1) - c3 * ua - 4 * c5 * ua * u2 ^2 - c5 * ua^2

	jacdiag = vcat(f1u, f2v)

	Δ + spdiagm(0 => jacdiag, n => f1v, -n => f2u)
end

Nx = 100
Ny = 100
n = Nx*Ny
lx = pi
ly = pi/2

Δ = Laplacian2D(Nx, Ny, lx, ly)[1]
par_cgl = (r = 0.5, μ = 0.1, ν = 1.0, c3 = -1.0, c5 = 1.0, Δ = blockdiag(Δ, Δ))
sol0 = 0.1rand(2Nx, Ny)
sol0_f = vec(sol0)

prob = BK.BifurcationProblem(Fcgl, sol0_f, par_cgl, (@optic _.r); J = Jcgl)

eigls = EigArpack(1.0, :LM)
opt_newton = NewtonPar(tol = 1e-9, verbose = true, eigsolver = eigls, max_iterations = 20)
opts_br = ContinuationPar(dsmax = 0.02, ds = 0.01, p_max = 2., nev = 15, newton_options = (@set opt_newton.verbose = false))
br = @time continuation(prob, PALC(), opts_br, verbosity = 0)

plot(br)

It requires some work to adapt it to GPU and I have “never” done it with sparse matrices.
Let us first forget about computing the stability of of the equilibria.

In BK, equilibria are computed using newton algorithm. Hence, you need to invert the jacobian which is here sparse. The function Jcgl compute the jacobian but you have to port it to GPU.
This should be “simple” if you rely on https://github.com/exanauts/CUDSS.jl.

Keep us updated!

Some useful links: