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_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_decic(): pm = ParticleMesh(BoxSize=8.0, Nmesh=(8, 8, 8), dtype='f8') engine = FastPMEngine(pm) code = CodeSegment(engine) code.r2c(complex='d_k', real='d_r') code.de_cic(deconvolved='decic', d_k='d_k') field = engine.pm.generate_whitenoise(seed=1234).c2r() eps = field.cnorm()**0.5 * 1e-3 check_grad(code, 'decic', 'd_r', init={'d_r': field}, eps=eps, 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_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_solve_linear_displacement(): engine = FastPMEngine(pm) code = CodeSegment(engine) code.r2c(real='source', complex='dlinear_k') code.solve_linear_displacement(source_k='dlinear_k', s='s') field = pm.generate_whitenoise(seed=1234, mode='real') eps = field.cnorm()**0.5 * 1e-3 check_grad(code, 's', 'source', init={'source': field}, eps=eps, rtol=1e-8) # from fastpm.operators import lpt1, lpt2source dlin_k, s = code.compute(['dlinear_k', 's'], init={'source': field})
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)