Particle-in-Cell advection ready to rock the GPU ๐
Warning
This package is still under development and the API is not stable yet.
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, AMDGPUBackendand load the required packages:
using JustPIC, JustPIC._2D
using GLMakieDefine 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), yvwhere 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)
endNow 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 stepwhere 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)
endwhere 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)
fThe development of this package is supported by the GPU4GEO PASC project.