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=}")
# compute hubble correction to scalar field effective mass
addot = expand.addot_friedmann_2(expand.a, energy["total"], energy["pressure"])
hubbleCorrection = -addot / expand.a
# effective masses of scalar fields
from pymbolic import var
from pymbolic.mapper.evaluator import evaluate_kw
fields = [var("f0")[i] for i in range(nscalars)]
d2Vd2f = [ps.diff(potential(fields), field, field) for field in fields]
eff_mass = [evaluate_kw(x, f0=f0) + hubbleCorrection for x in d2Vd2f]
modes = ps.RayleighGenerator(ctx,
fft,
dk,
volume,
seed=49279 * (decomp.rank + 1))
for fld in range(nscalars):
modes.init_WKB_fields(f[fld],
dfdt[fld],
norm=mphi**2,
omega_k=lambda k: np.sqrt(k**2 + eff_mass[fld]),
hubble=expand.hubble[0])
for i in range(nscalars):
f[i] += f0[i]
dfdt[i] += df0[i]
# re-initialize energy and expansion