Speeding up stencil computations on CPUs and GPUs
Project description
pystencils
Run blazingly fast stencil codes on numpy arrays.
pystencils uses sympy to define stencil operations, that can be executed on numpy arrays. Exploiting the stencil structure makes pystencils run faster than normal numpy code and even as Cython and numba, as demonstrated in this notebook.
Here is a code snippet that computes the average of neighboring cells:
import pystencils as ps
import numpy as np
f, g = ps.fields("f, g : [2D]")
stencil = ps.Assignment(g[0, 0],
(f[1, 0] + f[-1, 0] + f[0, 1] + f[0, -1]) / 4)
kernel = ps.create_kernel(stencil).compile()
f_arr = np.random.rand(1000, 1000)
g_arr = np.empty_like(f_arr)
kernel(f=f_arr, g=g_arr)
pystencils is mostly used for numerical simulations using finite difference or finite volume methods. It comes with automatic finite difference discretization for PDEs:
c, v = ps.fields("c, v(2): [2D]")
adv_diff_pde = ps.fd.transient(c) - ps.fd.diffusion(c, sp.symbols("D")) + ps.fd.advection(c, v)
discretize = ps.fd.Discretization2ndOrder(dx=1, dt=0.01)
discretization = discretize(adv_diff_pde)
Installation
pip install pystencils[interactive]
Without [interactive]
you get a minimal version with very little dependencies.
All options:
gpu
: use this if an NVIDIA GPU is available and CUDA is installedopencl
: basic OpenCL support (experimental)alltrafos
: pulls in additional dependencies for loop simplification e.g. libislbench_db
: functionality to store benchmark result in object databasesinteractive
: installs dependencies to work in Jupyter including image I/O, plotting etc.autodiff
: enable derivation of adjoint kernels and generation of Torch/Tensorflow operationsdoc
: packages to build documentation
Options can be combined e.g.
pip install pystencils[interactive,gpu,doc]
Documentation
Read the docs here and
check out the Jupyter notebooks in doc/notebooks
.
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.