def test_field(proc_shape):
if proc_shape != (1, 1, 1):
pytest.skip("test field only on one rank")
y = ps.Field("y", offset="h")
result = ps.index_fields(y)
assert result == parse("y[i + h, j + h, k + h]"), result
y = ps.Field("y", offset="h", indices=("a", "b", "c"))
result = ps.index_fields(y)
assert result == parse("y[a + h, b + h, c + h]"), result
y = ps.Field("y", ignore_prepends=True)
result = ps.index_fields(y, prepend_with=(0, 1))
assert result == parse("y[i, j, k]"), result
y = ps.Field("y[4, 5]", ignore_prepends=True)
result = ps.index_fields(y, prepend_with=(0, 1))
assert result == parse("y[4, 5, i, j, k]"), result
y = ps.Field("y", ignore_prepends=True)
result = ps.index_fields(y[2, 3], prepend_with=(0, 1))
assert result == parse("y[2, 3, i, j, k]"), result
y = ps.Field("y[4, 5]", ignore_prepends=True)
result = ps.index_fields(y[2, 3], prepend_with=(0, 1))
assert result == parse("y[2, 3, 4, 5, i, j, k]"), result
y = ps.Field("y", ignore_prepends=False)
result = ps.index_fields(y, prepend_with=(0, 1))
assert result == parse("y[0, 1, i, j, k]"), result
y = ps.Field("y[4, 5]", ignore_prepends=False)
result = ps.index_fields(y, prepend_with=(0, 1))
assert result == parse("y[0, 1, 4, 5, i, j, k]"), result
y = ps.Field("y", ignore_prepends=False)
result = ps.index_fields(y[2, 3], prepend_with=(0, 1))
assert result == parse("y[0, 1, 2, 3, i, j, k]"), result
y = ps.Field("y[4, 5]", ignore_prepends=False)
result = ps.index_fields(y[2, 3], prepend_with=(0, 1))
assert result == parse("y[0, 1, 2, 3, 4, 5, i, j, k]"), result
y = ps.Field("y", offset=("hx", "hy", "hz"))
result = ps.index_fields(shift_fields(y, (1, 2, 3)))
assert result == parse("y[i + hx + 1, j + hy + 2, k + hz + 3]"), result
y = ps.Field("y", offset=("hx", var("hy"), "hz"))
result = ps.index_fields(shift_fields(y, (1, 2, var("a"))))
assert result == parse("y[i + hx + 1, j + hy + 2, k + hz + a]"), result
def test_sympy_interop(proc_shape):
if proc_shape != (1, 1, 1):
pytest.skip("test field only on one rank")
from pystella.field.sympy import pymbolic_to_sympy, sympy_to_pymbolic
import sympy as sym
f = ps.Field("f", offset="h")
g = ps.Field("g", offset="h")
expr = f[0]**2 * g + 2 * g[1] * f
sympy_expr = pymbolic_to_sympy(expr)
new_expr = sympy_to_pymbolic(sympy_expr)
sympy_expr_2 = pymbolic_to_sympy(new_expr)
assert sym.simplify(sympy_expr - sympy_expr_2) == 0, \
"sympy <-> pymbolic conversion not invertible"
expr = f + shift_fields(f, (1, 2, 3))
sympy_expr = pymbolic_to_sympy(expr)
new_expr = sympy_to_pymbolic(sympy_expr)
sympy_expr_2 = pymbolic_to_sympy(new_expr)
assert sym.simplify(sympy_expr - sympy_expr_2) == 0, \
"sympy <-> pymbolic conversion not invertible with shifted indices"
# from pymbolic.functions import fabs, exp, exmp1
fabs = parse("math.fabs")
exp = parse("math.exp")
expm1 = parse("math.expm1")
x = sym.Symbol("x")
expr = sym.Abs(x)
assert sympy_to_pymbolic(expr) == fabs(var("x"))
expr = sym.exp(x)
assert sympy_to_pymbolic(expr) == exp(var("x"))
expr = sym.Function("expm1")(x) # pylint: disable=E1102
assert sympy_to_pymbolic(expr) == expm1(var("x"))
expr = sym.Function("aaa")(x) # pylint: disable=E1102
from pymbolic.primitives import Call, Variable
assert sympy_to_pymbolic(expr) == Call(Variable("aaa"), (Variable("x"), ))
def expand_stencil(f, coefs):
"""
Expands a stencil over a field.
:arg f: A :class:`~pystella.Field`.
:arg coefs: A :class:`dict` whose values are the coefficients of the stencil
at an offset given by the key. The keys must be 3-:class:`tuple`\\ s, and the
values may be :mod:`pymbolic` expressions or constants.
Example::
>>> f = Field("f", offset="h")
>>> coefs = {(1, 0, 0): 1, (-1, 0, 0): -1}
>>> stencil = expand_stencil(f, coefs)
>>> print(index_fields(stencil))
f[i + h + 1, j + h, k + h] + (-1)*f[i + h + -1, j + h, k + h]
"""
return sum([c * shift_fields(f, shift=offset) for offset, c in coefs.items()])
def test_get_field_args(proc_shape):
if proc_shape != (1, 1, 1):
pytest.skip("test field only on one rank")
from pystella import Field, DynamicField, get_field_args
x = Field("x", offset=(1, 2, 3))
y = Field("y", offset="h")
z = DynamicField("z", shape=(2, "a"))
import loopy as lp
true_args = [
lp.GlobalArg("x", shape="(Nx+2, Ny+4, Nz+6)", offset=lp.auto),
lp.GlobalArg("y", shape="(Nx+2*h, Ny+2*h, Nz+2*h)", offset=lp.auto),
lp.GlobalArg("z", shape="(2, a, Nx, Ny, Nz)", offset=lp.auto),
lp.GlobalArg("dzdx", shape="(2, a, 3, Nx, Ny, Nz)", offset=lp.auto),
]
def lists_equal(a, b):
equal = True
for x in a:
equal *= x in b
for x in b:
equal *= x in a
return equal
expressions = {x: y, y: x * z + z.pd[0]}
args = get_field_args(expressions)
assert lists_equal(args, true_args)
expressions = x * y + z + z.pd[2]
args = get_field_args(expressions)
assert lists_equal(args, true_args)
expressions = [x, y, y * z**2, 3 + z.pd[0] + z.pd[1]]
args = get_field_args(expressions)
assert lists_equal(args, true_args)
expressions = [shift_fields(x, (1, 2, 3)), y + z.pd[0], y * z**2]
args = get_field_args(expressions)
assert lists_equal(args, true_args)
def test_collect_field_indices(proc_shape):
if proc_shape != (1, 1, 1):
pytest.skip("test field only on one rank")
from pystella import Field, DynamicField
from pystella.field import collect_field_indices
x = Field("x", offset=(1, 2, 3))
y = Field("y", indices=("i", "x"), offset="h")
z = DynamicField("z", shape=(2, "a"))
expressions = {x: y, y: x * z + z.pd[0]}
indices = collect_field_indices(expressions)
assert indices == {"i", "j", "k", "x"}
expressions = [x, z]
indices = collect_field_indices(expressions)
assert indices == {"i", "j", "k"}
expressions = [shift_fields(x, (1, 2, 3)), y + z.pd[0], y * z**2]
indices = collect_field_indices(expressions)
assert indices == {"i", "j", "k", "x"}