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_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_total(): engine = ParticleMeshEngine(pm) code = CodeSegment(engine) numpy.random.seed(1234) code.total(x='s', y='l') #s = numpy.random.uniform(1, 2, size=engine.q.shape) #s = numpy.ones_like(engine.q) s = pm.generate_whitenoise(seed=1234, mode='real') l = code.compute(['l'], init={'s':s}) print(l) check_grad(code, 'l', 's', init={'s': s}, eps=1e-8, rtol=1e-4, toscalar=False)
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_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_features(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=(4, 4, 4), dtype='f8', comm=comm) engine = FastPMEngine(pm, shift=0.5, B=1) code = CodeSegment(engine) code.find_neighbours(field='whitenoise', features='features') field = pm.generate_whitenoise(seed=1234, unitary=True).c2r() field[...] = numpy.arange(field.csize).reshape(field.cshape)[field.slices] eps = field.cnorm()**0.5 * 1e-5 features = code.compute('features', init={'whitenoise': field}) check_grad(code, 'features', '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)