Welcome to pystella’s documentation!#
pystella
is a Python-based framework enabling the easy expression and solution
of partial differential equations.
Here’s a simple example which evolves the scalar wave equation (without doing
anything interesting):
import pyopencl as cl
import pyopencl.array as cla
import pyopencl.clrandom as clr
import pystella as ps
# set parameters
grid_shape = (32, 32, 32)
proc_shape = (1, 1, 1)
rank_shape = tuple(Ni // pi for Ni, pi in zip(grid_shape, proc_shape))
halo_shape = 1
dtype = "float64"
dx = tuple(10 / Ni for Ni in grid_shape)
dt = min(dx) / 10
# create pyopencl context, queue, and halo-sharer
ctx = ps.choose_device_and_make_context()
queue = cl.CommandQueue(ctx)
decomp = ps.DomainDecomposition(proc_shape, halo_shape, rank_shape)
# initialize arrays with random data
f = clr.rand(queue, tuple(ni + 2 * halo_shape for ni in rank_shape), dtype)
dfdt = clr.rand(queue, tuple(ni + 2 * halo_shape for ni in rank_shape), dtype)
lap_f = cla.zeros(queue, rank_shape, dtype)
# define system of equations
f_ = ps.DynamicField("f", offset="h") # don't overwrite f
rhs_dict = {
f_: f_.dot, # df/dt = \dot{f}
f_.dot: f_.lap # d\dot{f}/dt = \nabla^2 f
}
# create time-stepping and derivative-computing kernels
stepper = ps.LowStorageRK54(rhs_dict, dt=dt, halo_shape=halo_shape)
derivs = ps.FiniteDifferencer(decomp, halo_shape, dx)
t = 0.
# loop over time
while t < 1.:
for s in range(stepper.num_stages):
derivs(queue, fx=f, lap=lap_f)
stepper(s, queue=queue, f=f, dfdt=dfdt, lap_f=lap_f)
t += dt
pystella
uses loopy
for code generation and pyopencl
for
execution on CPUs and GPUs, with MPI parallelization across multiple OpenCL devices
via mpi4py
.
For a more detailed tutorial on the tools to generate OpenCL kernels provided by
loopy
and pystella
, see
codegen-tutorial.ipynb.
For a complete example which simulates scalar-field preheating after
cosmological inflation (optionally including gravitational waves), see
scalar_preheating.py
(but note that grid_size
and end_time
are small by default so testing is
faster).
Table of Contents#
Please check Installation to get started.