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")