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)