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