def test_solver_convergence(comm): from mpi4py import MPI pm = ParticleMesh(BoxSize=128., Nmesh=[8, 8, 8], comm=comm) pm_ref = ParticleMesh(BoxSize=128., Nmesh=[8, 8, 8], comm=MPI.COMM_SELF) Plin = LinearPower(Planck15, redshift=0, transfer='EisensteinHu') Q = pm_ref.generate_uniform_particle_grid(shift=0) def solve(pm): solver = Solver(pm, Planck15, B=1) q = numpy.array_split(Q, pm.comm.size)[pm.comm.rank] wn = solver.whitenoise(1234) dlin = solver.linear(wn, lambda k: Plin(k)) state = solver.lpt(dlin, q, a=0.1, order=2) state = solver.nbody(state, leapfrog([0.1, 0.5, 1.0])) d = {} for key in 'X', 'P', 'F': d[key] = numpy.concatenate(pm.comm.allgather(getattr(state, key)), axis=0) return d s = solve(pm) r = solve(pm_ref) assert_allclose(s['X'], r['X']) assert_allclose(s['P'], r['P']) assert_allclose(s['F'], r['F'])
def test_decompose(comm): pm = ParticleMesh(BoxSize=4.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') numpy.random.seed(1234) if comm.rank == 0: Npar = 1000 else: Npar = 0 pos = 4.0 * (numpy.random.uniform(size=(Npar, 3))) pos = pm.generate_uniform_particle_grid(shift=0.5) all_pos = numpy.concatenate(comm.allgather(pos), axis=0) for resampler in ['cic', 'tsc', 'db12']: def test(resampler): print(resampler) truth = numpy.zeros(pm.Nmesh, dtype='f8') affine = window.Affine(ndim=3, period=4) window.FindResampler(resampler).paint(truth, all_pos, transform=affine) truth = comm.bcast(truth) layout = pm.decompose(pos, smoothing=resampler) npos = layout.exchange(pos) real = pm.paint(npos, resampler=resampler) full = numpy.zeros(pm.Nmesh, dtype='f8') full[real.slices] = real full = comm.allreduce(full) #print(full.sum(), pm.Nmesh.prod()) #print(truth.sum(), pm.Nmesh.prod()) #print(comm.rank, npos, real.slices) assert_almost_equal(full, truth) # can't yield! test(resampler)
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_grid(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') grid, id = pm.generate_uniform_particle_grid(shift=0.5, return_id=True) assert len(id) == len(grid) allid = numpy.concatenate(comm.allgather(id), axis=0) # must be all unique assert len(numpy.unique(allid)) == len(allid) assert numpy.max(allid) == len(allid) - 1 assert numpy.min(allid) == 0
def test_fof_parallel_no_merge(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=[8, 8, 8], Nmesh=[8, 8, 8], comm=comm) Q = pm.generate_uniform_particle_grid() cat = ArrayCatalog({'Position' : Q}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh, comm=comm) fof = FOF(cat, linking_length=0.9, nmin=0) labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0) # one particle per group assert max(labels) == cat.csize - 1
def test_fof_parallel_no_merge(comm): CurrentMPIComm.set(comm) from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=[8, 8, 8], Nmesh=[8, 8, 8], comm=comm) Q = pm.generate_uniform_particle_grid() cat = ArrayCatalog({'Position': Q}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh) fof = FOF(cat, linking_length=0.9, nmin=0) labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0) # one particle per group assert max(labels) == cat.csize - 1
def test_fof_fully_connected(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=[4, 4, 4], Nmesh=[4, 4, 4], comm=comm) Q = pm.generate_uniform_particle_grid(shift=0) cat = ArrayCatalog({'Position': numpy.concatenate([Q], axis=0)}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh, comm=comm) fof = FOF(cat, linking_length=2, nmin=63, absolute=True) labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0) assert_array_equal(labels, 1)
def getPow(data1, data2=None): pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32]) q = pm.generate_uniform_particle_grid() den1 = pm.paint(q + data1.reshape([-1, 3])) if (data2 is not None): pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32]) den2 = pm.paint(q + data2.reshape([-1, 3])) temp = FFTPower(first=den1, second=den2, mode='1d', BoxSize=128, dk=dk) k, powspec = temp.power['k'], temp.power['power'] else: temp = FFTPower(den1, mode='1d', BoxSize=128, dk=dk) k, powspec = temp.power['k'], temp.power['power'] return [k, powspec.real]
def test_ncdm(comm): pm = ParticleMesh(BoxSize=512., Nmesh=[8, 8, 8], comm=comm) Plin = LinearPower(Planck15, redshift=0, transfer='EisensteinHu') solver = Solver(pm, Planck15, B=1) Q = pm.generate_uniform_particle_grid(shift=0) wn = solver.whitenoise(1234) dlin = solver.linear(wn, lambda k: Plin(k)) state = solver.lpt(dlin, Q, a=1.0, order=2) dnonlin = solver.nbody(state, leapfrog([0.1, 1.0])) dnonlin.save('nonlin')
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 getPow_dis(data1, data2=None): pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32]) q = pm.generate_uniform_particle_grid() den1 = pm.paint(q) power = 0 if (data2 is not None): pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32]) q = pm.generate_uniform_particle_grid() den2 = pm.paint(q) for ii in range(3): den1[:] = data1[:, :, :, ii] den2[:] = data2[:, :, :, ii] temp = FFTPower(first=den1, second=den2, mode='1d', BoxSize=128, dk=dk) k, power = temp.power['k'], power + temp.power['power'] else: for ii in range(3): den1[:] = data1[:, :, :, ii] temp = FFTPower(den1, mode='1d', BoxSize=128, dk=dk) k, power = temp.power['k'], power + temp.power['power'] return [k, power.real]
def test_grid_shifted(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') grid = pm.generate_uniform_particle_grid(shift=0.5) grid = grid + 4.0 assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod()) layout = pm.decompose(grid) g2 = layout.exchange(grid) real = pm.paint(grid, layout=layout) #print(real, g2 % 8, real.slices) assert_allclose(real, 1.0) grid = grid - 6.1 assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod()) layout = pm.decompose(grid) real = pm.paint(grid, layout=layout) assert_allclose(real, 1.0)
def test_fof_parallel_merge(comm): from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=[8, 8, 8], Nmesh=[8, 8, 8], comm=comm) Q = pm.generate_uniform_particle_grid(shift=0) Q1 = Q.copy() Q1[:] += 0.01 Q2 = Q.copy() Q2[:] -= 0.01 Q3 = Q.copy() Q3[:] += 0.02 cat = ArrayCatalog({'Position' : numpy.concatenate([Q, Q1, Q2, Q3], axis=0)}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh, comm=comm) fof = FOF(cat, linking_length=0.011 * 3 ** 0.5, nmin=0, absolute=True) labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0) assert max(labels) == pm.Nmesh.prod() - 1 assert all(numpy.bincount(labels) == 4)
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_fof_parallel_merge(comm): CurrentMPIComm.set(comm) from pmesh.pm import ParticleMesh pm = ParticleMesh(BoxSize=[32, 32, 32], Nmesh=[32, 32, 32], comm=comm) Q = pm.generate_uniform_particle_grid(shift=0) Q1 = Q.copy() Q1[:] += 0.01 Q2 = Q.copy() Q2[:] -= 0.01 Q3 = Q.copy() Q3[:] += 0.02 cat = ArrayCatalog({'Position' : numpy.concatenate([Q, Q1, Q2, Q3], axis=0)}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh) fof = FOF(cat, linking_length=0.011 * 3 ** 0.5, nmin=0, absolute=True) labels = numpy.concatenate(comm.allgather((fof.labels)), axis=0) assert max(labels) == pm.Nmesh.prod() - 1 assert all(numpy.bincount(labels) == 4)
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_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_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 set_up(): params = get_params() pm = ParticleMesh(Nmesh=params['Nmesh'], BoxSize=params['BoxSize'], comm=MPI.COMM_WORLD, resampler='cic') # generate initial conditions cosmo = Planck15.clone(P_k_max=30) x = pm.generate_uniform_particle_grid(shift=0.1) BoxSize2D = [deg / 180. * np.pi for deg in params['BoxSize2D']] logger = None sim = lightcone.WLSimulation(stages=numpy.linspace(0.1, 1.0, params['N_steps'], endpoint=True), cosmology=cosmo, pm=pm, boxsize2D=BoxSize2D, params=params, logger=None) kmaps = [sim.mappm.create('real', value=0.) for ii in range(1)] return pm, cosmo, x, kmaps, sim.DriftFactor, sim.mappm, sim
def test_solver(comm): pm = ParticleMesh(BoxSize=512., Nmesh=[8, 8, 8], comm=comm) solver = Solver(pm, Planck15, B=1) P_prm = Planck15.Primordial.get_pkprim tf = get_species_transfer_function_from_class(Planck15, 9) Q = pm.generate_uniform_particle_grid(shift=0) wn = solver.whitenoise(1234) prm = solver.primordial(wn, P_prm) ic = solver.lpt(prm, { '0': (Baryon, tf['d_b'], tf['dd_b']), '1': (CDM, tf['d_cdm'], tf['dd_cdm']), '4': (NCDM, tf['d_ncdm[0]'], tf['dd_ncdm[0]']), }, Q, a=0.1) print('0', ic.species['0'].S[0], ic.species['0'].P[0], ic.species['0'].Q[0]) print('1', ic.species['1'].S[0], ic.species['1'].P[0], ic.species['1'].Q[0]) print('4', ic.species['4'].S[0], ic.species['4'].P[0], ic.species['4'].Q[0]) c2 = CoreSolver(pm, Planck15, B=1) Pk = lambda k: Planck15.get_pk(k, z=0) dlin = c2.linear(wn, Pk) ic2 = c2.lpt(dlin, Q, 0.1, order=1) print(ic2.S[0], ic2.P[0], ic2.Q[0]) final2 = c2.nbody(ic2, leapfrog([0.1, 1.0])) final = solver.nbody(ic, leapfrog([0.1, 1.0])) print('0', final.species['0'].F[0]) print('1', final.species['1'].F[0]) print('4', final.species['4'].F[0]) print(final2.F[0]) final.to_catalog()
def Halos(bs, nc, seed, nstep, seed_hod, Omega_m, p_alpha, p_logMin, p_logM1, p_logM0, p_sigma_logM): ''' ''' # setup initial conditions Omegacdm = Omega_m-0.049, cosmo = NBlab.cosmology.Planck15.clone(Omega_cdm=Omegacdm, h = 0.6711, Omega_b = 0.049) power = NBlab.cosmology.LinearPower(cosmo, 0) klin = np.logspace(-4,2,1000) plin = power(klin) pkfunc = interpolate(klin, plin) # run the simulation pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc]) Q = pm.generate_uniform_particle_grid() stages = numpy.linspace(0.1, 1.0, nstep, endpoint=True) solver = Solver(pm, cosmo, B=2) wn = solver.whitenoise(seed) dlin = solver.linear(wn, pkfunc) state = solver.lpt(dlin, Q, stages[0]) state = solver.nbody(state, leapfrog(stages)) # create the catalogue cat = ArrayCatalog({'Position' : state.X, 'Velocity' : state.V, 'Displacement' : state.S, 'Density' : state.RHO}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh, M0 = Omega_m * 27.75e10 * bs**3 / (nc/2.0)**3) cat['KDDensity'] = KDDensity(cat).density # run FOF to construct Halo catalog fof = FOF(cat, linking_length=0.2, nmin=12) fofcat = fof.to_halos(particle_mass=cat.attrs['M0'], cosmo = cosmo, redshift = 0.0) return fofcat
Ng = args.Ng Lbox = args.Lbox print ("Seed = ",Rdm_seed) print ("Lbox = %.1f"%Lbox,"Ng = %d"%Ng,"RG = %.2f"%RG) # Choose cosmology # ------------------------------------------- cosmology = nbcosmos.WMAP9 mycosmo = Cosmos(FLRW=True,obj=cosmology) # generate linear density field at z=0 # ------------------------------------------- pm = ParticleMesh(BoxSize=Lbox, Nmesh=[Ng,Ng,Ng]) Q = pm.generate_uniform_particle_grid(shift=0) solver = Solver(pm,cosmology,B=1) wn = solver.whitenoise(seed = Rdm_seed) dlin = solver.linear(wn, lambda k: mycosmo.Pk_lin(k)) dx_field = dlin.c2r().value # dx_field is the density contrast field centered at 0 # initialize gsCR object, build xij^{-1} matrix for full 18 constraints # ------------------------------------------- fg = gsCR(mycosmo,Lbox=Lbox,Nmesh=Ng,RG=RG,CONS=['full']) fg.build_Xij_inv_matrix() print ("xij^{-1}:") print (fg.xij_tensor_inv) print ("*********************************************")
from fastpm.core import Solver as Solver from fastpm.core import leapfrog from nbodykit.cosmology import Planck15, EHPower from nbodykit.algorithms.fof import FOF from nbodykit.algorithms.fftpower import FFTPower from nbodykit.source import ArrayCatalog from pmesh.pm import ParticleMesh import numpy pm = ParticleMesh(BoxSize=512, Nmesh=[256, 256, 256], dtype='f8', resampler='tsc') Q = pm.generate_uniform_particle_grid() stages = numpy.linspace(0.1, 1.0, 20, endpoint=True) solver = Solver(pm, Planck15, B=2) solver_ncdm = SolverNCDM(pm, Planck15, B=2) wn = solver.whitenoise(400) dlin = solver.linear(wn, EHPower(Planck15, 0)) lpt = solver.lpt(dlin, Q, stages[0]) #lpt.S = numpy.float32(lpt.S) def monitor(action, ai, ac, af, state, event): if pm.comm.rank == 0: print(state.a['S'], state.a['P'], state.a['F'], state.S[0], state.P[0],
k[mask] = 0. return power #time steps in scale factor stages = np.linspace(0.1, 1., args.Nstep, endpoint=True) a_output = 1. / (np.array(args.output_redshift) + 1.) #IC 2LPT t = time.time() solver_IC = Solver(pm_IC, Planck15, B=2) wn = solver_IC.whitenoise(2695896) dlin = solver_IC.linear(wn, Power) Q = pm.generate_uniform_particle_grid(shift=0) solver = Solver(pm, Planck15, B=2) state = solver.lpt(dlin, Q, stages[0], order=2) if pm.comm.rank == 0: print('Finish generating initial conditions with 2LPT. Time:', time.time() - t) X0 = state.X V0 = state.V a0 = np.array(stages[0]) def monitor(action, ai, ac, af, state, event): if not state.synchronized: return
def test_grid(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[4, 4, 4], comm=comm, dtype='f8') grid = pm.generate_uniform_particle_grid(shift=0.5) assert_array_equal(pm.comm.allreduce(grid.shape[0]), pm.Nmesh.prod()) real = pm.paint(grid) assert_array_equal(real, 1.0)
def main(): """ Convert grids4plots 3d grids to 2d slices. """ ##################################### # PARSE COMMAND LINE ARGS ##################################### ap = ArgumentParser() ap.add_argument('--inbasepath', type=str, default='$SCRATCH/perr/grids4plots/', help='Input base path.') ap.add_argument('--outbasepath', type=str, default='$SCRATCH/perr/grids4plots/', help='Output base path.') ap.add_argument( '--inpath', type=str, default= 'main_calc_Perr_2020_Sep_22_18:44:31_time1600800271.dill', # laptop #default='main_calc_Perr_2020_Aug_26_02:49:57_time1598410197.dill', # cluster help='Input path.') ap.add_argument('--Rsmooth', type=float, default=2.0, help='3D Gaussian smoothing applied to field.') # min and max index included in output. inclusive. ap.add_argument('--ixmin', type=int, default=5, help='xmin of output. must be between 0 and Ngrid.') ap.add_argument('--ixmax', type=int, default=5, help='xmax of output') ap.add_argument('--iymin', type=int, default=0, help='ymin of output') ap.add_argument('--iymax', type=int, default=-1, help='ymax of output') ap.add_argument('--izmin', type=int, default=0, help='zmin of output') ap.add_argument('--izmax', type=int, default=-1, help='zmax of output') cmd_args = ap.parse_args() verbose = True ##################################### # START PROGRAM ##################################### comm = CurrentMPIComm.get() rank = comm.rank path = os.path.join(os.path.expandvars(cmd_args.inbasepath), cmd_args.inpath) if rank == 0: print('path: ', path) initialized_slicecat = False for fname in os.listdir(path): if fname.startswith('SLICE'): continue full_fname = os.path.join(path, fname) print('%d Reading %s' % (rank, full_fname)) inmesh = BigFileMesh(full_fname, dataset='tmp4storage', header='header') Ngrid = inmesh.attrs['Ngrid'] boxsize = inmesh.attrs['boxsize'] # apply smoothing mesh = apply_smoothing(inmesh, mode='Gaussian', R=cmd_args.Rsmooth) del inmesh # convert indices to modulo ngrid ixmin = cmd_args.ixmin % Ngrid ixmax = cmd_args.ixmax % Ngrid iymin = cmd_args.iymin % Ngrid iymax = cmd_args.iymax % Ngrid izmin = cmd_args.izmin % Ngrid izmax = cmd_args.izmax % Ngrid # convert to boxsize units (Mpc/h) xmin = float(ixmin) / float(Ngrid) * boxsize xmax = float(ixmax + 1) / float(Ngrid) * boxsize ymin = float(iymin) / float(Ngrid) * boxsize ymax = float(iymax + 1) / float(Ngrid) * boxsize zmin = float(izmin) / float(Ngrid) * boxsize zmax = float(izmax + 1) / float(Ngrid) * boxsize if not initialized_slicecat: # Generate catalog with positions of slice points, to readout mesh there. # First generate all 3D points. Then keep only subset in slice. # THen readout mesh at those points. # use pmesh generate_uniform_particle_grid # http://rainwoodman.github.io/pmesh/pmesh.pm.html?highlight= # readout#pmesh.pm.ParticleMesh.generate_uniform_particle_grid partmesh = ParticleMesh(BoxSize=boxsize, Nmesh=[Ngrid, Ngrid, Ngrid]) ptcles = partmesh.generate_uniform_particle_grid(shift=0.0, dtype='f8') #print("type ptcles", type(ptcles), ptcles.shape) #print("head ptcles:", ptcles[:5,:]) dtype = np.dtype([('Position', ('f8', 3))]) # number of rows is given by number of ptcles on this rank uni_cat_array = np.empty((ptcles.shape[0], ), dtype=dtype) uni_cat_array['Position'] = ptcles uni_cat = ArrayCatalog(uni_cat_array, comm=None, BoxSize=boxsize * np.ones(3), Nmesh=[Ngrid, Ngrid, Ngrid]) del ptcles del uni_cat_array print("%d: Before cut: local Nptcles=%d, global Nptcles=%d" % (comm.rank, uni_cat.size, uni_cat.csize)) # only keep points in the slice uni_cat = uni_cat[(uni_cat['Position'][:, 0] >= xmin) & (uni_cat['Position'][:, 0] < xmax) & (uni_cat['Position'][:, 1] >= ymin) & (uni_cat['Position'][:, 1] < ymax) & (uni_cat['Position'][:, 2] >= zmin) & (uni_cat['Position'][:, 2] < zmax)] print("%d: After cut: local Nptcles=%d, global Nptcles=%d" % (comm.rank, uni_cat.size, uni_cat.csize)) initialized_slicecat = True # read out full 3D mesh at catalog positions. this is a numpy array slicecat = readout_mesh_at_cat_pos(mesh=mesh, cat=uni_cat, readout_window='nearest') if rank == 0: print('slicecat type:', type(slicecat)) slicecat = GatherArray(slicecat, comm, root=0) if rank == 0: if not slicecat.shape == ((ixmax - ixmin + 1) * (iymax - iymin + 1) * (izmax - izmin + 1), ): raise Exception( 'Unexpected shape of particles read out on slice: %s' % str(slicecat.shape)) slicecat = slicecat.reshape( (ixmax - ixmin + 1, iymax - iymin + 1, izmax - izmin + 1)) print('slicecat shape:', slicecat.shape) if verbose: print('slicecat:', slicecat) # convert to a mesh. assume full numpy array sits on rank 0. Lx = xmax - xmin Ly = ymax - ymin Lz = zmax - zmin if Lx == 0.: Lx = boxsize / float(Ngrid) if Ly == 0.: Ly = boxsize / float(Ngrid) if Lz == 0.: Lz = boxsize / float(Ngrid) BoxSize_slice = np.array([Lx, Ly, Lz]) slicemesh = ArrayMesh(slicecat, BoxSize=BoxSize_slice, root=0) outshape = slicemesh.compute(mode='real').shape if verbose: print('slicemesh: ', slicemesh.compute(mode='real')) # write to disk outpath = os.path.join( os.path.expandvars(cmd_args.outbasepath), cmd_args.inpath, 'SLICE_R%g_%d-%d_%d-%d_%d-%d/' % (cmd_args.Rsmooth, ixmin, ixmax, iymin, iymax, izmin, izmax)) if rank == 0: if not os.path.exists(outpath): os.makedirs(outpath) full_outfname = os.path.join(outpath, fname) if rank == 0: print('Writing %s' % full_outfname) slicemesh.save(full_outfname) if rank == 0: print('Wrote %s' % full_outfname)
def weigh_and_shift_uni_cat(delta_for_weights, displacements, Nptcles_per_dim, out_Ngrid, boxsize, internal_scale_factor_for_weights=None, out_scale_factor=None, cosmo_params=None, weighted_CIC_mode=None, uni_cat_generator='pmesh', plot_slices=False, verbose=False, return_catalog=False): """ Make uniform catalog, weigh by delta_for_weights, displace by displacements and interpolate to grid which we return. Parameters ---------- delta_for_weights : None or pmesh.pm.RealField object Particles are weighted by 1+delta_for_weights. If None, use weight=1. displacements : list [Psi_x, Psi_y, Psi_z] where Psi_i are pmesh.pm.RealField objects, holding the displacement field in different directions (on the grid). If None, do not shift. uni_cat_generator : string If 'pmesh', use pmesh for generating uniform catalog. If 'manual', use an old serial code. Returns ------- delta_shifted : FieldMesh object Density delta_shifted of shifted weighted particles (normalized to mean of 1 i.e. returning 1+delta). Returned if return_catalog=False. attrs : meshsource attrs Attrs of delta_shifted. Returned if return_catalog=False. catalog_shifted : ArrayCatalog object Shifted catalog, returned if return_catalog=True. """ comm = CurrentMPIComm.get() # ###################################################################### # Generate uniform catalog with Nptcles_per_dim^3 particles on regular # grid # ###################################################################### if uni_cat_generator == 'pmesh': # use pmesh generate_uniform_particle_grid # http://rainwoodman.github.io/pmesh/pmesh.pm.html?highlight= # readout#pmesh.pm.ParticleMesh.generate_uniform_particle_grid pmesh = ParticleMesh( BoxSize=boxsize, Nmesh=[Nptcles_per_dim, Nptcles_per_dim, Nptcles_per_dim]) ptcles = pmesh.generate_uniform_particle_grid(shift=0.0, dtype='f8') #print("type ptcles", type(ptcles), ptcles.shape) #print("head ptcles:", ptcles[:5,:]) dtype = np.dtype([('Position', ('f8', 3))]) # number of rows is given by number of ptcles on this rank uni_cat_array = np.empty((ptcles.shape[0], ), dtype=dtype) uni_cat_array['Position'] = ptcles uni_cat = ArrayCatalog( uni_cat_array, comm=None, BoxSize=boxsize * np.ones(3), Nmesh=[Nptcles_per_dim, Nptcles_per_dim, Nptcles_per_dim]) print("%d: local Nptcles=%d, global Nptcles=%d" % (comm.rank, uni_cat.size, uni_cat.csize)) del ptcles del uni_cat_array elif uni_cat_generator == 'manual': # Old serial code to generate regular grid and scatter across ranks. if comm.rank == 0: # Code copied from do_rec_v1.py and adopted # Note that nbkit UniformCatalog is random catalog, but we want a catalog # where each ptcle sits at grid points of a regular grid. # This is what we call 'regular uniform' catalog. Np = Nptcles_per_dim dtype = np.dtype([('Position', ('f8', 3))]) # Have Np**3 particles, and each particle has position x,y,z and weight 'Weight' uni_cat_array = np.empty((Np**3, ), dtype=dtype) # x components in units such that box ranges from 0 to 1. Note dx=1/Np. #x_components_1d = np.linspace(0.0, (Np-1)*(L/float(Np)), num=Np, endpoint=True)/L x_components_1d = np.linspace(0.0, (Np - 1) / float(Np), num=Np, endpoint=True) ones_1d = np.ones(x_components_1d.shape) # Put particles on the regular grid print("%d: Fill regular uniform catalog" % comm.rank) uni_cat_array['Position'][:, 0] = np.einsum('a,b,c->abc', x_components_1d, ones_1d, ones_1d).reshape( (Np**3, )) uni_cat_array['Position'][:, 1] = np.einsum('a,b,c->abc', ones_1d, x_components_1d, ones_1d).reshape( (Np**3, )) uni_cat_array['Position'][:, 2] = np.einsum('a,b,c->abc', ones_1d, ones_1d, x_components_1d).reshape( (Np**3, )) print("%d: Done filling regular uniform catalog" % comm.rank) # in nbkit0.3 units must be in Mpc/h uni_cat_array['Position'] *= boxsize else: uni_cat_array = None # Scatter across all ranks print("%d: Scatter array" % comm.rank) from nbodykit.utils import ScatterArray uni_cat_array = ScatterArray(uni_cat_array, comm, root=0, counts=None) print("%d: Scatter array done. Shape: %s" % (comm.rank, str(uni_cat_array.shape))) # Save in ArrayCatalog object uni_cat = ArrayCatalog(uni_cat_array) uni_cat.attrs['BoxSize'] = np.ones(3) * boxsize uni_cat.attrs['Nmesh'] = np.ones(3) * Nptcles_per_dim uni_cat.attrs['Nmesh_internal'] = np.ones(3) * Nmesh_orig else: raise Exception('Invalid uni_cat_generator %s' % uni_cat_generator) ######################################################################## # Set weight of particles in uni_cat to delta (interpolated to ptcle # positions) ######################################################################## if delta_for_weights is None: # set all weights to 1 uni_cat['Mass'] = np.ones(uni_cat['Position'].shape[0]) else: # weight by delta_for_weights nbkit03_utils.interpolate_pm_rfield_to_catalog( delta_for_weights, uni_cat, catalog_column_to_save_to='Mass') print("%d: rms Mass: %g" % (comm.rank, np.sqrt(np.mean(np.array(uni_cat['Mass'])**2)))) # optionally plot weighted uniform cat before shifting if plot_slices: # paint the original uni_cat to a grid and plot slice import matplotlib.pyplot as plt tmp_meshsource = uni_cat.to_mesh(Nmesh=out_Ngrid, value='Mass', window='cic', compensated=False, interlaced=False) # paint to get delta(a_internal) tmp_outfield = tmp_meshsource.paint(mode='real') # linear rescale factor from internal_scale_factor_for_weights to # out_scale_factor rescalefac = nbkit03_utils.linear_rescale_fac( internal_scale_factor_for_weights, out_scale_factor, cosmo_params=cosmo_params) tmp_outfield = 1.0 + rescalefac * (tmp_outfield - 1.0) tmp_mesh = FieldMesh(tmp_outfield) plt.imshow(tmp_mesh.preview(Nmesh=32, axes=(0, 1))) if comm.rank == 0: plt_fname = 'inmesh_Np%d_Nm%d_Ng%d.pdf' % (Nptcles_per_dim, Nmesh_orig, out_Ngrid) plt.savefig(plt_fname) print("Made %s" % plt_fname) del tmp_meshsource, rescalefac, tmp_outfield, tmp_mesh # ###################################################################### # Shift uniform catalog particles by Psi (changes uni_cat) # ###################################################################### nbkit03_utils.shift_catalog_by_psi_grid( cat=uni_cat, in_displacement_rfields=displacements, pos_column='Position', pos_units='Mpc/h', displacement_units='Mpc/h', boxsize=boxsize, verbose=verbose) #del Psi_rfields if return_catalog: # return shifted catalog return uni_cat else: # return density of shifted catalog, delta_shifted # ###################################################################### # paint shifted catalog to grid, using field_to_shift as weights # ###################################################################### print("%d: paint shifted catalog to grid using mass weights" % comm.rank) # this gets 1+delta if weighted_CIC_mode == 'sum': delta_shifted, attrs = paint_utils.weighted_paint_cat_to_delta( uni_cat, weight='Mass', Nmesh=out_Ngrid, weighted_paint_mode=weighted_CIC_mode, normalize=True, # compute 1+delta verbose=verbose, to_mesh_kwargs={ 'window': 'cic', 'compensated': False, 'interlaced': False }) # this get rho elif weighted_CIC_mode == 'avg': delta_shifted, attrs = paint_utils.mass_avg_weighted_paint_cat_to_rho( uni_cat, weight='Mass', Nmesh=out_Ngrid, verbose=verbose, to_mesh_kwargs={ 'window': 'cic', 'compensated': False, 'interlaced': False }) else: raise Exception('Invalid weighted_CIC_mode %s' % weighted_CIC_mode) # ###################################################################### # rescale to output redshift # ###################################################################### if internal_scale_factor_for_weights != out_scale_factor: # linear rescale factor from internal_scale_factor_for_weights to # out_scale_factor rescalefac = nbkit03_utils.linear_rescale_fac( internal_scale_factor_for_weights, out_scale_factor, cosmo_params=cosmo_params) delta_shifted *= rescalefac # print some info: if comm.rank == 0: print( "%d: Linear rescalefac from a=%g to a=%g, rescalefac=%g" % (comm.rank, internal_scale_factor_for_weights, out_scale_factor, rescalefac)) raise Exception( 'Check if rescaling of delta_shifted is correct. Looks like 1+delta.' ) if verbose: print("%d: delta_shifted: min, mean, max, rms(x-1):" % comm.rank, np.min(delta_shifted), np.mean(delta_shifted), np.max(delta_shifted), np.mean((delta_shifted - 1.)**2)**0.5) # get 1+deta mesh from field #outmesh = FieldMesh(1 + out_delta) # print some info: this makes code never finish (race condition maybe?) #nbkit03_utils.rfield_print_info(outfield, comm, 'outfield: ') return delta_shifted, attrs
class Data_Generation(): log=classlogger def __init__(self, parameters_box, parameters_igm, use_baryonic=False, transfer="EisensteinHu", path_transfer=None, column=None, cosmo=None): self.redshift= parameters_box['redshift'] self.width=parameters_box['width'] chosen_cosmo=parameters_box['cosmology'] #Use the baryonic power spectrum instead #Then the resulting overdensity field does not correspond to cdm-density, but to baryonic density #Jeans length needed [h**-1 Mpc] #Note that Jeans length scales with redshift #For proper values look: Choudhury, Sraianda ... 2001 self.use_baryonic=use_baryonic self.Jeans_length=parameters_box['jeans_length'] #specifies specific cosmological model if chosen_cosmo=='Planck15': self.cosmo = cosmology.Planck15 elif chosen_cosmo=='UserDefined': self.cosmo = cosmo else: print(chosen_cosmo, 'not implemented right now') print('Continue with Planck15 cosmology instead') self.cosmo = cosmology.Planck15 self.comoving_distance=self.cosmo.comoving_distance(self.redshift) #linear power spectrum Plin = LinearPower(self.cosmo, self.redshift, transfer, path_transfer, column) if self.use_baryonic: self.Plin=lambda k: 1/(1+self.Jeans_length**2*k**2)**2*Plin(k) self.Plin.redshift=self.redshift self.Plin.sigma8=self.cosmo.sigma8 self.Plin.cosmo=self.cosmo if self.log.isEnabledFor(logging.INFO): self.log.info('Start computation of baryonic density field') else: self.Plin=Plin if self.log.isEnabledFor(logging.INFO): self.log.info('Start computation of cdm density field') #Parameters describing the box self.background_field_x=parameters_box['background_box'][0] self.background_field_y=parameters_box['background_box'][1] self.background_sampling_x=parameters_box['background_sampling'][0] self.background_sampling_y=parameters_box['background_sampling'][1] if self.log.isEnabledFor(logging.INFO): self.log.info('Background has been specified:') self.log.info('-->Comoving length: {}'.format([self.background_field_x, self.background_field_y])) self.log.info('-->sampled:{}'.format([self.background_sampling_x, self.background_sampling_y])) #initializes pseudo-gaussian random generator self.seed=parameters_box['seed'] #if True, the seed gaussian has unitary_amplitude self.unitary_amplitude=True self.inverted_phase=False #Also compute displacement self.compute_displacement=True self.BoxSize=[self.background_field_x, self.background_field_y, self.width] self.N_pix=parameters_box['nr_pix'] self.bias=2 #The following values describe the state of the IGM and have to match whatever specified in the inversion self.beta=parameters_igm['beta'] self.alpha=2-1.4*self.beta self.tildeC=parameters_igm['tilde_c'] if self.log.isEnabledFor(logging.INFO): self.log.info('Generator initialized with parameters') self.log.info('-->Cosmology: {}'.format(chosen_cosmo)) self.log.info('-->Box size: {}'.format(self.BoxSize)) self.log.info('-->Number pixels: {}'.format(self.N_pix)) self.log.info('-->Redshift: {}'.format(self.redshift)) self.log.info('-->seed: {}'.format(self.seed)) self.log.info('-->Jeans: {}'.format(self.Jeans_length)) # the particle mesh for gridding purposes _Nmesh = np.empty(3, dtype='i8') _Nmesh[0] = self.background_sampling_x _Nmesh[1] = self.background_sampling_y _Nmesh[2] = self.N_pix self.pm = ParticleMesh(BoxSize=self.BoxSize, Nmesh=_Nmesh, dtype='f4') return #Compute Density Field. The function strongly matches the implementation in nbodykit. #Return: #delta: Density Perturbation (rho/rho_med) #comoving: Comoving distance to each point at line of sight [h^-1 Mpc] #vel: Peculiar velocity field [km/s] def Compute_Density_Field(self): # growth rate to do RSD in the Zel'dovich approximation f = self.cosmo.scale_independent_growth_rate(self.redshift) if self.log.isEnabledFor(logging.INFO): self.log.info('Growth rate is {}'.format(f)) #Lineat density and displacement field delta, disp = gaussian_real_fields(self.pm, self.Plin, self.seed, unitary_amplitude=self.unitary_amplitude, inverted_phase=self.inverted_phase, compute_displacement=self.compute_displacement) if self.log.isEnabledFor(logging.INFO): self.log.info('Gaussian field generated') # apply the lognormal transformation to the initial conditions density # this creates a positive-definite delta (necessary for Poisson sampling) lagrangian_bias = self.bias - 1 delta = lognormal_transform(delta, bias=lagrangian_bias) if self.log.isEnabledFor(logging.INFO): self.log.info('Density field projected to lognormal') #Compute peculiar velocity field from linear displacement field if self.compute_displacement: self.displacement=disp velocity_norm = f * 100 * self.cosmo.efunc(self.redshift) / (1+self.redshift) vel_x = velocity_norm * disp[0] vel_y = velocity_norm * disp[1] vel_z = velocity_norm * disp[2] vel=[vel_x.value, vel_y.value, vel_z.value] if self.log.isEnabledFor(logging.INFO): self.log.info('Velocity field computed') #Computes comoving distances along line of sight comoving=self.cosmo.comoving_distance(self.redshift)-self.width+self.width/self.N_pix*np.linspace(1, self.N_pix, self.N_pix) if self.log.isEnabledFor(logging.INFO): self.log.info('Density field succesfully computed') return [delta, vel, comoving] #Computes the Covariance matrix, depending on the linear autocorrelation #Correlation may look different in nonlinear and mildly nonlinear regime (due to lognormal transform) #In case of non biased gaussian fields (expectation value equal to zero) the output matches the covariance matrix #diff: The difference value between two points in real space #Autocorrelation is given by the inverse Fourier transform of power spectrum (Wiener-Khinchin-Theorem) #comoving: vector of comoving coordinates on which the density field is evaluated #WARNING: Due to rebinning the comoving field could have changed and is not equal binned anymore #As Covariance matrix is symmetric, only the upper half is computed, then both merged together by transposition def Compute_Linear_Covariance(self, comoving): if self.log.isEnabledFor(logging.INFO): self.log.info('Computation of prior covariance started') Covariance=np.zeros((np.size(comoving), np.size(comoving))) k = np.linspace(10**(-5), 10, 10**6) size = len(k)//2 Pk = self.Plin(k) fourier_coeff = np.abs(np.fft.fftn(Pk)[0:size+1]) frqs = np.linspace(0, 0.1*size, size+1) cf_lin = Spline(frqs, fourier_coeff) diff=np.zeros((np.size(comoving), np.size(comoving))) for i in range(np.size(comoving)): for j in range(np.size(comoving)): diff[i, j]=np.abs(comoving[i]-comoving[j]) Covariance=cf_lin(diff) Covariance /= Covariance[0, 0] if self.log.isEnabledFor(logging.INFO): self.log.info('Prior Covariance computed') return Covariance #This Function computes the selection of a single line of sight from the delta sample #i, j: indices of background sources at maximal redshift #Return: Density-field along one line-of-sight, evaluated at the pixels of comoving def Select_LOS(self, delta, vel, index): i=index[0].astype(int) j=index[1].astype(int) if self.log.isEnabledFor(logging.INFO): self.log.info('Line of sight selected: {}'.format(index)) return [delta[i, j, :], vel[i, j, :]] #Poisson sample to overdensity field, #matches the nbodykit routine def PoissonSample(self, delta, parameters_sampling): nbar=parameters_sampling['nbar'] seed1=parameters_sampling['seed1'] seed2=parameters_sampling['seed2'] comm = self.pm.comm # mean number of objects per cell H = self.BoxSize / self.pm.Nmesh overallmean = H.prod() * nbar # number of objects in each cell (per rank, as a RealField) cellmean = delta * overallmean # create a random state with the input seed rng = MPIRandomState(seed=seed1, comm=comm, size=delta.size) # generate poissons. Note that we use ravel/unravel to # maintain MPI invariane. Nravel = rng.poisson(lam=cellmean.ravel()) N = self.pm.create(type='real') N.unravel(Nravel) Ntot = N.csum() if self.log.isEnabledFor(logging.INFO): self.log.info('Poisson sampling done, total number of objects is {}'.format(Ntot)) pos_mesh = self.pm.generate_uniform_particle_grid(shift=0.0) disp_mesh = np.empty_like(pos_mesh) # no need to do decompose because pos_mesh is strictly within the # local volume of the RealField. N_per_cell = N.readout(pos_mesh, resampler='nnb') for i in range(N.ndim): disp_mesh[:, i] = self.displacement[i].readout(pos_mesh, resampler='nnb') # fight round off errors, if any N_per_cell = np.int64(N_per_cell + 0.5) pos = pos_mesh.repeat(N_per_cell, axis=0) disp = disp_mesh.repeat(N_per_cell, axis=0) del pos_mesh del disp_mesh if self.log.isEnabledFor(logging.INFO): self.log.info("Catalog produced. Assigning in cell shift.") # FIXME: after pmesh update, remove this orderby = np.int64(pos[:, 0] / H[0] + 0.5) for i in range(1, delta.ndim): orderby[...] *= self.pm.Nmesh[i] orderby[...] += np.int64(pos[:, i] / H[i] + 0.5) # sort by ID to maintain MPI invariance. pos = mpsort.sort(pos, orderby=orderby, comm=comm) disp = mpsort.sort(disp, orderby=orderby, comm=comm) if self.log.isEnabledFor(logging.INFO): self.log.info("Sorting done") rng_shift = MPIRandomState(seed=seed2, comm=comm, size=len(pos)) in_cell_shift = rng_shift.uniform(0, H[i], itemshape=(delta.ndim,)) pos[...] += in_cell_shift pos[...] %= self.pm.BoxSize if self.log.isEnabledFor(logging.INFO): self.log.info("Catalog shifted.") #Catalog needs to be shifted in z-coordinate, such that pos and comoving match pos[...,0]+=(self.comoving_distance-self.width)*np.ones(pos.shape[0]) return pos, disp #Projects baryonic density field to neutral hydrogen overdensity field def Find_Neutral_Hydrogen_Fraction(self, density, redshift): density_h1=self.tildeC*(1+redshift)**6*density**(self.alpha) if self.log.isEnabledFor(logging.INFO): self.log.info("Density field projected to neutral hydrogen density field") return density_h1 def Project_Velocity(self, vel, direction): return vel[direction]
def look_den_slice(net1, net2, net3, s, title): matplotlib.rc('xtick', labelsize=12) matplotlib.rc('ytick', labelsize=12) matplotlib.rc('font', size=12) cmap = 'coolwarm' assert net1.shape[-1] == 3 pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32]) q = pm.generate_uniform_particle_grid() net1 = pm.paint(q + net1.reshape([-1, 3])) pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32]) net2 = pm.paint(q + net2.reshape([-1, 3])) pm = ParticleMesh(BoxSize=128, Nmesh=[32, 32, 32]) net3 = pm.paint(q + net3.reshape([-1, 3])) fig = plt.figure(figsize=(6, 4)) amp_low = min(np.percentile(net1[:, :, s], 5), np.percentile(net2[:, :, s], 5), np.percentile(net3[:, :, s], 5)) amp_high = max(np.percentile(net1[:, :, s], 95), np.percentile(net2[:, :, s], 95), np.percentile(net3[:, :, s], 95)) amp_low1 = min(np.percentile((net1[:, :, s] - net2[:, :, s]), 5), np.percentile((net1[:, :, s] - net3[:, :, s]), 5)) amp_high1 = max(np.percentile((net1[:, :, s] - net2[:, :, s]), 95), np.percentile((net1[:, :, s] - net3[:, :, s]), 95)) amp = max(np.abs(amp_low1), amp_high1) gs = GridSpec(2, 4, height_ratios=[1, 1], width_ratios=[1, 1, 1, 0.1]) plt.subplot(gs[0, 0]) im = plt.imshow(net1[:, :, s], cmap=cmap, vmin=amp_low, vmax=amp_high) #plt.xticks([0,16,32],[0,64,128]) #plt.yticks([16,0],[64,128]) plt.axis('off') plt.title('fastPM') plt.subplot(gs[0, 1]) plt.imshow(net2[:, :, s], cmap=cmap, vmin=amp_low, vmax=amp_high) #plt.xticks([0,16,32],[0,64,128]) #plt.yticks([16,0],[64,128]) plt.axis('off') plt.title('2LPT') plt.subplot(gs[0, 2]) plt.imshow(net3[:, :, s], cmap=cmap, vmin=amp_low, vmax=amp_high) #plt.xticks([0,16,32],[0,64,128]) #plt.yticks([16,0],[64,128]) plt.axis('off') plt.title('U-Net') cbax = plt.subplot(gs[0, 3]) cbar = fig.colorbar(mappable=im, cax=cbax, orientation='vertical', ticklocation='right') cbar.ax.tick_params(labelsize=12) plt.subplot(gs[1, 1]) im = plt.imshow((net1[:, :, s] - net2[:, :, s]), cmap=cmap, vmin=-amp, vmax=amp) #plt.xticks([0,16,32],[0,64,128]) #plt.yticks([16,0],[64,128]) plt.axis('off') plt.title(r'fastPM $-$ 2LPT') plt.subplot(gs[1, 2]) plt.imshow((net1[:, :, s] - net3[:, :, s]), cmap=cmap, vmin=-amp, vmax=amp) #plt.xticks([0,16,32],[0,64,128]) #plt.yticks([16,0],[64,128]) plt.axis('off') plt.title(r'fastPM $-$ U-Net') cbax = plt.subplot(gs[1, 3]) cbar = fig.colorbar(mappable=im, cax=cbax, orientation='vertical', ticklocation='right') cbar.ax.tick_params(labelsize=12) fig.subplots_adjust(hspace=.2)
def saveIC_bt2(IC_path, dx_field, Lbox, Ng, cosmology, redshift=99): """ Use Zel-dovich approximation to back-scale the linear density field to initial redshift, paint baryon and dm particles from the grid, and save initial condition in MP-Gadget/bluetides-ii format Parameters --------- : dx_field : linear density field at z=0 : Lbox : BoxSize, in Mpc/h, will be converted to kpc/h in the IC output : Ng : Number of grid on which to paint the particle : cosmology : nbodykit.cosmology, or astropy.cosmology.FLRW : redshift : redshift of the initial condition """ mesh = ArrayMesh(dx_field, BoxSize=Lbox) # density contrast field centered at zero dk_field = mesh.compute(mode='complex') shift_gas = -0.5 * (cosmology.Om0 - cosmology.Ob0) / cosmology.Om0 shift_dm = 0.5 * cosmology.Ob0 / cosmology.Om0 pm = ParticleMesh(BoxSize=Lbox, Nmesh=[Ng, Ng, Ng]) solver = Solver(pm, cosmology, B=1) Q_gas = pm.generate_uniform_particle_grid(shift=shift_gas) Q_dm = pm.generate_uniform_particle_grid(shift=shift_dm) scale_a = 1. / (1 + redshift) state_gas = solver.lpt(dk_field, Q_gas, a=scale_a, order=1) # order = 1 : Zel-dovich state_dm = solver.lpt(dk_field, Q_dm, a=scale_a, order=1) state = state_dm with FileMPI(state.pm.comm, IC_path, create=True) as ff: m0 = state.cosmology.rho_crit(0) * state.cosmology.Omega0_b * ( Lbox**3) / state.csize m1 = state.cosmology.rho_crit(0) * ( state.cosmology.Om0 - state.cosmology.Omega0_b) * (Lbox** 3) / state.csize with ff.create('Header') as bb: bb.attrs['BoxSize'] = Lbox * 1000 bb.attrs['HubbleParam'] = state.cosmology.h bb.attrs['MassTable'] = [m0, m1, 0, 0, 0, 0] bb.attrs['OmegaM'] = state.cosmology.Om0 bb.attrs['OmegaB'] = state.cosmology.Omega0_b bb.attrs['OmegaL'] = state.cosmology.Omega0_lambda bb.attrs['Time'] = state.a['S'] bb.attrs['TotNumPart'] = [state.csize, state.csize, 0, 0, 0, 0] ff.create_from_array('1/Position', 1000 * periodic_wrap(state.X, Lbox)) # in kpc/h ff.create_from_array( '1/Velocity', state.V / np.sqrt(state.a['S'])) # old gadget convention for IC dmID = np.arange(state.csize) ff.create_from_array('1/ID', dmID) ####################################################### ff.create_from_array('0/Position', 1000 * periodic_wrap(state_gas.X, Lbox)) ff.create_from_array('0/Velocity', state_gas.V / np.sqrt(state.a['S'])) gasID = np.arange(state.csize, 2 * state.csize) ff.create_from_array('0/ID', gasID) print("IC generated!") print("*********************************************") return state
rank = pm.comm.rank wsize = pm.comm.size comm = pm.comm if rank == 0: print(args) print(outfolder) cube_size = nc / ncube shift = cube_size cube_length = cube_size * bs / nc pmsmall = ParticleMesh(BoxSize=bs / ncube, Nmesh=[cube_size, cube_size, cube_size], dtype=np.float32, comm=MPI.COMM_SELF) gridsmall = pmsmall.generate_uniform_particle_grid(shift=0) for zz in [3.5, 4.0]: aa = 1 / (1 + zz) cats, meshes = setupuvmesh(zz, suff=suff, sim=sim, pm=pm, profile=profile, stellar=stellar) cencat, satcat = cats h1meshfid, h1mesh, lmesh, uvmesh = meshes if ncsim == 10240: dm = BigFileMesh(scratchyf+sim+'/fastpm_%0.4f/'%aa+\ '/1-mesh/N%04d'%nc,'').paint() else: dm = BigFileMesh(project+sim+'/fastpm_%0.4f/'%aa+\
def func_gal_catalogue(bs, nc, seed, nstep, seed_hod, Omega_m, p_alpha, p_logMin, p_logM1, p_logM0, p_sigma_logM): folder = "L%04d_N%04d_S%04d_%02dstep" % (bs, nc, seed, nstep) # setup initial conditions Omegacdm = Omega_m - 0.049, cosmo = cosmology.Planck15.clone(Omega_cdm=Omegacdm, h=0.6711, Omega_b=0.049) power = cosmology.LinearPower(cosmo, 0) klin = np.logspace(-4, 2, 1000) plin = power(klin) pkfunc = interpolate(klin, plin) # run the simulation pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc]) Q = pm.generate_uniform_particle_grid() stages = numpy.linspace(0.1, 1.0, nstep, endpoint=True) solver = Solver(pm, cosmo, B=2) wn = solver.whitenoise(seed) dlin = solver.linear(wn, pkfunc) state = solver.lpt(dlin, Q, stages[0]) state = solver.nbody(state, leapfrog(stages)) # create the catalogue cat = ArrayCatalog( { 'Position': state.X, 'Velocity': state.V, 'Displacement': state.S, 'Density': state.RHO }, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh, M0=Omega_m * 27.75e10 * bs**3 / (nc / 2.0)**3) cat['KDDensity'] = KDDensity(cat).density cat.save('%s/Matter' % (folder), ('Position', 'Velocity', 'Density', 'KDDensity')) # run FOF fof = FOF(cat, linking_length=0.2, nmin=12) fofcat = fof.to_halos(particle_mass=cat.attrs['M0'], cosmo=cosmo, redshift=0.0) fofcat.save('%s/FOF' % (folder), ('Position', 'Velocity', 'Mass', 'Radius')) # run HOD params = { 'alpha': p_alpha, 'logMmin': p_logMin, 'logM1': p_logM1, 'logM0': p_logM0, 'sigma_logM': p_sigma_logM } halos = HaloCatalog(fofcat, cosmo=cosmo, redshift=0.0, mdef='vir') halocat = halos.to_halotools(halos.attrs['BoxSize']) hod = HODCatalog(halocat, seed=seed_hod, **params) hod.save('%s/HOD' % (folder), ('Position', 'Velocity')) return folder, cat, fofcat, hod