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],