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.

Indices and tables#