def test_histogram(ctx_factory, grid_shape, proc_shape, dtype, num_bins,
timing=False):
if ctx_factory:
ctx = ctx_factory()
else:
ctx = ps.choose_device_and_make_context()
queue = cl.CommandQueue(ctx)
h = 1
mpi = ps.DomainDecomposition(proc_shape, h, grid_shape=grid_shape)
rank_shape, _ = mpi.get_rank_shape_start(grid_shape)
if np.dtype(dtype) in (np.dtype("float64"), np.dtype("complex128")):
max_rtol, avg_rtol = 1e-10, 1e-11
else:
max_rtol, avg_rtol = 5e-4, 5e-5
from pymbolic import var
_fx = ps.Field("fx")
histograms = {
"count": (var("abs")(_fx) * num_bins, 1),
"squared": (var("abs")(_fx) * num_bins, _fx**2),
}
hist = ps.Histogrammer(mpi, histograms, num_bins, dtype, rank_shape=rank_shape)
rng = clr.ThreefryGenerator(ctx, seed=12321)
fx = rng.uniform(queue, rank_shape, dtype)
fx_h = fx.get()
result = hist(queue, fx=fx)
res = result["count"]
assert np.sum(res.astype("int64")) == np.product(grid_shape), \
f"Count histogram doesn't sum to grid_size ({np.sum(res)})"
bins = np.linspace(0, 1, num_bins+1).astype(dtype)
weights = np.ones_like(fx_h)
np_res = np.histogram(fx_h, bins=bins, weights=weights)[0]
np_res = mpi.allreduce(np_res)
max_err, avg_err = get_errs(res, np_res)
assert max_err < max_rtol and avg_err < avg_rtol, \
f"Histogrammer inaccurate for grid_shape={grid_shape}" \
f": {max_err=}, {avg_err=}"
res = result["squared"]
np_res = np.histogram(fx_h, bins=bins, weights=fx_h**2)[0]
np_res = mpi.allreduce(np_res)
max_err, avg_err = get_errs(res, np_res)
assert max_err < max_rtol and avg_err < avg_rtol, \
f"Histogrammer with weights inaccurate for grid_shape={grid_shape}" \
f": {max_err=}, {avg_err=}"
if timing:
from common import timer
t = timer(lambda: hist(queue, fx=fx))
print(f"histogram took {t:.3f} ms for {grid_shape=}, {dtype=}")
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 test_field_diff(proc_shape):
if proc_shape != (1, 1, 1):
pytest.skip("test field only on one rank")
from pystella import diff
y = ps.Field("y")
assert diff(y, y) == 1
assert diff(y[0], y[0]) == 1
assert diff(y[0], y[1]) == 0
y = ps.DynamicField("y")
assert diff(y, y) == 1
assert diff(y[0], y[0]) == 1
assert diff(y[0], y[1]) == 0
import pymbolic.primitives as pp
assert diff(y**3, y, "t") == pp.Product((3, 2, y, y.d(0)))
assert diff(y**3, "t", y) == pp.Product((3, y.d(0), 2, y))
for i, x in enumerate(["t", "x", "y", "z"]):
assert diff(y, x) == y.d(i)
assert diff(y[1, 3], x) == y.d(1, 3, i)
assert diff(y[1]**2, x) == 2 * y[1] * y.d(1, i)
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_step(ctx_factory, proc_shape, dtype, Stepper):
if proc_shape != (1, 1, 1):
pytest.skip("test step only on one rank")
if ctx_factory:
ctx = ctx_factory()
else:
ctx = ps.choose_device_and_make_context()
queue = cl.CommandQueue(ctx)
from pystella.step import LowStorageRKStepper
is_low_storage = LowStorageRKStepper in Stepper.__bases__
rank_shape = (1, 1, 8)
init_vals = np.linspace(1, 3, 8)
if is_low_storage:
y = cla.zeros(queue, rank_shape, dtype)
y[0, 0, :] = init_vals
y0 = y.copy()
else:
num_copies = Stepper.num_copies
y = cla.zeros(queue, (num_copies, ) + rank_shape, dtype)
y[0, 0, 0, :] = init_vals
y0 = y[0].copy()
dtlist = [.1, .05, .025]
for n in [-1., -2., -3., -4.]:
max_errs = {}
for dt in dtlist:
def sol(y0, t):
return ((-1 + n) * (-t + y0**(1 - n) / (-1 + n)))**(1 /
(1 - n))
_y = ps.Field("y")
rhs_dict = {_y: _y**n}
stepper = Stepper(rhs_dict,
dt=dt,
halo_shape=0,
rank_shape=rank_shape)
if is_low_storage:
y[0, 0, :] = init_vals
else:
y[0, 0, 0, :] = init_vals
t = 0
errs = []
while t < .1:
for s in range(stepper.num_stages):
stepper(s, queue=queue, y=y, filter_args=True)
t += dt
if is_low_storage:
errs.append(cla.max(clm.fabs(1. - sol(y0, t) / y)).get())
else:
errs.append(
cla.max(clm.fabs(1. - sol(y0, t) / y[0])).get())
max_errs[dt] = np.max(errs)
order = stepper.expected_order
print(f"{order=}, {n=}")
print(max_errs)
print([
max_errs[a] / max_errs[b] for a, b in zip(dtlist[:-1], dtlist[1:])
])
order = stepper.expected_order
rtol = dtlist[-1]**order if dtype == np.float64 else 1e-1
assert list(max_errs.values())[-1] < rtol, \
f"Stepper solution inaccurate for {n=}"
for a, b in zip(dtlist[:-1], dtlist[1:]):
assert max_errs[a] / max_errs[b] > .9 * (a/b)**order, \
f"Stepper convergence failing for {n=}"
def test_low_storage_edge_codegen_and_tmp_alloc(ctx_factory,
proc_shape,
dtype=None):
if proc_shape != (1, 1, 1):
pytest.skip("test step only on one rank")
if ctx_factory:
ctx = ctx_factory()
else:
ctx = ps.choose_device_and_make_context()
queue = cl.CommandQueue(ctx)
from pystella import LowStorageRK54
from pymbolic import parse
rhs_dict = {
parse("y[i, j, k]"): 1,
}
stepper = LowStorageRK54(rhs_dict, dt=.1, halo_shape=0)
y = cla.zeros(queue, (8, 8, 8), "complex128")
tmp_arrays = stepper.get_tmp_arrays_like(y=y)
assert tmp_arrays["_y_tmp"].shape == y.shape
assert tmp_arrays["_y_tmp"].dtype == y.dtype
stepper(0, queue=queue, y=y)
tmp_for_check = stepper.tmp_arrays["_y_tmp"]
stepper(1, queue=queue, y=y)
assert tmp_for_check is stepper.tmp_arrays["_y_tmp"]
rhs_dict = {
parse("y"): 1,
}
stepper = LowStorageRK54(rhs_dict,
args=[lp.GlobalArg("y", shape=tuple())],
dt=.1,
halo_shape=0)
y = np.zeros(1)
tmp_arrays = stepper.get_tmp_arrays_like(y=y)
assert tmp_arrays["_y_tmp"].shape == y.shape
assert tmp_arrays["_y_tmp"].dtype == y.dtype
# kernel won't work
# stepper(0, queue=queue, y=y)
# tmp_for_check = stepper.tmp_arrays["_y_tmp"]
# stepper(1, queue=queue, y=y)
# assert tmp_for_check is stepper.tmp_arrays["_y_tmp"]
rhs_dict = {
ps.Field(parse("y[0, 0]"), shape=(2, 2)): 1,
ps.Field(parse("y[0, 1]"), shape=(2, 2)): 1,
ps.Field(parse("y[1, 0]"), shape=(2, 2)): 1,
ps.Field(parse("y[1, 1]"), shape=(2, 2)): 1,
}
stepper = LowStorageRK54(rhs_dict, dt=.1, halo_shape=0)
y = cla.zeros(queue, (2, 2, 12, 12, 12), "float64")
tmp_arrays = stepper.get_tmp_arrays_like(y=y)
assert tmp_arrays["_y_tmp"].shape == y.shape
assert tmp_arrays["_y_tmp"].dtype == y.dtype
stepper(0, queue=queue, y=y)
tmp_for_check = stepper.tmp_arrays["_y_tmp"]
stepper(1, queue=queue, y=y)
assert tmp_for_check is stepper.tmp_arrays["_y_tmp"]
rhs_dict = {
ps.Field("y", shape=(1, 2))[0, 1]: 1,
}
stepper = LowStorageRK54(rhs_dict, dt=.1, halo_shape=0)
y = cla.zeros(queue, (1, 2, 12, 12, 12), "float64")
tmp_arrays = stepper.get_tmp_arrays_like(y=y)
assert tmp_arrays["_y_tmp"].shape == y.shape
assert tmp_arrays["_y_tmp"].dtype == y.dtype
stepper(0, queue=queue, y=y)
tmp_for_check = stepper.tmp_arrays["_y_tmp"]
stepper(1, queue=queue, y=y)
assert tmp_for_check is stepper.tmp_arrays["_y_tmp"]
rhs_dict = {
ps.Field("y"): 1,
ps.Field("z"): 1,
}
stepper = LowStorageRK54(rhs_dict, dt=.1, halo_shape=0)
y = cla.zeros(queue, (12, 12, 12), "float64")
z = cla.zeros(queue, (12, 12, 12), "complex128")
tmp_arrays = stepper.get_tmp_arrays_like(y=y, z=z)
assert tmp_arrays["_y_tmp"].shape == y.shape
assert tmp_arrays["_y_tmp"].dtype == y.dtype
assert tmp_arrays["_z_tmp"].shape == z.shape
assert tmp_arrays["_z_tmp"].dtype == z.dtype
# create output function
if decomp.rank == 0:
from pystella.output import OutputFile
out = OutputFile(ctx=ctx, runfile=__file__)
else:
out = None
statistics = ps.FieldStatistics(decomp,
halo_shape,
rank_shape=rank_shape,
grid_size=grid_size)
spectra = ps.PowerSpectra(decomp, fft, dk, volume)
projector = ps.Projector(fft, halo_shape, dk, dx)
hist = ps.FieldHistogrammer(decomp, 1000, dtype, rank_shape=rank_shape)
a_sq_rho = (3 * mpl**2 * ps.Field("hubble", indices=[])**2 / 8 / np.pi)
rho_dict = {ps.Field("rho"): scalar_sector.stress_tensor(0, 0) / a_sq_rho}
compute_rho = ps.ElementWiseMap(rho_dict,
halo_shape=halo_shape,
rank_shape=rank_shape)
def output(step_count, t, energy, expand, f, dfdt, lap_f, dfdx, hij, dhijdt,
lap_hij):
if step_count % 4 == 0:
f_stats = statistics(f)
if decomp.rank == 0:
out.output("energy",
t=t,
a=expand.a[0],