def test_c2rr2c():
field = pm.generate_whitenoise(seed=1234, type='real')
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(real='r', complex='c')
code.c2r(complex='c', real='r')
check_grad(code, 'r', 'r', init={'r': field}, eps=1e-4, rtol=1e-8)
def test_c2rr2c():
field = pm.generate_whitenoise(seed=1234, mode='real')
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(real='r', complex='c')
code.c2r(complex='c', real='r')
check_grad(code, 'r', 'r', init={'r': field}, eps=1e-4, rtol=1e-8)
def test_compute():
def transfer(k): return 2.0
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(real='r', complex='c')
code.transfer(complex='c', tf=transfer)
code.c2r(complex='c', real='r')
code.to_scalar(x='r', y='sum')
field = pm.generate_whitenoise(seed=1234).c2r()
norm = code.compute('sum', init={'r': field})
assert_allclose(norm, field.cnorm() * 4)
def test_transfer_real():
def transfer(k):
return k[0]
field = pm.generate_whitenoise(seed=1234, mode='real')
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(complex='c', real='r')
code.transfer(complex='c', tf=transfer)
code.c2r(complex='c', real='r')
check_grad(code, 'r', 'r', init={'r': field}, eps=1e-4, rtol=1e-8)
def test_compute():
def transfer(k): return 2.0
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(real='r', complex='c')
code.transfer(complex='c', tf=transfer)
code.c2r(complex='c', real='r')
code.to_scalar(x='r', y='sum')
field = pm.generate_whitenoise(seed=1234).c2r()
norm = code.compute('sum', init={'r': field})
assert_allclose(norm, field.cnorm() * 4)
def test_transfer_real():
def transfer(k):
return k[0]
field = pm.generate_whitenoise(seed=1234, type='real')
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(complex='c', real='r')
code.transfer(complex='c', tf=transfer)
code.c2r(complex='c', real='r')
check_grad(code, 'r', 'r', init={'r': field}, eps=1e-4, 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)
def test_generate_2nd_order_source():
engine = FastPMEngine(pm)
code = CodeSegment(engine)
code.r2c(real='source', complex='source_k')
code.generate_2nd_order_source(source_k='source_k', source2_k='source2_k')
code.c2r(complex='source2_k', real='source2')
field = pm.generate_whitenoise(seed=1234).c2r()
# from fastpm.operators import lpt1, lpt2source
# dlin2_k = lpt2source(field.r2c())
source2_k = code.compute('source2', init={'source': field}).r2c()
# assert_allclose(dlin2_k[...], source2_k[...], atol=1e-7, rtol=1e-4)
check_grad(code,
'source2',
'source',
init={'source': field},
eps=1e-4,
rtol=1e-8)
def test_vjp():
def transfer(k):
k2 = sum(ki **2 for ki in k)
k2[k2 == 0] = 1.0
# return 1 / k2
return 2.0
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(real='r', complex='c')
code.transfer(complex='c', tf=transfer)
code.c2r(complex='c', real='r')
code.multiply(x1='r', x2=Literal(0.1), y='r')
code.to_scalar(x='r', y='sum')
field = pm.generate_whitenoise(seed=1234).c2r()
norm, tape = code.compute('sum', init={'r': field}, return_tape=True)
assert_allclose(norm, field.cnorm() * 4 * 0.1 ** 2)
vjp = tape.get_vjp()
_r = vjp.compute('_r', init={'_sum': 1.0})
assert_allclose(_r, field * 4 * 2 * 0.1 * 0.1)
def test_vjp():
def transfer(k):
k2 = sum(ki **2 for ki in k)
k2[k2 == 0] = 1.0
# return 1 / k2
return 2.0
engine = ParticleMeshEngine(pm)
code = CodeSegment(engine)
code.r2c(real='r', complex='c')
code.transfer(complex='c', tf=transfer)
code.c2r(complex='c', real='r')
code.multiply(x1='r', x2=Literal(0.1), y='r')
code.to_scalar(x='r', y='sum')
field = pm.generate_whitenoise(seed=1234).c2r()
norm, tape = code.compute('sum', init={'r': field}, return_tape=True)
assert_allclose(norm, field.cnorm() * 4 * 0.1 ** 2)
vjp = tape.get_vjp()
_r = vjp.compute('_r', init={'_sum': 1.0})
assert_allclose(_r, field * 4 * 2 * 0.1 * 0.1)
