pystella 2020.1

A code generator for grid-based PDE solving on CPUs and GPUs

pystella enables the easy expression of both PDE systems and the algorithms to solve them in high-performance computing environments within Python. It provides interfaces to generate custom computational kernels via loopy which can then be executed from Python on (multiple) CPUs or GPUs using pyopencl and mpi4py. Moreover, pystella implements a number of algorithms for PDE time evolution and spatial discretization that can be readily applied to a variety of physical systems.

Its features include:

• code generation for performant element-wise kernels, stencil-based computations, reductions, and histograms

• distributed domain decomposition and grid boundary sychronization

• time-stepping algorithms, including low-storage Runge-Kutta schemes

• finite difference and spectral collocation methods for spatial derivatives

• a geometric multigrid solver for generic sets of nonlinear boundary-value problems (in beta)

• wrappers to OpenCL-based Fast Fourier Transforms (FFTs) and distributed CPU FFTs

• Fourier space methods for field analysis and solving Poisson’s equation

All of the above functionality is configured to run at high performance by default, as are the interfaces for generating custom kernels (though this is entirely user-configurable!). Additionally, the provided functionality is intended to work seamlessly whether running in distributed- (i.e., multiple devices) or shared-memory (i.e., a single device) contexts without sacrificing performance in either case.

pystella was designed to simulate preheating and gravitational wave production after cosmological inflation and provides a simple way to specify models of this process. However, pystella is also designed to be sufficiently abstract as to provide a good framework for most systems that can be discretized onto grids (e.g., lattice field theory, (magneto)hydrodynamics, Einstein’s equations, electromagnetism, etc.). The preheating-specific components can be viewed as examples for the symbolic representation of arbitrary physical systems as an interface to its code generation routines. pystella provides entry points at varying levels of abstraction—so if you like the idea of pystella but the algorithms you require are not implemented, you can create new interfaces (or extend existing ones) for your purposes with ease. (Better yet, consider contributing a PR!)

pystella is fully documented and is licensed under the liberal MIT license.