def test_generate_WKB(ctx_factory, grid_shape, proc_shape, dtype, random, 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) fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype) L = (10,)*3 volume = np.product(L) dk = tuple(2 * np.pi / Li for Li in L) modes = ps.RayleighGenerator(ctx, fft, dk, volume) # only checking that this call is successful fk, dfk = modes.generate_WKB(queue, random=random) if timing: ntime = 10 from common import timer t = timer(lambda: modes.generate_WKB(queue, random=random), ntime=ntime) print(f"{random=} set_modes took {t:.3f} ms for {grid_shape=}")
def test_generate(ctx_factory, grid_shape, proc_shape, dtype, random, 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) fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype) num_bins = int(sum(Ni**2 for Ni in grid_shape)**.5 / 2 + .5) + 1 L = (10,)*3 volume = np.product(L) dk = tuple(2 * np.pi / Li for Li in L) spectra = ps.PowerSpectra(mpi, fft, dk, volume) modes = ps.RayleighGenerator(ctx, fft, dk, volume, seed=5123) kbins = min(dk) * np.arange(0, num_bins) test_norm = 1 / 2 / np.pi**2 / np.product(grid_shape)**2 for exp in [-1, -2, -3]: def power(k): return k**exp fk = modes.generate(queue, random=random, norm=1, field_ps=power) spectrum = spectra.norm * spectra.bin_power(fk, queue=queue, k_power=3)[1:-1] true_spectrum = test_norm * kbins[1:-1]**3 * power(kbins[1:-1]) err = np.abs(1 - spectrum / true_spectrum) tol = .1 if num_bins < 64 else .3 assert (np.max(err[num_bins//3:-num_bins//3]) < tol and np.average(err[1:]) < tol), \ f"init power spectrum incorrect for {random=}, k**{exp}" if random: fx = fft.idft(cla.to_device(queue, fk)).real if isinstance(fx, cla.Array): fx = fx.get() grid_size = np.product(grid_shape) avg = mpi.allreduce(np.sum(fx)) / grid_size var = mpi.allreduce(np.sum(fx**2)) / grid_size - avg**2 skew = mpi.allreduce(np.sum(fx**3)) / grid_size - 3 * avg * var - avg**3 skew /= var**1.5 assert skew < tol, \ f"init power spectrum has large skewness for k**{exp}" if timing: ntime = 10 from common import timer t = timer(lambda: modes.generate(queue, random=random), ntime=ntime) print(f"{random=} set_modes took {t:.3f} ms for {grid_shape=}")
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_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_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_spectral_poisson(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) fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype) 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_2(k, dx): return - k**2 derivs = ps.SpectralCollocator(fft, dk) else: from pystella.derivs import SecondCenteredDifference get_evals_2 = SecondCenteredDifference(h).get_eigenvalues derivs = ps.FiniteDifferencer(mpi, h, dx, stream=False) solver = ps.SpectralPoissonSolver(fft, dk, dx, get_evals_2) pencil_shape = tuple(ni + 2*h for ni in rank_shape) statistics = ps.FieldStatistics(mpi, 0, rank_shape=rank_shape, grid_size=np.product(grid_shape)) fx = cla.empty(queue, pencil_shape, dtype) rho = clr.rand(queue, rank_shape, dtype) rho -= statistics(rho)["mean"] lap = cla.empty(queue, rank_shape, dtype) rho_h = rho.get() for m_squared in (0, 1.2, 19.2): solver(queue, fx, rho, m_squared=m_squared) fx_h = fx.get() if h > 0: fx_h = fx_h[h:-h, h:-h, h:-h] derivs(queue, fx=fx, lap=lap) diff = np.fabs(lap.get() - rho_h - m_squared * fx_h) max_err = np.max(diff) / cla.max(clm.fabs(rho)) avg_err = np.sum(diff) / cla.sum(clm.fabs(rho)) max_rtol = 1e-12 if dtype == np.float64 else 1e-4 avg_rtol = 1e-13 if dtype == np.float64 else 1e-5 assert max_err < max_rtol and avg_err < avg_rtol, \ f"solution inaccurate for {h=}, {grid_shape=}, {proc_shape=}" if timing: from common import timer time = timer(lambda: solver(queue, fx, rho, m_squared=m_squared), ntime=10) if mpi.rank == 0: print(f"poisson took {time:.3f} ms for {grid_shape=}, {proc_shape=}")
def test_relax(ctx_factory, grid_shape, proc_shape, h, dtype, Solver, timing=False): if min(grid_shape) < 128: pytest.skip("test_relax needs larger grids, for now") 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)) mpi = ps.DomainDecomposition(proc_shape, h, rank_shape) L = 10 dx = L / grid_shape[0] dk = 2 * np.pi / L fft = ps.DFT(mpi, ctx, queue, grid_shape, dtype) spectra = ps.PowerSpectra(mpi, fft, (dk, ) * 3, L**3) statistics = ps.FieldStatistics(mpi, h, rank_shape=rank_shape, grid_size=np.product(grid_shape)) def get_laplacian(f): from pystella.derivs import _lap_coefs, centered_diff lap_coefs = _lap_coefs[h] from pymbolic import var return sum([ centered_diff(f, lap_coefs, direction=mu, order=2) for mu in range(1, 4) ]) / var("dx")**2 test_problems = {} from pystella import Field f = Field("f", offset="h") rho = Field("rho", offset="h") test_problems[f] = (get_laplacian(f), rho) f = Field("f2", offset="h") rho = Field("rho2", offset="h") test_problems[f] = (get_laplacian(f) - f, rho) solver = Solver(mpi, queue, test_problems, halo_shape=h, dtype=dtype, fixed_parameters=dict(omega=1 / 2)) def zero_mean_array(): f0 = clr.rand(queue, grid_shape, dtype) f = clr.rand(queue, tuple(ni + 2 * h for ni in rank_shape), dtype) mpi.scatter_array(queue, f0, f, root=0) avg = statistics(f)["mean"] f = f - avg mpi.share_halos(queue, f) return f f = zero_mean_array() rho = zero_mean_array() tmp = cla.zeros_like(f) f2 = zero_mean_array() rho2 = zero_mean_array() tmp2 = cla.zeros_like(f) num_iterations = 1000 errors = {"f": [], "f2": []} first_mode_zeroed = {"f": [], "f2": []} for i in range(0, num_iterations, 2): solver(mpi, queue, iterations=2, dx=np.array(dx), f=f, tmp_f=tmp, rho=rho, f2=f2, tmp_f2=tmp2, rho2=rho2) err = solver.get_error(queue, f=f, r_f=tmp, rho=rho, f2=f2, r_f2=tmp2, rho2=rho2, dx=np.array(dx)) for k, v in err.items(): errors[k].append(v) for key, resid in zip(["f", "f2"], [tmp, tmp2]): spectrum = spectra(resid, k_power=0) if mpi.rank == 0: max_amp = np.max(spectrum) first_zero = np.argmax(spectrum[1:] < 1e-30 * max_amp) first_mode_zeroed[key].append(first_zero) for k, errs in errors.items(): errs = np.array(errs) iters = np.arange(1, errs.shape[0] + 1) assert (errs[10:, 0] * iters[10:] / errs[0, 0] < 1.).all(), \ "relaxation not converging at least linearly for " \ f"{grid_shape=}, {h=}, {proc_shape=}" first_mode_zeroed = mpi.bcast(first_mode_zeroed, root=0) for k, x in first_mode_zeroed.items(): x = np.array(list(x))[2:] assert (x[1:] <= x[:-1]).all() and np.min(x) < np.max(x) / 5, \ f"relaxation not smoothing error {grid_shape=}, {h=}, {proc_shape=}"
mchi = 0. gsq = 2.5e-7 sigma = 0. lambda4 = 0. f0 = [.193 * mpl, 0] # units of mpl df0 = [-.142231 * mpl, 0] # units of mpl end_time = 1 end_scale_factor = 20 Stepper = ps.LowStorageRK54 gravitational_waves = True # whether to simulate gravitational waves ctx = ps.choose_device_and_make_context() queue = cl.CommandQueue(ctx) decomp = ps.DomainDecomposition(proc_shape, halo_shape, rank_shape) fft = ps.DFT(decomp, ctx, queue, grid_shape, dtype) if halo_shape == 0: derivs = ps.SpectralCollocator(fft, dk) else: derivs = ps.FiniteDifferencer(decomp, halo_shape, dx, rank_shape=rank_shape) def potential(f): phi, chi = f[0], f[1] unscaled = (mphi**2 / 2 * phi**2 + mchi**2 / 2 * chi**2 + gsq / 2 * phi**2 * chi**2 + sigma / 2 * phi * chi**2 + lambda4 / 4 * chi**4) return unscaled / mphi**2