def test_linear_vs_whitenoise():
pm = ParticleMesh(BoxSize=8.0, Nmesh=(4, 4), dtype='f8')
engine = FastPMEngine(pm)
code = CodeSegment(engine)
def pk(k):
return 0.0 * k + 1.0
code.create_linear_field(whitenoise='whitenoise',
powerspectrum=pk,
dlinear_k='dlinear_k')
code.create_whitenoise(whitenoise='whitenoise2',
powerspectrum=pk,
dlinear_k='dlinear_k')
def tf1(k):
k2 = sum(ki**2 for ki in k)
r = (pk(k2**0.5) / engine.pm.BoxSize.prod())**0.5
r[k2 == 0] = 1.0
return r
def tf2(k):
k2 = sum(ki**2 for ki in k)
r = (pk(k2**0.5) / engine.pm.BoxSize.prod())**-0.5
r[k2 == 0] = 1.0
return r
field = engine.pm.generate_whitenoise(seed=1234).c2r()
eps = field.cnorm()**0.5 * 1e-3
dlineark, whitenoise2 = code.compute(['dlinear_k', 'whitenoise2'],
init={'whitenoise': field})
dlineark2 = field.r2c().apply(lambda k, v: tf1(k) * v)
assert_allclose(dlineark, dlineark2)
assert_allclose(field, whitenoise2)
def test_force():
pm = ParticleMesh(BoxSize=8.0, Nmesh=(8, 8, 8), dtype='f8')
engine = FastPMEngine(pm)
code = CodeSegment(engine)
code.create_linear_field(whitenoise='whitenoise',
powerspectrum=pk,
dlinear_k='dlinear_k')
code.solve_lpt(pt=pt,
dlinear_k='dlinear_k',
aend=0.1,
s='s',
v='v',
s1='s1',
s2='s2')
field = engine.pm.generate_whitenoise(seed=1234).c2r()
s = code.compute('s', init={'whitenoise': field})
code = CodeSegment(engine)
code.force(s='s', force='force', force_factor=1.0)
eps = (pm.comm.allreduce(
(s**2).sum()) / pm.comm.allreduce(len(s)))**0.5 * 1e-3
s = s.clip(2 * eps * pm.BoxSize / pm.Nmesh,
(1 - 2 * eps) * pm.BoxSize / pm.Nmesh)
check_grad(code, 'force', 's', init={'s': s}, eps=eps, rtol=1e-8)
def test_solve_fastpm_linear_growth():
pm = ParticleMesh(BoxSize=1024.0, Nmesh=(128, 128, 128), dtype='f8')
engine = FastPMEngine(pm)
code = CodeSegment(engine)
code.create_linear_field(whitenoise='whitenoise',
powerspectrum=pk,
dlinear_k='dlinear_k')
code.solve_fastpm(pt=pt,
dlinear_k='dlinear_k',
asteps=[0.1, 0.5, 1.0],
s='s',
v='v',
s1='s1',
s2='s2')
code.get_x(s='s', x='x')
code.paint_simple(x='x', density='density')
field = pm.generate_whitenoise(seed=1234, unitary=True).c2r()
density, dlinear_k, s = code.compute(['density', 'dlinear_k', 's'],
init={'whitenoise': field})
density_k = density.r2c()
p_lin = FFTPower(FieldMesh(dlinear_k), mode='1d')
p_nonlin = FFTPower(FieldMesh(density), mode='1d')
# the simulation shall do a linear growth
t1 = abs((p_nonlin.power['power'] / p_lin.power['power'])**0.5)
assert_allclose(t1[1:4], 1.0, rtol=5e-2)
def test_solve_lpt():
pm = ParticleMesh(BoxSize=8.0, Nmesh=(4, 4, 4), dtype='f8')
engine = FastPMEngine(pm)
code = CodeSegment(engine)
code.create_linear_field(whitenoise='whitenoise',
powerspectrum=pk,
dlinear_k='dlinear_k')
code.solve_lpt(pt=pt,
dlinear_k='dlinear_k',
aend=0.1,
s='s',
v='v',
s1='s1',
s2='s2')
field = pm.generate_whitenoise(seed=1234).c2r()
s1, s2 = code.compute(['s1', 's2'], init={'whitenoise': field})
dlin_k = code.compute('dlinear_k', init={'whitenoise': field})
s1, tape = code.compute('s1', init={'whitenoise': field}, return_tape=True)
# from fastpm.operators import lpt1, lpt2source
# s1_truth = lpt1(dlin_k, engine.q, resampler='cic')
# dlin2_k = lpt2source(dlin_k)
# s2_truth = lpt1(dlin2_k, engine.q, resampler='cic')
# assert_allclose(s1, s1_truth, rtol=1e-4)
# assert_allclose(s2, s2_truth, rtol=1e-4)
eps = field.cnorm()**0.5 * 1e-4
check_grad(code,
's1',
'whitenoise',
init={'whitenoise': field},
eps=eps,
rtol=1e-8)
check_grad(code,
's2',
'whitenoise',
init={'whitenoise': field},
eps=eps,
rtol=1e-8)
check_grad(code,
's',
'whitenoise',
init={'whitenoise': field},
eps=eps,
rtol=1e-8)
check_grad(code,
'v',
'whitenoise',
init={'whitenoise': field},
eps=eps,
rtol=1e-8)
def test_solve_fastpm():
pm = ParticleMesh(BoxSize=8.0, Nmesh=(4, 4, 4), dtype='f8')
engine = FastPMEngine(pm, shift=0.5, B=1)
code = CodeSegment(engine)
code.create_linear_field(whitenoise='whitenoise',
powerspectrum=pk,
dlinear_k='dlinear_k')
code.solve_fastpm(pt=pt,
dlinear_k='dlinear_k',
asteps=[0.1, 1.0],
s='s',
v='v',
s1='s1',
s2='s2')
# code.solve_fastpm(pt=pt, dlinear_k='dlinear_k', asteps=[1.0], s='s', v='v', s1='s1', s2='s2')
code.get_x(s='s', x='x')
code.paint_simple(x='x', density='density')
field = pm.generate_whitenoise(seed=1234, unitary=True).c2r()
eps = field.cnorm()**0.5 * 1e-5
check_grad(code,
's1',
'whitenoise',
init={'whitenoise': field},
eps=eps,
rtol=1e-8)
check_grad(code,
's',
'whitenoise',
init={'whitenoise': field},
eps=eps,
rtol=1e-8)
check_grad(code,
'v',
'whitenoise',
init={'whitenoise': field},
eps=eps,
rtol=1e-8)
check_grad(code,
'density',
'whitenoise',
init={'whitenoise': field},
eps=eps,
rtol=1e-8)
def test_create_linear_field():
pm = ParticleMesh(BoxSize=8.0, Nmesh=(8, 8, 8), dtype='f8')
engine = FastPMEngine(pm)
code = CodeSegment(engine)
code.create_linear_field(whitenoise='whitenoise',
powerspectrum=pk,
dlinear_k='dlinear_k')
code.c2r(complex='dlinear_k', real='dlinear')
field = engine.pm.generate_whitenoise(seed=1234).c2r()
eps = field.cnorm()**0.5 * 1e-3
check_grad(code,
'dlinear',
'whitenoise',
init={'whitenoise': field},
eps=eps,
rtol=1e-8)
