Forked from
waLBerla / waLBerla
945 commits behind the upstream repository.
-
Markus Holzer authoredMarkus Holzer authored
Poisson.py 621 B
import sympy as sp
import pystencils as ps
from pystencils_walberla import CodeGeneration, generate_sweep
with CodeGeneration() as ctx:
# ----- Solving the 2D Poisson equation with rhs --------------------------
dx = sp.Symbol("dx")
dy = sp.Symbol("dy")
src, dst, rhs = ps.fields("src, src_tmp, rhs: [2D]", layout='fzyx')
@ps.kernel
def kernel_func():
src[0, 0] @= ((dy**2 * (src[1, 0] + src[-1, 0]))
+ (dx**2 * (src[0, 1] + src[0, -1]))
- (rhs[0, 0] * dx**2 * dy**2)) / (2 * (dx**2 + dy**2))
generate_sweep(ctx, 'Poisson', kernel_func)