Particle-in-Cell advection ready to rock the GPU 🚀
The first step is to chose whether we want to run our simulation on the CPU or Nvidia or AMD GPUs. This is done by setting the backend
variable to either CUDABackend
, AMDGPUBackend
or CPUBackend
. In the following we will assume that we are running on a Nvidia GPU.
const backend = CUDABackend # Options: CPUBackend, CUDABackend, AMDGPUBackend
and load the required packages:
using JustPIC, JustPIC._2D
using GLMakie
Define domain and grids of the domain:
# number of grid points
n = 257
# number of cells
nx = ny = n-1
# domain size
Lx = Ly = 1.0
# nodal vertices
xvi = xv, yv = range(0, Lx, length=n), range(0, Ly, length=n)
# grid spacing
dxi = dx, dy = xv[2] - xv[1], yv[2] - yv[1]
# nodal centers
xci = xc, yc = range(0+dx/2, Lx-dx/2, length=n-1), range(0+dy/2, Ly-dy/2, length=n-1)
# staggered grid velocity nodal locations
grid_vx = xv, expand_range(yc)
grid_vy = expand_range(xc), yv
where expand_range
is a helper function that expands a range by one element on each side:
function expand_range(x::AbstractRange)
dx = x[2] - x[1]
n = length(x)
x1, x2 = extrema(x)
xI = round(x1-dx; sigdigits=5)
xF = round(x2+dx; sigdigits=5)
range(xI, xF, length=n+2)
end
Now we can initialize the particles object
nxcell = 24 # initial number of particles per cell
max_xcell = 48 # maximum number of particles per cell
min_xcell = 12 # minimum number of particles per cell
particles = init_particles(
backend, nxcell, max_xcell, min_xcell, xvi, dxi, (nx, ny)
)
The velocity field is defined by the stream function
vx_stream(x, y) = 250 * sin(π*x) * cos(π*y)
vy_stream(x, y) = -250 * cos(π*x) * sin(π*y)
and therefore the velocity field (with ghost nodes) on the staggered grid is given by
Vx = TA(backend)([vx_stream(x, y) for x in grid_vx[1], y in grid_vx[2]]);
Vy = TA(backend)([vy_stream(x, y) for x in grid_vy[1], y in grid_vy[2]]);
V = Vx, Vy
dt = min(dx / maximum(abs.(Vx)), dy / maximum(abs.(Vy))) # time step
where TA(backend)
is a type alias for either Array
, CuArray
or ROCArray
.
We save the initial particle positions:
pxv = particles.coords[1].data;
pyv = particles.coords[2].data;
idxv = particles.index.data;
p = [(pxv[idxv][1], pyv[idxv][1])]
chose the advection scheme
advection_scheme = RungeKutta2()
and finally perform the time iterations:
# Advection test
particle_args = ()
niter = 750
for iter in 1:niter
# advect particles
advection!(particles, advection_scheme, V, (grid_vx, grid_vy), dt)
# shuffle particles in the memory to keep the spatial locality tight
move_particles!(particles, xvi, particle_args)
# save particle position
pxv = particles.coords[1].data;
pyv = particles.coords[2].data;
idxv = particles.index.data;
p_i = (pxv[idxv][1], pyv[idxv][1])
push!(p, p_i)
end
where particle_args
is an empty tuple, but typically it contains the fields that are advected with the particles (e.g. temperature). At last, we plot the particle trajectory on top of the stream function:
using GLMakie
g(x) = Point2f(
vx_stream(x[1], x[2]),
vy_stream(x[1], x[2])
)
f, ax, = streamplot(g, xvi...)
lines!(ax, p, color=:red)
f
The development of this package is supported by the GPU4GEO PASC project.