def test_2d_2d(comm): import pfft pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], np=pfft.split_size_2d(comm.size), comm=comm, dtype='f8') real = pm.generate_whitenoise(seed=123, mode='real') complex = pm.generate_whitenoise(seed=123, mode='complex') assert_array_equal(real, complex.c2r())
def test_2d_2d(comm): import pfft pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], np=pfft.split_size_2d(comm.size), comm=comm, dtype='f8') pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], np=pfft.split_size_2d(comm.size), comm=comm, dtype='f8') assert pm._use_padded == False real = pm.generate_whitenoise(seed=123, type='real') complex = pm.generate_whitenoise(seed=123, type='complex') assert_array_equal(real, complex.c2r()) real2 = pm2.generate_whitenoise(seed=123, type='real') assert real2.shape[:2] == real.shape
def test_cdot_types(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex') comp2 = pm.generate_whitenoise(1239, type='untransposedcomplex') with pytest.raises(TypeError): norm1 = comp1.cdot(comp2) norm2 = comp2.cdot(comp1) # this should work though fragile, because only on the Nmesh and 1 rank # the shape of values is the same. norm1 = comp1.cdot(comp2.value) norm2 = comp2.cdot(comp1.value)
def test_cnorm_log(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex', mean=1.0) norm2 = comp1.cnorm(norm = lambda x: numpy.log(x.real ** 2 + x.imag ** 2)) norm3 = (numpy.log(abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value))) ** 2)).sum() assert_allclose(norm2, norm3)
def test_whitenoise_mean(comm): # the whitenoise shall preserve the large scale. pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8') complex1 = pm0.generate_whitenoise(seed=8, unitary=True, mean=1.0) assert_allclose(complex1.c2r().cmean(), 1.0)
def test_whitenoise(comm): # the whitenoise shall preserve the large scale. pm0 = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8, 8], comm=comm, dtype='f8') pm1 = ParticleMesh(BoxSize=8.0, Nmesh=[16, 16, 16], comm=comm, dtype='f8') pm2 = ParticleMesh(BoxSize=8.0, Nmesh=[32, 32, 32], comm=comm, dtype='f8') complex1_down = ComplexField(pm0) complex2_down = ComplexField(pm0) complex1 = pm1.generate_whitenoise(seed=8, unitary=True) complex2 = pm2.generate_whitenoise(seed=8, unitary=True) complex1.resample(complex1_down) complex2.resample(complex2_down) mask1 = complex1_down.value != complex2_down.value assert_array_equal(complex1_down.value, complex2_down.value)
def test_kdk(comm): # Or use an object from astropy. cosmo = lambda : None cosmo.Om0 = 0.3 cosmo.Ode0 = 0.7 cosmo.Ok0 = 0.0 pm = ParticleMesh(BoxSize=128.0, Nmesh=(4,4,4), comm=comm, dtype='f8') vm = fastpm.Evolution(pm, B=1, shift=0.5) dlink = pm.generate_whitenoise(12345, mode='complex', unitary=True) dlink, ampl = _addampl(dlink) data = dlink.c2r() data[...] = 0 sigma = data.copy() sigma[...] = 1.0 code = vm.code() code.Decompress(C='dlin_k') code.KDKSimulation(dlin_k='dlin_k', mesh='mesh', cosmo=cosmo, astart=0.1, aend=1.0, Nsteps=5) code.Decompose() code.Paint() code.Residual(data_x=data, sigma_x=sigma) code.Chi2(variable='residual') _test_model(code, dlink, ampl)
def test_preview(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='real') preview = comp1.preview(axes=(0, 1, 2)) preview = comp1.preview(Nmesh=4, axes=(0, 1, 2)) for ind1 in numpy.ndindex(*(list(comp1.cshape))): assert_allclose(preview[ind1], comp1.cgetitem(ind1)) preview1 = comp1.preview(Nmesh=4, axes=(0, 1)) previewsum1 = preview.sum(axis=2) assert_allclose(preview1, previewsum1) preview2 = comp1.preview(Nmesh=4, axes=(1, 2)) previewsum2 = preview.sum(axis=0) assert_allclose(preview2, previewsum2) preview3 = comp1.preview(Nmesh=4, axes=(0, 2)) previewsum3 = preview.sum(axis=1) assert_allclose(preview3, previewsum3) preview4 = comp1.preview(Nmesh=4, axes=(2, 0)) previewsum4 = preview.sum(axis=1).T assert_allclose(preview4, previewsum4) preview5 = comp1.preview(Nmesh=4, axes=(0, )) previewsum5 = preview.sum(axis=(1, 2)) assert_allclose(preview5, previewsum5) preview6 = comp1.preview(Nmesh=8, axes=(0, ))
def test_nody(): """ Checking end to end nbody """ a0 = 0.1 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) solver = Solver(pm, Planck15, B=1) stages = np.linspace(0.1, 1.0, 10, endpoint=True) # Generate initial state with fastpm whitec = pm.generate_whitenoise(100, mode='complex', unitary=False) lineark = whitec.apply(lambda k, v: Planck15.get_pklin( sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5) statelpt = solver.lpt(lineark, grid, a0, order=1) finalstate = solver.nbody(statelpt, leapfrog(stages)) final_cube = pm.paint(finalstate.X) # Same thing with flowpm tlinear = tf.expand_dims(np.array(lineark.c2r()), 0) state = tfpm.lpt_init(tlinear, a0, order=1) state = tfpm.nbody(state, stages, nc) tfread = pmutils.cic_paint(tf.zeros_like(tlinear), state[0]).numpy() assert_allclose(final_cube, tfread[0], atol=1.2)
def test_gravity(comm): from pmesh.pm import ParticleMesh import fastpm.operators as operators pm = ParticleMesh(BoxSize=4.0, Nmesh=(4, 4), comm=comm, method='cic', dtype='f8') vm = fastpm.Evolution(pm, shift=0.5) dlink = pm.generate_whitenoise(12345, mode='complex') code = vm.code() code.Decompress(C='dlin_k') code.LPTDisplace(dlin_k='dlin_k', D1=1.0, v1=0, D2=0.0, v2=0.0) code.ForcePaint() code.Force(factor=0.1) code.Chi2(variable='f') print(code) # FIXME: without the shift some particles have near zero dx1. # or near 1 dx1. # the gradient is not well approximated by the numerical if # any of the left or right value shifts beyond the support of # the window. # dlink, ampl = _addampl(dlink) _test_model(code, dlink, ampl)
def test_cnorm_grad(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex') def objective(comp1): return comp1.cnorm() grad_comp1 = comp1 * 2 grad_comp1.decompress_vjp(grad_comp1) ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])): dx1, c1 = perturb(comp1, ind1, dx) ng1 = (objective(c1) - objective(comp1)) / dx ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))): print (ind1, 'a', ag1, 'n', ng1) comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5, atol=1e-6)
def test_paint(comm): from pmesh.pm import ParticleMesh # initialize a random white noise field in real space pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm) real = pm.generate_whitenoise(type='real', seed=3333) # a RealField # FFT to a ComplexField complex = real.r2c() # gather to all ranks --> shape is now (8,8,8) on all ranks data = numpy.concatenate(comm.allgather(numpy.array(real.ravel())))\ .reshape(real.cshape) # mesh from data array realmesh = ArrayMesh(data, BoxSize=128., comm=comm) # paint() must yield the RealField assert_array_equal(real, realmesh.to_field(mode='real')) # gather complex to all ranks --> shape is (8,8,5) on all ranks cdata = numpy.concatenate(comm.allgather(numpy.array(complex.ravel())))\ .reshape(complex.cshape) # mesh from complex array cmesh = ArrayMesh(cdata, BoxSize=128., comm=comm) # paint() yields the ComplexField assert_allclose(complex, cmesh.to_field(mode='complex'), atol=1e-8)
def test_c2r_vjp(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8') real = pm.generate_whitenoise(1234, mode='real') comp = real.r2c() def objective(comp): real = comp.c2r() obj = (real.value ** 2).sum() return comm.allreduce(obj) grad_real = RealField(pm) grad_real[...] = real[...] * 2 grad_comp = ComplexField(pm) grad_comp = grad_real.c2r_vjp(grad_real) grad_comp.decompress_vjp(grad_comp) ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(grad_comp.cshape) + [2])): dx1, c1 = perturb(comp, ind1, dx) ng1 = (objective(c1) - objective(comp)) / dx ag1 = grad_comp.cgetitem(ind1) * dx1 / dx comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5)
def test_c2r_vjp(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8') real = pm.generate_whitenoise(1234, type='real', mean=1.0) comp = real.r2c() def objective(comp): real = comp.c2r() obj = (real.value ** 2).sum() return comm.allreduce(obj) grad_real = RealField(pm) grad_real[...] = real[...] * 2 grad_comp = ComplexField(pm) grad_comp = grad_real.c2r_vjp(grad_real) grad_comp.decompress_vjp(grad_comp) ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(grad_comp.cshape) + [2])): dx1, c1 = perturb(comp, ind1, dx) ng1 = (objective(c1) - objective(comp)) / dx ag1 = grad_comp.cgetitem(ind1) * dx1 / dx comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5)
def test_cnorm_grad(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, mode='complex') def objective(comp1): return comp1.cnorm() grad_comp1 = comp1 * 2 grad_comp1.decompress_vjp(grad_comp1) ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])): dx1, c1 = perturb(comp1, ind1, dx) ng1 = (objective(c1) - objective(comp1)) / dx ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))): print (ind1, 'a', ag1, 'n', ng1) comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5, atol=1e-6)
def test_preview(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='real') preview = comp1.preview(axes=(0, 1, 2)) preview = comp1.preview(Nmesh=4, axes=(0, 1, 2)) for ind1 in numpy.ndindex(*(list(comp1.cshape))): assert_allclose(preview[ind1], comp1.cgetitem(ind1)) preview1 = comp1.preview(Nmesh=4, axes=(0, 1)) previewsum1 = preview.sum(axis=2) assert_allclose(preview1, previewsum1) preview2 = comp1.preview(Nmesh=4, axes=(1, 2)) previewsum2 = preview.sum(axis=0) assert_allclose(preview2, previewsum2) preview3 = comp1.preview(Nmesh=4, axes=(0, 2)) previewsum3 = preview.sum(axis=1) assert_allclose(preview3, previewsum3) preview4 = comp1.preview(Nmesh=4, axes=(2, 0)) previewsum4 = preview.sum(axis=1).T assert_allclose(preview4, previewsum4) preview5 = comp1.preview(Nmesh=4, axes=(0,)) previewsum5 = preview.sum(axis=(1, 2)) assert_allclose(preview5, previewsum5) preview6 = comp1.preview(Nmesh=8, axes=(0,))
def test_cnorm_log(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex', mean=1.0) norm2 = comp1.cnorm(norm=lambda x: numpy.log(x.real**2 + x.imag**2)) norm3 = (numpy.log(abs(numpy.fft.fftn(numpy.fft.irfftn( comp1.value)))**2)).sum() assert_allclose(norm2, norm3)
def test_cdot_cnorm(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex') norm1 = comp1.cdot(comp1) norm2 = comp1.cnorm() norm3 = (abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value)))**2).sum() assert_allclose(norm2, norm3) assert_allclose(norm2, norm1)
def test_cdot_cnorm(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex') norm1 = comp1.cdot(comp1) norm2 = comp1.cnorm() norm3 = (abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value))) ** 2).sum() assert_allclose(norm2, norm3) assert_allclose(norm2, norm1)
def test_whitenoise_untransposed(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f4') f1 = pm.generate_whitenoise(seed=3333, type='untransposedcomplex') f2 = pm.generate_whitenoise(seed=3333, type='transposedcomplex') f1r = numpy.concatenate(comm.allgather(numpy.array(f1.ravel()))) f2r = numpy.concatenate(comm.allgather(numpy.array(f2.ravel()))) assert_array_equal(f1r, f2r) # this should have asserted r2c transforms as well. r1 = f1.c2r() r2 = f2.c2r() r1r = numpy.concatenate(comm.allgather(numpy.array(r1.ravel()))) r2r = numpy.concatenate(comm.allgather(numpy.array(r2.ravel()))) assert_array_equal(r1r, r2r)
def make_data(bs, nc, seed, nsteps, path='', z=0): #initiate pm = ParticleMesh(BoxSize=bs, Nmesh=(nc, nc, nc), dtype='f8') ofolder = path + 'z%02d/L%04d_N%04d_S%04d_%02dstep/'%(z*10, bs, nc, seed, nsteps) # if pm.comm.rank == 0: print('Make data for seed = %d'%seed) print('Data to be saved at path - %s'%ofolder) klin, plin = numpy.loadtxt('ics_matterpow_0.dat', unpack = True) pk = interpolate(klin, plin) cosmo = Planck15.clone(Omega_cdm = 0.2685, h = 0.6711, Omega_b = 0.049) #pk = EHPower(Planck15, redshift=0) #cosmo = Planck15 s_truth = pm.generate_whitenoise(seed, type='complex')\ .apply(lambda k, v: v * (pk(sum(ki **2 for ki in k) **0.5) / v.BoxSize.prod()) ** 0.5)\ .c2r() #dynamics aa = 1.0/(1+z) if pm.comm.rank == 0: print('Evolve to redshift = %0.1f, scale factor = %0.2f'%(z, aa)) stages = numpy.linspace(0.1, aa, nsteps, endpoint=True) start = time() dynamic_model = NBodyModel(cosmo, pm, B=2, steps=stages) #dynamic_model = LPTModel(cosmo, pm, B=2, steps=stages) #Save data X, V, final = dynamic_model.get_code().compute(['X', 'V', 'final'], init={'parameters':s_truth}) end = time() print('Time taken = ', end-start) save_map(s_truth, ofolder + 'mesh', 's') save_map(final, ofolder + 'mesh', 'd') if pm.comm.rank == 0: print('X, V computed') cat = ArrayCatalog({'Position': X, 'Velocity' : V}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh) kdd = KDDensity(cat).density cat['KDDensity'] = kdd cat['InitPosition'] = dynamic_model.get_code().engine.q cat.save(ofolder + 'dynamic/1', ('InitPosition', 'Position', 'Velocity', 'KDDensity')) if pm.comm.rank == 0: print('dynamic model created') #FOF fof = FOF(cat, linking_length=0.2, nmin=12) fofcat = fof.find_features(peakcolumn='KDDensity') fofcat['Mass'] = fofcat['Length'] * cosmo.Om0 * 27.7455 * pm.BoxSize.prod() / pm.Nmesh.prod() fofcat.save(ofolder+'FOF', ('CMPosition', 'CMVelocity', 'PeakPosition', 'PeakVelocity', 'Length', 'Mass'))
def test_cdot_grad(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex', mean=1) comp2 = pm.generate_whitenoise(1235, type='complex', mean=1) def objective(comp1, comp2): return comp1.cdot(comp2).real y = objective(comp1, comp2) grad_comp2 = comp1.cdot_vjp(1.0) grad_comp1 = comp2.cdot_vjp(1.0) grad_comp1.decompress_vjp(grad_comp1) grad_comp2.decompress_vjp(grad_comp2) print(grad_comp1) print("comp1") ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])): dx1, c1 = perturb(comp1, ind1, dx) ng1 = (objective(c1, comp2) - objective(comp1, comp2)) / dx ag1 = grad_comp1.cgetitem(ind1) * dx1 / dx if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))): print (ind1, 'a', ag1, 'n', ng1) comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5) print("comp2") ng = [] ag = [] ind = [] dx = 1e-7 for ind1 in numpy.ndindex(*(list(comp1.cshape) + [2])): dx1, c2 = perturb(comp2, ind1, dx) ng1 = (objective(comp1, c2) - objective(comp1, comp2)) / dx ag1 = grad_comp2.cgetitem(ind1) * dx1 / dx if abs(ag1 - ng1) > 1e-5 * max((abs(ag1), abs(ng1))): print (ind1, 'a', ag1, 'n', ng1) comm.barrier() ng.append(ng1) ag.append(ag1) ind.append(ind1) assert_allclose(ng, ag, rtol=1e-5)
def test_cdot(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='complex') comp2 = pm.generate_whitenoise(1239, type='complex') norm1 = comp1.cdot(comp2) norm2 = comp2.cdot(comp1) norm_r = comp1.c2r().cdot(comp2.c2r()) / pm.Nmesh.prod() assert_allclose(norm2.real, norm_r) assert_allclose(norm1.real, norm_r) assert_allclose(norm1.real, norm2.real) assert_allclose(norm1.imag, -norm2.imag)
def test_prior(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=128.0, Nmesh=(4, 4), comm=comm) vm = fastpm.Evolution(pm, shift=0.5) dlink = pm.generate_whitenoise(1234, mode='complex') code = vm.code() code.Prior(dlin_k='dlin_k', powerspectrum=lambda k : 1.0) code.CopyVariable(x='prior', y='chi2') dlink, ampl = _addampl(dlink) _test_model(code, dlink, ampl)
def simulate(comm, ns): pm = ParticleMesh(BoxSize=ns.BoxSize, Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh], dtype='f8', comm=comm) gaussian = pm.generate_whitenoise(ns.seed, unitary=True) time_steps = numpy.linspace(ns.ainit, ns.afinal, ns.steps, endpoint=True) Q = pm.generate_uniform_particle_grid(shift=0) print(Q.min(axis=0), Q.max(axis=0)) def convolve(k, v): kmag = sum(ki**2 for ki in k)**0.5 ampl = (PowerSpectrum(kmag) / v.BoxSize.prod())**0.5 return v * ampl dlinear = gaussian.apply(convolve) DX1 = numpy.zeros_like(Q) layout = pm.decompose(Q) # Fill it in one dimension at a time. for d in range(pm.ndim): DX1[..., d] = dlinear \ .apply(dx1_transfer(d)) \ .c2r().readout(Q, layout=layout) a0 = time_steps[0] # 1-LPT Displacement and Veloicty; scaled back from z=0 to the first time step. S = DX1 * pt.D1(a=a0) V = S * a0**2 * pt.f1(a0) * pt.E(a0) state = State(Q, S, V) fpm = ParticleMesh(BoxSize=pm.BoxSize, Nmesh=pm.Nmesh * ns.boost, resampler='tsc', dtype='f8') ns.scheme(fpm, state, time_steps, ns.factors) r = Result() r.Q = Q r.DX1 = DX1 r.S = S r.V = V r.dlinear = dlinear return r
def test_paint(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm) real = pm.generate_whitenoise(type='real', seed=3333) complex = real.r2c() realmesh = FieldMesh(real) complexmesh = FieldMesh(complex) assert_array_equal(real, realmesh.to_field(mode='real')) assert_array_equal(complex, complexmesh.to_field(mode='complex')) assert_allclose(complex, realmesh.to_field(mode='complex')) assert_allclose(real, complexmesh.to_field(mode='real'))
def test_paint(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(Nmesh=(8, 8, 8), BoxSize=(128, 128, 128.), comm=comm) real = pm.generate_whitenoise(mode='real', seed=3333) complex = real.r2c() realmesh = FieldMesh(real) complexmesh = FieldMesh(complex) assert_array_equal(real, realmesh.to_field(mode='real')) assert_array_equal(complex, complexmesh.to_field(mode='complex')) assert_allclose(complex, realmesh.to_field(mode='complex')) assert_allclose(real, complexmesh.to_field(mode='real'))
def test_lpt(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=128.0, Nmesh=(4, 4), comm=comm) vm = fastpm.Evolution(pm, shift=0.5) dlink = pm.generate_whitenoise(1234, mode='complex') code = vm.code() code.Decompress(C='dlin_k') code.LPTDisplace(dlin_k='dlin_k', D1=1.0, v1=0, D2=0.0, v2=0.0) code.Chi2(variable='s') dlink, ampl = _addampl(dlink) _test_model(code, dlink, ampl)
def simulate(comm, ns): pm = ParticleMesh(BoxSize=ns.BoxSize, Nmesh=[ns.Nmesh, ns.Nmesh, ns.Nmesh], dtype='f8', comm=comm) gaussian = pm.generate_whitenoise(ns.seed, unitary=True) time_steps = numpy.linspace(ns.ainit, ns.afinal, ns.steps, endpoint=True) Q = pm.generate_uniform_particle_grid(shift=0) print(Q.min(axis=0), Q.max(axis=0)) def convolve(k, v): kmag = sum(ki**2 for ki in k) ** 0.5 ampl = (PowerSpectrum(kmag) / v.BoxSize.prod()) ** 0.5 return v * ampl dlinear = gaussian.apply(convolve) DX1 = numpy.zeros_like(Q) layout = pm.decompose(Q) # Fill it in one dimension at a time. for d in range(pm.ndim): DX1[..., d] = dlinear \ .apply(dx1_transfer(d)) \ .c2r().readout(Q, layout=layout) a0 = time_steps[0] # 1-LPT Displacement and Veloicty; scaled back from z=0 to the first time step. S = DX1 * pt.D1(a=a0) V = S * a0 ** 2 * pt.f1(a0) * pt.E(a0) state = State(Q, S, V) fpm = ParticleMesh(BoxSize=pm.BoxSize, Nmesh=pm.Nmesh * ns.boost, resampler='tsc', dtype='f8') ns.scheme(fpm, state, time_steps, ns.factors) r = Result() r.Q = Q r.DX1 = DX1 r.S = S r.V = V r.dlinear = dlinear return r
def test_lpt2(): """ Checking lpt2_source, this also checks the laplace and gradient kernels """ pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) whitec = pm.generate_whitenoise(100, mode='complex', unitary=False) lineark = whitec.apply(lambda k, v: Planck15.get_pklin( sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5) # Compute lpt1 from fastpm with matching kernel order source = fpmops.lpt2source(lineark).c2r() # Same thing from tensorflow tfsource = tfpm.lpt2_source( pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0))) tfread = pmutils.c2r3d(tfsource).numpy() assert_allclose(source, tfread[0], atol=1e-5)
def test_resample(comm): if comm.size == 4: comm.barrier() from pmesh.pm import ParticleMesh # testing the gradient of down sampling, which we use the most. pm = ParticleMesh(BoxSize=128.0, Nmesh=(8, 8), comm=comm) vm = fastpm.Evolution(pm, shift=0.5) dlink = pm.generate_whitenoise(1234, mode='complex') dlink, ampl = _addampl(dlink) code = vm.code() code.Decompress(C='dlin_k') code.C2R(C='dlin_k', R='mesh') code.Resample(Neff=4) code.Chi2(variable='mesh') _test_model(code, dlink, ampl)
def test_lpt1(): """ Checking lpt1, this also checks the laplace and gradient kernels """ pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) whitec = pm.generate_whitenoise(100, mode='complex', unitary=False) lineark = whitec.apply(lambda k, v: Planck15.get_pklin( sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5) # Compute lpt1 from fastpm with matching kernel order lpt = fpmops.lpt1(lineark, grid) # Same thing from tensorflow with tf.Session() as sess: state = tfpm.lpt1( pmutils.r2c3d(tf.expand_dims(tf.constant(lineark.c2r()), axis=0)), grid.reshape((1, -1, 3)) * nc / bs) tfread = sess.run(state) assert_allclose(lpt, tfread[0] * bs / nc, atol=1e-5)
def test_lpt_init(): """ Checking lpt init """ a0 = 0.1 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) solver = Solver(pm, Planck15, B=1) # Generate initial state with fastpm whitec = pm.generate_whitenoise(100, mode='complex', unitary=False) lineark = whitec.apply(lambda k, v: Planck15.get_pklin( sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5) statelpt = solver.lpt(lineark, grid, a0, order=1) # Same thing with flowpm tlinear = tf.expand_dims(np.array(lineark.c2r()), 0) tfread = tfpm.lpt_init(tlinear, a0, order=1).numpy() assert_allclose(statelpt.X, tfread[0, 0] * bs / nc, rtol=1e-2)
def test_lpt1_64(): """ Checking lpt1, this also checks the laplace and gradient kernels This variant of the test checks that it works for cubes of size 64 """ nc = 64 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') grid = pm.generate_uniform_particle_grid(shift=0).astype(np.float32) whitec = pm.generate_whitenoise(100, mode='complex', unitary=False) lineark = whitec.apply(lambda k, v: Planck15.get_pklin( sum(ki**2 for ki in k)**0.5, 0)**0.5 * v / v.BoxSize.prod()**0.5) # Compute lpt1 from fastpm with matching kernel order lpt = fpmops.lpt1(lineark, grid) # Same thing from tensorflow tfread = tfpm.lpt1( pmutils.r2c3d(tf.expand_dims(np.array(lineark.c2r()), axis=0)), grid.reshape((1, -1, 3)) * nc / bs).numpy() assert_allclose(lpt, tfread[0] * bs / nc, atol=5e-5)
def test_lpt_prior(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=128.0, Nmesh=(4, 4), comm=comm) vm = fastpm.Evolution(pm, shift=0.5) dlink = pm.generate_whitenoise(1234, mode='complex') code = vm.code() code.Decompress(C='dlin_k') code.LPTDisplace(dlin_k='dlin_k', D1=1.0, v1=0, D2=0.0, v2=0.0) code.Decompose() code.Paint() code.Chi2(variable='mesh') code.Prior(powerspectrum=lambda k : 1.0) code.Multiply(a='prior', b='prior', f=0.1) code.Add(a='prior', b='chi2', c='chi2') dlink, ampl = _addampl(dlink) _test_model(code, dlink, ampl)
def test_transpose(comm): pm = ParticleMesh(BoxSize=[8.0, 16.0, 32.0], Nmesh=[4, 6, 8], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, type='real') comp1t = comp1.ctranspose([0, 1, 2]) assert_array_equal(comp1t.Nmesh, comp1.Nmesh) assert_array_equal(comp1t.BoxSize, comp1.BoxSize) assert_array_equal(comp1t.cnorm(), comp1.cnorm()) comp1t = comp1.ctranspose([1, 2, 0]) assert_array_equal(comp1t.Nmesh, comp1.Nmesh[[1, 2, 0]]) assert_array_equal(comp1t.BoxSize, comp1.BoxSize[[1, 2, 0]]) comp1tt = comp1t.ctranspose([1, 2, 0]) comp1ttt = comp1tt.ctranspose([1, 2, 0]) assert_allclose(comp1ttt, comp1)
def test_cdot_cnorm(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') comp1 = pm.generate_whitenoise(1234, mode='complex') norm1 = comp1.cdot(comp1) norm2 = comp1.cnorm() norm3 = (abs(numpy.fft.fftn(numpy.fft.irfftn(comp1.value)))**2).sum() assert_allclose(norm2, norm3) # p-2 norm of the full complex matrix (norm2) is # NORM = Self + Lower + Upper # CDOT = Self + Lower # Upper = Lower self_conj_sum = 0 for ind1 in numpy.ndindex(*(list(comp1.cshape))): c = [(4 - i) % 4 for i, n in zip(ind1, comp1.cshape)] ci = numpy.ravel_multi_index(c, comp1.Nmesh) ii = numpy.ravel_multi_index(ind1, comp1.Nmesh) v = comp1.cgetitem(ind1) if ci == ii: self_conj_sum += abs(v)**2 assert_allclose(norm1 * 2, norm2 + self_conj_sum)
def test_1d(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8], comm=comm, dtype='f8') real = pm.generate_whitenoise(seed=123, type='real') complex = pm.generate_whitenoise(seed=123, type='complex') assert_array_equal(real, complex.c2r())
def test_readout_gradients(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic') real = pm.generate_whitenoise(1234, type='real') def objective(real, pos, layout): value = real.readout(pos, layout=layout) obj = (value ** 2).sum() return comm.allreduce(obj) def forward_gradient(real, pos, layout, v_real=None, v_pos=None): value = real.readout(pos, layout=layout) v_value = real.readout_jvp(pos, v_self=v_real, v_pos=v_pos, layout=layout) return comm.allreduce((v_value * value * 2).sum()) def backward_gradient(real, pos, layout): value = real.readout(pos, layout=layout) return real.readout_vjp(pos, v=value * 2, layout=layout) pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T pos += real.start # avoid sitting at the pmesh points # cic gradient is zero on them, the numerical gradient fails. pos += 0.5 pos *= pm.BoxSize / pm.Nmesh layout = pm.decompose(pos) obj = objective(real, pos, layout) grad_real, grad_pos = backward_gradient(real, pos, layout) ng = [] f*g = [] bag = [] ind = [] dx = 1e-6 for ind1 in numpy.ndindex(*grad_real.cshape): dx1, r1 = perturb(real, ind1, dx) ng1 = (objective(r1, pos, layout) - obj) bag1 = grad_real.cgetitem(ind1) * dx1 fag1 = forward_gradient(real, pos, layout, v_real=r1 - real) # print (dx1, dx, ind1, ag1, ng1) ng.append(ng1) f*g.append(fag1) bag.append(bag1) ind.append(ind1) assert_allclose(bag, f*g, rtol=1e-7) assert_allclose(ng, bag, rtol=1e-4) # FIXME ng = [] bag = [] f*g = [] ind = [] dx = 1e-6 for ind1 in numpy.ndindex((real.csize, real.ndim)): dx1, pos1, old = perturb_pos(pos, ind1, dx, comm) layout1 = pm.decompose(pos1) ng1 = (objective(real, pos1, layout1) - obj) bag1 = get_pos(grad_pos, ind1, comm) * dx1 fag1 = forward_gradient(real, pos, layout, v_pos=pos1 - pos) # print ('pos', old, ind1, 'a', ag1, 'n', ng1) ng.append(ng1) bag.append(bag1) f*g.append(fag1) ind.append(ind1) assert_allclose(bag, f*g, rtol=1e-7) assert_allclose(ng, bag, rtol=1e-4)
def test_paint_gradients(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8', resampler='cic') real = pm.generate_whitenoise(1234, type='real') def objective(pos, mass, layout): real = pm.paint(pos, mass=mass, layout=layout) obj = (real[...] ** 2).sum() return comm.allreduce(obj) def forward_gradient(pos, mass, layout, v_pos=None, v_mass=None): real = pm.paint(pos, mass=mass, layout=layout) jvp = pm.paint_jvp(pos, mass=mass, v_mass=v_mass, v_pos=v_pos, layout=layout) return comm.allreduce((jvp * real * 2)[...].sum()) def backward_gradient(pos, mass, layout): real = pm.paint(pos, mass=mass, layout=layout) return pm.paint_vjp(real * 2, pos, mass=mass, layout=layout) pos = numpy.array(numpy.indices(real.shape), dtype='f8').reshape(real.value.ndim, -1).T pos += real.start numpy.random.seed(9999) # avoid sitting at the pmesh points # cic gradient is zero on them, the numerical gradient fails. pos += (numpy.random.uniform(size=pos.shape)) * 0.8 + 0.1 pos *= pm.BoxSize / pm.Nmesh mass = numpy.ones(len(pos)) * 2 layout = pm.decompose(pos) obj = objective(pos, mass, layout) grad_pos, grad_mass = backward_gradient(pos, mass, layout) ng = [] f*g = [] bag = [] ind = [] dx = 1e-6 for ind1 in numpy.ndindex(real.csize): dx1, mass1, old = perturb_mass(mass, ind1[0], dx, comm) ng1 = (objective(pos, mass1, layout) - obj) bag1 = get_mass(grad_mass, ind1[0], comm) * dx1 fag1 = forward_gradient(pos, mass, layout, v_mass=mass1 - mass) # print (dx1, dx, ind1, fag1, bag1, ng1) ng.append(ng1) f*g.append(fag1) bag.append(bag1) ind.append(ind1) assert_allclose(bag, f*g, rtol=1e-7) assert_allclose(ng, bag, rtol=1e-4) # FIXME ng = [] bag = [] f*g = [] ind = [] dx = 1e-6 for ind1 in numpy.ndindex((real.csize, real.ndim)): dx1, pos1, old = perturb_pos(pos, ind1, dx, comm) layout1 = pm.decompose(pos1) ng1 = (objective(pos1, mass, layout1) - obj) bag1 = get_pos(grad_pos, ind1, comm) * dx1 fag1 = forward_gradient(pos, mass, layout, v_pos=pos1 - pos) # print ('pos', old, ind1, 'v_pos', (pos1 - pos)[ind1[0]], 'f*g', fag1, 'bag', bag1, 'n', ng1) ng.append(ng1) bag.append(bag1) f*g.append(fag1) ind.append(ind1) assert_allclose(bag, f*g, rtol=1e-7) assert_allclose(ng, bag, rtol=1e-4)