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_vector_projector(ctx_factory, grid_shape, proc_shape, h, dtype, timing=False): if ctx_factory: ctx = ctx_factory() else: ctx = ps.choose_device_and_make_context() queue = cl.CommandQueue(ctx) mpi = ps.DomainDecomposition(proc_shape, h, grid_shape=grid_shape) rank_shape, _ = mpi.get_rank_shape_start(grid_shape) pencil_shape = tuple(ni + 2 * h for ni in rank_shape) L = (10, 8, 11.5) dx = tuple(Li / Ni for Li, Ni in zip(L, grid_shape)) dk = tuple(2 * np.pi / Li for Li in L) fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype) cdtype = fft.cdtype if h > 0: stencil = FirstCenteredDifference(h) project = ps.Projector(fft, stencil.get_eigenvalues, dk, dx) derivs = ps.FiniteDifferencer(mpi, h, dx) else: project = ps.Projector(fft, lambda k, dx: k, dk, dx) derivs = ps.SpectralCollocator(fft, dk) vector_x = cla.empty(queue, (3, ) + pencil_shape, dtype) div = cla.empty(queue, rank_shape, dtype) pdx = cla.empty(queue, (3, ) + rank_shape, dtype) def get_divergence_error(vector): for mu in range(3): fft.idft(vector[mu], vector_x[mu]) derivs.divergence(queue, vector_x, div) derivs(queue, fx=vector_x[0], pdx=pdx[0]) derivs(queue, fx=vector_x[1], pdy=pdx[1]) derivs(queue, fx=vector_x[2], pdz=pdx[2]) norm = sum([clm.fabs(pdx[mu]) for mu in range(3)]) max_err = cla.max(clm.fabs(div)) / cla.max(norm) avg_err = cla.sum(clm.fabs(div)) / cla.sum(norm) return max_err, avg_err max_rtol = 1e-11 if dtype == np.float64 else 1e-4 avg_rtol = 1e-13 if dtype == np.float64 else 1e-5 k_shape = fft.shape(True) vector = cla.empty(queue, (3, ) + k_shape, cdtype) for mu in range(3): vector[mu] = make_data(queue, fft).astype(cdtype) project.transversify(queue, vector) max_err, avg_err = get_divergence_error(vector) assert max_err < max_rtol and avg_err < avg_rtol, \ f"transversify failed for {grid_shape=}, {h=}: {max_err=}, {avg_err=}" plus = make_data(queue, fft).astype(cdtype) minus = make_data(queue, fft).astype(cdtype) project.pol_to_vec(queue, plus, minus, vector) if isinstance(fft, gDFT): assert all(is_hermitian(vector[i]) for i in range(3)), \ f"pol->vec is non-hermitian for {grid_shape=}, {h=}" max_err, avg_err = get_divergence_error(vector) assert max_err < max_rtol and avg_err < avg_rtol, \ f"pol_to_vec result not transverse for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" vector_h = vector.get() vector_2 = cla.zeros_like(vector) project.transversify(queue, vector, vector_2) vector_2_h = vector_2.get() max_err, avg_err = get_errs(vector_h, vector_2_h) assert max_err < max_rtol and avg_err < avg_rtol, \ f"pol->vector != its own transverse proj. for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" plus1 = cla.zeros_like(plus) minus1 = cla.zeros_like(minus) project.vec_to_pol(queue, plus1, minus1, vector) if isinstance(fft, gDFT): assert is_hermitian(plus1) and is_hermitian(minus1), \ f"polarizations aren't hermitian for {grid_shape=}, {h=}" max_err, avg_err = get_errs(plus1.get(), plus.get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"pol->vec->pol (plus) is not identity for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" max_err, avg_err = get_errs(minus1.get(), minus.get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"pol->vec->pol (minus) is not identity for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" project.vec_to_pol(queue, vector[0], vector[1], vector) max_err, avg_err = get_errs(plus1.get(), vector[0].get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"in-place pol->vec->pol (plus) not identity for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" max_err, avg_err = get_errs(minus1.get(), vector[1].get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"in-place pol->vec->pol (minus) not identity for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" # reset and test longitudinal component for mu in range(3): vector[mu] = make_data(queue, fft).astype(cdtype) fft.idft(vector[mu], vector_x[mu]) long = cla.zeros_like(minus) project.decompose_vector(queue, vector, plus1, minus1, long) long_x = cla.empty(queue, pencil_shape, dtype) fft.idft(long, long_x) div_true = cla.empty(queue, rank_shape, dtype) derivs.divergence(queue, vector_x, div_true) derivs(queue, fx=long_x, grd=pdx) div_long = cla.empty(queue, rank_shape, dtype) if h != 0: pdx_h = cla.empty(queue, (3, ) + pencil_shape, dtype) for mu in range(3): mpi.restore_halos(queue, pdx[mu], pdx_h[mu]) derivs.divergence(queue, pdx_h, div_long) else: derivs.divergence(queue, pdx, div_long) max_err, avg_err = get_errs(div_true.get(), div_long.get()) assert max_err < 1e-6 and avg_err < 1e-11, \ f"lap(longitudinal) != div vector for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" if timing: from common import timer ntime = 10 t = timer(lambda: project.transversify(queue, vector), ntime=ntime) print(f"transversify took {t:.3f} ms for {grid_shape=}") t = timer(lambda: project.pol_to_vec(queue, plus, minus, vector), ntime=ntime) print(f"pol_to_vec took {t:.3f} ms for {grid_shape=}") t = timer(lambda: project.vec_to_pol(queue, plus, minus, vector), ntime=ntime) print(f"vec_to_pol took {t:.3f} ms for {grid_shape=}") t = timer( lambda: project.decompose_vector(queue, vector, plus, minus, long), ntime=ntime) print(f"decompose_vector took {t:.3f} ms for {grid_shape=}")
def test_tensor_projector(ctx_factory, grid_shape, proc_shape, h, dtype, timing=False): if ctx_factory: ctx = ctx_factory() else: ctx = ps.choose_device_and_make_context() queue = cl.CommandQueue(ctx) mpi = ps.DomainDecomposition(proc_shape, h, grid_shape=grid_shape) rank_shape, _ = mpi.get_rank_shape_start(grid_shape) L = (10, 8, 11.5) dx = tuple(Li / Ni for Li, Ni in zip(L, grid_shape)) dk = tuple(2 * np.pi / Li for Li in L) fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype) cdtype = fft.cdtype if h > 0: stencil = FirstCenteredDifference(h) project = ps.Projector(fft, stencil.get_eigenvalues, dk, dx) derivs = ps.FiniteDifferencer(mpi, h, dx) else: project = ps.Projector(fft, lambda k, dx: k, dk, dx) derivs = ps.SpectralCollocator(fft, dk) vector_x = cla.empty(queue, (3, ) + tuple(ni + 2 * h for ni in rank_shape), dtype) div = cla.empty(queue, rank_shape, dtype) pdx = cla.empty(queue, (3, ) + rank_shape, dtype) def get_divergence_errors(hij): max_errors = [] avg_errors = [] for i in range(1, 4): for mu in range(3): fft.idft(hij[tensor_id(i, mu + 1)], vector_x[mu]) derivs.divergence(queue, vector_x, div) derivs(queue, fx=vector_x[0], pdx=pdx[0]) derivs(queue, fx=vector_x[1], pdy=pdx[1]) derivs(queue, fx=vector_x[2], pdz=pdx[2]) norm = sum([clm.fabs(pdx[mu]) for mu in range(3)]) max_errors.append(cla.max(clm.fabs(div)) / cla.max(norm)) avg_errors.append(cla.sum(clm.fabs(div)) / cla.sum(norm)) return np.array(max_errors), np.array(avg_errors) max_rtol = 1e-11 if dtype == np.float64 else 1e-4 avg_rtol = 1e-13 if dtype == np.float64 else 1e-5 def get_trace_errors(hij_h): trace = sum([hij_h[tensor_id(i, i)] for i in range(1, 4)]) norm = np.sqrt( sum(np.abs(hij_h[tensor_id(i, i)])**2 for i in range(1, 4))) trace = np.abs(trace[norm != 0]) / norm[norm != 0] trace = trace[trace < .9] return np.max(trace), np.sum(trace) / trace.size k_shape = fft.shape(True) hij = cla.empty(queue, shape=(6, ) + k_shape, dtype=cdtype) for mu in range(6): hij[mu] = make_data(queue, fft).astype(cdtype) project.transverse_traceless(queue, hij) hij_h = hij.get() if isinstance(fft, gDFT): assert all(is_hermitian(hij_h[i]) for i in range(6)), \ f"TT projection is non-hermitian for {grid_shape=}, {h=}" max_err, avg_err = get_divergence_errors(hij) assert all(max_err < max_rtol) and all(avg_err < avg_rtol), \ f"TT projection not transverse for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" max_err, avg_err = get_trace_errors(hij_h) assert max_err < max_rtol and avg_err < avg_rtol, \ f"TT projected tensor isn't traceless for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" plus = make_data(queue, fft).astype(cdtype) minus = make_data(queue, fft).astype(cdtype) project.pol_to_tensor(queue, plus, minus, hij) if isinstance(fft, gDFT): assert all(is_hermitian(hij[i]) for i in range(6)), \ f"pol->tensor is non-hermitian for {grid_shape=}, {h=}" max_err, avg_err = get_divergence_errors(hij) assert all(max_err < max_rtol) and all(avg_err < avg_rtol), \ f"pol->tensor not transverse for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" hij_h = hij.get() max_err, avg_err = get_trace_errors(hij_h) assert max_err < max_rtol and avg_err < avg_rtol, \ f"pol->tensor isn't traceless for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" hij_2 = cla.zeros_like(hij) project.transverse_traceless(queue, hij, hij_2) hij_h_2 = hij_2.get() max_err, avg_err = get_errs(hij_h, hij_h_2) assert max_err < max_rtol and avg_err < avg_rtol, \ f"pol->tensor != its own TT projection for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" plus1 = cla.zeros_like(plus) minus1 = cla.zeros_like(minus) project.tensor_to_pol(queue, plus1, minus1, hij) if isinstance(fft, gDFT): assert is_hermitian(plus1) and is_hermitian(minus1), \ f"polarizations aren't hermitian for {grid_shape=}, {h=}" max_err, avg_err = get_errs(plus1.get(), plus.get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"pol->tensor->pol (plus) is not identity for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" max_err, avg_err = get_errs(minus1.get(), minus.get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"pol->tensor->pol (minus) is not identity for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" project.tensor_to_pol(queue, hij[0], hij[1], hij) max_err, avg_err = get_errs(plus1.get(), hij[0].get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"in-place pol->tensor->pol (plus) not identity for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" max_err, avg_err = get_errs(minus1.get(), hij[1].get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"in-place pol->tensor->pol (minus) not identity for {grid_shape=}, {h=}" \ f": {max_err=}, {avg_err=}" if timing: from common import timer ntime = 10 t = timer(lambda: project.transverse_traceless(queue, hij), ntime=ntime) print(f"TT projection took {t:.3f} ms for {grid_shape=}") t = timer(lambda: project.pol_to_tensor(queue, plus, minus, hij), ntime=ntime) print(f"pol->tensor took {t:.3f} ms for {grid_shape=}") t = timer(lambda: project.tensor_to_pol(queue, plus, minus, hij), ntime=ntime) print(f"tensor->pol took {t:.3f} ms for {grid_shape=}")
def test_stencil(ctx_factory, grid_shape, proc_shape, dtype, stream, h=1, timing=False): if ctx_factory: ctx = ctx_factory() else: ctx = ps.choose_device_and_make_context() queue = cl.CommandQueue(ctx) rank_shape = tuple(Ni // pi for Ni, pi in zip(grid_shape, proc_shape)) from pymbolic import var x = var("x") y = var("y") i, j, k = var("i"), var("j"), var("k") map_dict = {} map_dict[y[i, j, k]] = (x[i + h + h, j + h, k + h] + x[i + h, j + h + h, k + h] + x[i + h, j + h, k + h + h] + x[i - h + h, j + h, k + h] + x[i + h, j - h + h, k + h] + x[i + h, j + h, k - h + h]) if stream: try: stencil_map = ps.StreamingStencil(map_dict, prefetch_args=["x"], halo_shape=h) except: # noqa pytest.skip("StreamingStencil unavailable") else: stencil_map = ps.Stencil(map_dict, h, prefetch_args=["x"]) x = clr.rand(queue, tuple(ni + 2 * h for ni in rank_shape), dtype) y = clr.rand(queue, rank_shape, dtype) x_h = x.get() y_true = (x_h[2 * h:, h:-h, h:-h] + x_h[h:-h, 2 * h:, h:-h] + x_h[h:-h, h:-h, 2 * h:] + x_h[:-2 * h, h:-h, h:-h] + x_h[h:-h, :-2 * h, h:-h] + x_h[h:-h, h:-h, :-2 * h]) stencil_map(queue, x=x, y=y) max_rtol = 5e-14 if dtype == np.float64 else 1e-5 avg_rtol = 5e-14 if dtype == np.float64 else 1e-5 max_err, avg_err = get_errs(y_true, y.get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"y innaccurate for {grid_shape=}, {h=}, {proc_shape=}" \ f": {max_err=}, {avg_err=}" if timing: from common import timer t = timer(lambda: stencil_map(queue, x=x, y=y)[0]) print( f"stencil took {t:.3f} ms for {grid_shape=}, {h=}, {proc_shape=}") bandwidth = (x.nbytes + y.nbytes) / 1024**3 / t * 1000 print(f"Bandwidth = {bandwidth} GB/s")
def test_dft(ctx_factory, grid_shape, proc_shape, dtype, use_fftw, timing=False): if not use_fftw and np.product(proc_shape) > 1: pytest.skip("Must use mpi4py-fft on more than one rank.") ctx = ctx_factory() queue = cl.CommandQueue(ctx) h = 1 mpi = ps.DomainDecomposition(proc_shape, h, grid_shape=grid_shape) mpi0 = ps.DomainDecomposition(proc_shape, 0, grid_shape=grid_shape) rank_shape, _ = mpi.get_rank_shape_start(grid_shape) fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype, use_fftw=use_fftw) grid_size = np.product(grid_shape) rdtype = fft.rdtype if fft.is_real: np_dft = np.fft.rfftn np_idft = np.fft.irfftn else: np_dft = np.fft.fftn np_idft = np.fft.ifftn rtol = 1e-11 if dtype in ("float64", "complex128") else 2e-3 rng = clr.ThreefryGenerator(ctx, seed=12321 * (mpi.rank + 1)) fx = rng.uniform(queue, rank_shape, rdtype) + 1e-2 if not fft.is_real: fx = fx + 1j * rng.uniform(queue, rank_shape, rdtype) fx = fx.get() fk = fft.dft(fx) if isinstance(fk, cla.Array): fk = fk.get() fk, _fk = fk.copy(), fk # hang on to one that fftw won't overwrite fx2 = fft.idft(_fk) if isinstance(fx2, cla.Array): fx2 = fx2.get() fx_glb = np.empty(shape=grid_shape, dtype=dtype) for root in range(mpi.nranks): mpi0.gather_array(queue, fx, fx_glb, root=root) fk_glb_np = np.ascontiguousarray(np_dft(fx_glb)) fx2_glb_np = np.ascontiguousarray(np_idft(fk_glb_np)) if use_fftw: fk_np = fk_glb_np[fft.fft.local_slice(True)] fx2_np = fx2_glb_np[fft.fft.local_slice(False)] else: fk_np = fk_glb_np fx2_np = fx2_glb_np max_err, avg_err = get_errs(fx, fx2 / grid_size) assert max_err < rtol, \ f"IDFT(DFT(f)) != f for {grid_shape=}, {proc_shape=}: {max_err=}, {avg_err=}" max_err, avg_err = get_errs(fk_np, fk) assert max_err < rtol, \ f"DFT disagrees with numpy for {grid_shape=}, {proc_shape=}:"\ f" {max_err=}, {avg_err=}" max_err, avg_err = get_errs(fx2_np, fx2 / grid_size) assert max_err < rtol, \ f"IDFT disagrees with numpy for {grid_shape=}, {proc_shape=}:"\ f" {max_err=}, {avg_err=}" fx_cl = cla.empty(queue, rank_shape, dtype) pencil_shape = tuple(ni + 2 * h for ni in rank_shape) fx_cl_halo = cla.empty(queue, pencil_shape, dtype) fx_np = np.empty(rank_shape, dtype) fx_np_halo = np.empty(pencil_shape, dtype) fk_cl = cla.empty(queue, fft.shape(True), fft.fk.dtype) fk_np = np.empty(fft.shape(True), fft.fk.dtype) # FIXME: check that these actually produce the correct result fx_types = { "cl": fx_cl, "cl halo": fx_cl_halo, "np": fx_np, "np halo": fx_np_halo, "None": None } fk_types = {"cl": fk_cl, "np": fk_np, "None": None} # run all of these to ensure no runtime errors even if no timing ntime = 20 if timing else 1 from common import timer if mpi.rank == 0: print(f"N = {grid_shape}, ", "complex" if np.dtype(dtype).kind == "c" else "real") from itertools import product for (a, input_), (b, output) in product(fx_types.items(), fk_types.items()): t = timer(lambda: fft.dft(input_, output), ntime=ntime) if mpi.rank == 0: print(f"dft({a}, {b}) took {t:.3f} ms") for (a, input_), (b, output) in product(fk_types.items(), fx_types.items()): t = timer(lambda: fft.idft(input_, output), ntime=ntime) if mpi.rank == 0: print(f"idft({a}, {b}) took {t:.3f} ms")
def test_scalar_energy(ctx_factory, grid_shape, proc_shape, h, dtype, timing=False): if ctx_factory: ctx = ctx_factory() else: ctx = ps.choose_device_and_make_context() queue = cl.CommandQueue(ctx) mpi = ps.DomainDecomposition(proc_shape, h, grid_shape=grid_shape) rank_shape, _ = mpi.get_rank_shape_start(grid_shape) grid_size = np.product(grid_shape) nscalars = 2 def potential(f): phi, chi = f[0], f[1] return 1 / 2 * phi**2 + 1 / 2 * chi**2 + 1 / 2 * phi**2 * chi**2 scalar_sector = ps.ScalarSector(nscalars, potential=potential) scalar_energy = ps.Reduction(mpi, scalar_sector, rank_shape=rank_shape, grid_size=grid_size, halo_shape=h) pencil_shape = tuple(ni + 2 * h for ni in rank_shape) f = clr.rand(queue, (nscalars, ) + pencil_shape, dtype) dfdt = clr.rand(queue, (nscalars, ) + pencil_shape, dtype) lap = clr.rand(queue, (nscalars, ) + rank_shape, dtype) energy = scalar_energy(queue, f=f, dfdt=dfdt, lap_f=lap, a=np.array(1.)) kin_test = [] grad_test = [] for fld in range(nscalars): df_h = dfdt[fld].get() rank_sum = np.sum(df_h[h:-h, h:-h, h:-h]**2) kin_test.append(1 / 2 * mpi.allreduce(rank_sum) / grid_size) f_h = f[fld].get() lap_h = lap[fld].get() rank_sum = np.sum(-f_h[h:-h, h:-h, h:-h] * lap_h) grad_test.append(1 / 2 * mpi.allreduce(rank_sum) / grid_size) energy_test = {} energy_test["kinetic"] = np.array(kin_test) energy_test["gradient"] = np.array(grad_test) phi = f[0].get()[h:-h, h:-h, h:-h] chi = f[1].get()[h:-h, h:-h, h:-h] pot_rank = np.sum(potential([phi, chi])) energy_test["potential"] = np.array(mpi.allreduce(pot_rank) / grid_size) max_rtol = 1e-14 if dtype == np.float64 else 1e-5 avg_rtol = 1e-14 if dtype == np.float64 else 1e-5 for key, value in energy.items(): max_err, avg_err = get_errs(value, energy_test[key]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"{key} inaccurate for {nscalars=}, {grid_shape=}, {proc_shape=}" \ f": {max_err=}, {avg_err=}" if timing: from common import timer t = timer(lambda: scalar_energy( queue, a=np.array(1.), f=f, dfdt=dfdt, lap_f=lap)) if mpi.rank == 0: print(f"scalar energy took {t:.3f} " f"ms for {nscalars=}, {grid_shape=}, {proc_shape=}")
def test_spectra(ctx_factory, grid_shape, proc_shape, dtype, L, timing=False): ctx = ctx_factory() 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) fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype) L = L or (3, 5, 7) dk = tuple(2 * np.pi / Li for Li in L) cdtype = fft.cdtype spec = ps.PowerSpectra(mpi, fft, dk, np.product(L), bin_width=min(dk) + .001) # FIXME: bin_width=min(dk) sometimes disagrees to O(.1%) with numpy... assert int(np.sum(spec.bin_counts)) == np.product(grid_shape), \ "bin counts don't sum to total number of points/modes" k_power = 2. fk = make_data(*fft.shape(True)).astype(cdtype) fk_d = cla.to_device(queue, fk) spectrum = spec.bin_power(fk_d, k_power=k_power) bins = np.arange(-.5, spec.num_bins + .5) * spec.bin_width sub_k = list(x.get() for x in fft.sub_k.values()) kvecs = np.meshgrid(*sub_k, indexing="ij", sparse=False) kmags = np.sqrt(sum((dki * ki)**2 for dki, ki in zip(dk, kvecs))) if fft.is_real: counts = 2. * np.ones_like(kmags) counts[kvecs[2] == 0] = 1 counts[kvecs[2] == grid_shape[-1] // 2] = 1 else: counts = 1. * np.ones_like(kmags) if np.dtype(dtype) in (np.dtype("float64"), np.dtype("complex128")): max_rtol = 1e-8 avg_rtol = 1e-11 else: max_rtol = 2e-2 avg_rtol = 2e-4 bin_counts2 = spec.bin_power(np.ones_like(fk), queue=queue, k_power=0) max_err, avg_err = get_errs(bin_counts2, np.ones_like(bin_counts2)) assert max_err < max_rtol and avg_err < avg_rtol, \ f"bin counting disagrees between PowerSpectra and np.histogram" \ f" for {grid_shape=}: {max_err=}, {avg_err=}" hist = np.histogram(kmags, bins=bins, weights=np.abs(fk)**2 * counts * kmags**k_power)[0] hist = mpi.allreduce(hist) / spec.bin_counts # skip the Nyquist mode and the zero mode max_err, avg_err = get_errs(spectrum[1:-2], hist[1:-2]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"power spectrum inaccurate for {grid_shape=}: {max_err=}, {avg_err=}" if timing: from common import timer t = timer(lambda: spec.bin_power(fk_d, k_power=k_power)) print(f"power spectrum took {t:.3f} ms for {grid_shape=}, {dtype=}")
def test_pol_spectra(ctx_factory, grid_shape, proc_shape, dtype, timing=False): ctx = ctx_factory() if np.dtype(dtype).kind != "f": dtype = "float64" 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) fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype) L = (10, 8, 7) dk = tuple(2 * np.pi / Li for Li in L) dx = tuple(Li / Ni for Li, Ni in zip(L, grid_shape)) cdtype = fft.cdtype spec = ps.PowerSpectra(mpi, fft, dk, np.product(L)) k_power = 2. fk = make_data(*fft.shape(True)).astype(cdtype) fk = make_hermitian(fk, fft).astype(cdtype) plus = cla.to_device(queue, fk) fk = make_data(*fft.shape(True)).astype(cdtype) fk = make_hermitian(fk, fft).astype(cdtype) minus = cla.to_device(queue, fk) plus_ps_1 = spec.bin_power(plus, queue=queue, k_power=k_power) minus_ps_1 = spec.bin_power(minus, queue=queue, k_power=k_power) project = ps.Projector(fft, h, dk, dx) vector = cla.empty(queue, (3, ) + fft.shape(True), cdtype) project.pol_to_vec(queue, plus, minus, vector) project.vec_to_pol(queue, plus, minus, vector) plus_ps_2 = spec.bin_power(plus, k_power=k_power) minus_ps_2 = spec.bin_power(minus, k_power=k_power) max_rtol = 1e-8 if dtype == np.float64 else 1e-2 avg_rtol = 1e-11 if dtype == np.float64 else 1e-4 max_err, avg_err = get_errs(plus_ps_1[1:-2], plus_ps_2[1:-2]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"plus power spectrum inaccurate for {grid_shape=}: {max_err=}, {avg_err=}" max_err, avg_err = get_errs(minus_ps_1[1:-2], minus_ps_2[1:-2]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"minus power spectrum inaccurate for {grid_shape=}: {max_err=}, {avg_err=}" vec_sum = sum( spec.bin_power(vector[mu], k_power=k_power) for mu in range(3)) pol_sum = plus_ps_1 + minus_ps_1 max_err, avg_err = get_errs(vec_sum[1:-2], pol_sum[1:-2]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"polarization power spectrum inaccurate for {grid_shape=}" \ f": {max_err=}, {avg_err=}" # reset for mu in range(3): fk = make_data(*fft.shape(True)).astype(cdtype) fk = make_hermitian(fk, fft).astype(cdtype) vector[mu].set(fk) long = cla.zeros_like(plus) project.decompose_vector(queue, vector, plus, minus, long, times_abs_k=True) plus_ps = spec.bin_power(plus, k_power=k_power) minus_ps = spec.bin_power(minus, k_power=k_power) long_ps = spec.bin_power(long, k_power=k_power) vec_sum = sum( spec.bin_power(vector[mu], k_power=k_power) for mu in range(3)) dec_sum = plus_ps + minus_ps + long_ps max_err, avg_err = get_errs(vec_sum[1:-2], dec_sum[1:-2]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"decomp power spectrum inaccurate for {grid_shape=}: {max_err=}, {avg_err=}" hij = cl.clrandom.rand(queue, (6, ) + rank_shape, dtype) gw_spec = spec.gw(hij, project, 1.3) gw_pol_spec = spec.gw_polarization(hij, project, 1.3) max_rtol = 1e-14 if dtype == np.float64 else 1e-2 avg_rtol = 1e-11 if dtype == np.float64 else 1e-4 pol_sum = gw_pol_spec[0] + gw_pol_spec[1] max_err, avg_err = get_errs(gw_spec[1:-2], pol_sum[1:-2]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"gw pol don't add up to gw for {grid_shape=}: {max_err=}, {avg_err=}"
def test_gradient_laplacian(ctx_factory, grid_shape, proc_shape, h, dtype, stream, timing=False): if h == 0 and stream is True: pytest.skip("no streaming spectral") ctx = ctx_factory() queue = cl.CommandQueue(ctx) mpi = ps.DomainDecomposition(proc_shape, h, grid_shape=grid_shape) rank_shape, start = mpi.get_rank_shape_start(grid_shape) L = (3, 5, 7) dx = tuple(Li / Ni for Li, Ni in zip(L, grid_shape)) dk = tuple(2 * np.pi / Li for Li in L) if h == 0: def get_evals_1(k, dx): return k def get_evals_2(k, dx): return -k**2 fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype) derivs = ps.SpectralCollocator(fft, dk) else: from pystella.derivs import FirstCenteredDifference, SecondCenteredDifference get_evals_1 = FirstCenteredDifference(h).get_eigenvalues get_evals_2 = SecondCenteredDifference(h).get_eigenvalues if stream: try: derivs = ps.FiniteDifferencer(mpi, h, dx, rank_shape=rank_shape, stream=stream) except: # noqa pytest.skip("StreamingStencil unavailable") else: derivs = ps.FiniteDifferencer(mpi, h, dx, rank_shape=rank_shape) pencil_shape = tuple(ni + 2 * h for ni in rank_shape) # set up test data fx_h = np.empty(pencil_shape, dtype) kvec = np.array(dk) * np.array([-5, 4, -3]).astype(dtype) xvec = np.meshgrid(*[ dxi * np.arange(si, si + ni) for dxi, si, ni in zip(dx, start, rank_shape) ], indexing="ij") phases = sum(ki * xi for ki, xi in zip(kvec, xvec)) if h > 0: fx_h[h:-h, h:-h, h:-h] = np.sin(phases) else: fx_h[:] = np.sin(phases) fx_cos = np.cos(phases) fx = cla.to_device(queue, fx_h) lap = cla.empty(queue, rank_shape, dtype) grd = cla.empty(queue, (3, ) + rank_shape, dtype) derivs(queue, fx=fx, lap=lap, grd=grd) eff_kmag_sq = sum( get_evals_2(kvec_i, dxi) for dxi, kvec_i in zip(dx, kvec)) lap_true = eff_kmag_sq * np.sin(phases) max_rtol = 1e-9 if dtype == np.float64 else 3e-4 avg_rtol = 1e-11 if dtype == np.float64 else 5e-5 # filter small values dominated by round-off error mask = np.abs(lap_true) > 1e-11 max_err, avg_err = get_errs(lap_true[mask], lap.get()[mask]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"lap inaccurate for {h=}, {grid_shape=}, {proc_shape=}:" \ f" {max_err=}, {avg_err=}" for i in range(3): eff_k = get_evals_1(kvec[i], dx[i]) pdi_true = eff_k * fx_cos # filter small values dominated by round-off error mask = np.abs(pdi_true) > 1e-11 max_err, avg_err = get_errs(pdi_true[mask], grd[i].get()[mask]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"pd{i} inaccurate for {h=}, {grid_shape=}, {proc_shape=}:" \ f" {max_err=}, {avg_err=}" vec = cla.empty(queue, (3, ) + pencil_shape, dtype) for mu in range(3): vec[mu] = fx div = cla.empty(queue, rank_shape, dtype) derivs.divergence(queue, vec, div) div_true = sum(grd[i] for i in range(3)).get() # filter small values dominated by round-off error mask = np.abs(div_true) > 1e-11 max_err, avg_err = get_errs(div_true[mask], div.get()[mask]) assert max_err < max_rtol and avg_err < avg_rtol, \ f"div inaccurate for {h=}, {grid_shape=}, {proc_shape=}:" \ f" {max_err=}, {avg_err=}" if timing: from common import timer base_args = dict(queue=queue, fx=fx) div_args = dict(queue=queue, vec=vec, div=div) if h == 0: import pyopencl.tools as clt pool = clt.MemoryPool(clt.ImmediateAllocator(queue)) base_args["allocator"] = pool div_args["allocator"] = pool times = {} times["gradient and laplacian"] = timer( lambda: derivs(lap=lap, grd=grd, **base_args)) times["gradient"] = timer(lambda: derivs(grd=grd, **base_args)) times["laplacian"] = timer(lambda: derivs(lap=lap, **base_args)) times["pdx"] = timer(lambda: derivs(pdx=grd[0], **base_args)) times["pdy"] = timer(lambda: derivs(pdy=grd[1], **base_args)) times["pdz"] = timer(lambda: derivs(pdz=grd[2], **base_args)) times["divergence"] = timer(lambda: derivs.divergence(**div_args)) if mpi.rank == 0: print(f"{grid_shape=}, {h=}, {proc_shape=}") for key, val in times.items(): print(f"{key} took {val:.3f} ms")
def test_field_histogram(ctx_factory, grid_shape, proc_shape, dtype, 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) pencil_shape = tuple(Ni + 2 * h for Ni in rank_shape) num_bins = 432 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 hist = ps.FieldHistogrammer(mpi, num_bins, dtype, rank_shape=rank_shape, halo_shape=h) rng = clr.ThreefryGenerator(ctx, seed=12321) fx = rng.uniform(queue, (2, 2)+pencil_shape, dtype, a=-1.2, b=3.) fx_h = fx.get()[..., h:-h, h:-h, h:-h] result = hist(fx) outer_shape = fx.shape[:-3] from itertools import product slices = list(product(*[range(n) for n in outer_shape])) for slc in slices: res = result["linear"][slc] np_res = np.histogram(fx_h[slc], bins=result["linear_bins"][slc])[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"linear Histogrammer inaccurate for grid_shape={grid_shape}" \ f": {max_err=}, {avg_err=}" res = result["log"][slc] bins = result["log_bins"][slc] # avoid FPA comparison issues # numpy sometimes doesn't count the actual maximum/minimum eps = 1e-14 if np.dtype(dtype) == np.dtype("float64") else 1e-4 bins[0] *= (1 - eps) bins[-1] *= (1 + eps) np_res = np.histogram(np.abs(fx_h[slc]), bins=bins)[0] np_res = mpi.allreduce(np_res) norm = np.maximum(np.abs(res), np.abs(np_res)) norm[norm == 0.] = 1. max_err, avg_err = get_errs(res, np_res) assert max_err < max_rtol and avg_err < avg_rtol, \ f"log Histogrammer inaccurate for grid_shape={grid_shape}" \ f": {max_err=}, {avg_err=}" if timing: from common import timer t = timer(lambda: hist(fx[0, 0])) print(f"field histogram took {t:.3f} ms for {grid_shape=}, {dtype=}")
def test_elementwise(ctx_factory, grid_shape, proc_shape, dtype, timing=False): if ctx_factory: ctx = ctx_factory() else: ctx = ps.choose_device_and_make_context() queue = cl.CommandQueue(ctx) rank_shape = tuple(Ni // pi for Ni, pi in zip(grid_shape, proc_shape)) from pymbolic import var a = var("a") b = var("b") from pystella.field import Field x = Field("x") y = Field("y") z = Field("z") tmp_dict = {a[0]: x + 2, a[1]: 2 + x * y, b: x + y / 2} map_dict = {x: a[0] * y**2 * x + a[1] * b, z: z + a[1] * b} single_insn = {x: y + z} ew_map = ps.ElementWiseMap(map_dict, tmp_instructions=tmp_dict) x = clr.rand(queue, rank_shape, dtype=dtype) y = clr.rand(queue, rank_shape, dtype=dtype) z = clr.rand(queue, rank_shape, dtype=dtype) a0 = x + 2 a1 = 2 + x * y b = x + y / 2 x_true = a0 * y**2 * x + a1 * b z_true = z + a1 * b ew_map(queue, x=x, y=y, z=z) max_rtol = 5e-14 if dtype == np.float64 else 1e-5 avg_rtol = 5e-14 if dtype == np.float64 else 1e-5 max_err, avg_err = get_errs(x_true.get(), x.get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"x innaccurate for {grid_shape=}, {proc_shape=}: {max_err=}, {avg_err=}" max_err, avg_err = get_errs(z_true.get(), z.get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"z innaccurate for {grid_shape=}, {proc_shape=}: {max_err=}, {avg_err=}" # test success of single instruction ew_map_single = ps.ElementWiseMap(single_insn) ew_map_single(queue, x=x, y=y, z=z) x_true = y + z max_err, avg_err = get_errs(x_true.get(), x.get()) assert max_err < max_rtol and avg_err < avg_rtol, \ f"x innaccurate for {grid_shape=}, {proc_shape=}: {max_err=}, {avg_err=}" if timing: from common import timer t = timer(lambda: ew_map(queue, x=x, y=y, z=z)[0]) print( f"elementwise map took {t:.3f} ms for {grid_shape=}, {proc_shape=}" ) bandwidth = 5 * x.nbytes / 1024**3 / t * 1000 print(f"Bandwidth = {bandwidth:.1f} GB/s")