def createmesh(self, bs, nc, positions, weights): '''use this to create mesh of HI ''' pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc]) mesh = pm.create(mode='real', value=0) comm = pm.comm rankweight = sum([wt.sum() for wt in weights]) totweight = comm.allreduce(rankweight) for wt in weights: wt /= totweight / float(nc)**3 lay = pm.decompose(positions[0]) mesh.paint(positions[0], mass=weights[0], layout=lay, hold=True) if len(positions) > 1: for i in range(self.nsat): shift = np.random.normal(0, positions[1]) pos = positions[0] + shift lay = pm.decompose(pos) mesh.paint(pos, mass=weights[1] / self.nsat, layout=lay, hold=True) return mesh
def test_inplace_fft(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') real = RealField(pm) numpy.random.seed(1234) if comm.rank == 0: Npar = 100 else: Npar = 0 pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape( -1, len(pm.Nmesh)) * (7, 7) pos %= (pm.Nmesh + 1) layout = pm.decompose(pos) npos = layout.exchange(pos) real = pm.paint(npos) complex = real.r2c() complex2 = real.r2c(out=Ellipsis) assert real._base in complex2._base assert_almost_equal(numpy.asarray(complex), numpy.asarray(complex2), decimal=7) real = complex2.c2r() real2 = complex2.c2r(out=Ellipsis) assert real2._base in complex2._base assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_c2c(comm): # this test requires pfft-python 0.1.16. pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='complex128') numpy.random.seed(1234) if comm.rank == 0: Npar = 100 else: Npar = 0 pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(-1, len(pm.Nmesh)) * (7, 7) pos %= (pm.Nmesh + 1) layout = pm.decompose(pos) npos = layout.exchange(pos) real = pm.paint(npos) complex = real.r2c() real2 = complex.c2r() assert numpy.iscomplexobj(real) assert numpy.iscomplexobj(real2) assert numpy.iscomplexobj(complex) assert_array_equal(complex.cshape, pm.Nmesh) assert_array_equal(real2.cshape, pm.Nmesh) assert_array_equal(real.cshape, pm.Nmesh) real.readout(npos) assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_c2c(comm): # this test requires pfft-python 0.1.16. pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='complex128') numpy.random.seed(1234) if comm.rank == 0: Npar = 100 else: Npar = 0 pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape( -1, len(pm.Nmesh)) * (7, 7) pos %= (pm.Nmesh + 1) layout = pm.decompose(pos) npos = layout.exchange(pos) real = pm.paint(npos) complex = real.r2c() real2 = complex.c2r() assert numpy.iscomplexobj(real) assert numpy.iscomplexobj(real2) assert numpy.iscomplexobj(complex) assert_array_equal(complex.cshape, pm.Nmesh) assert_array_equal(real2.cshape, pm.Nmesh) assert_array_equal(real.cshape, pm.Nmesh) real.readout(npos) assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def test_inplace_fft(comm): pm = ParticleMesh(BoxSize=8.0, Nmesh=[8, 8], comm=comm, dtype='f8') real = RealField(pm) numpy.random.seed(1234) if comm.rank == 0: Npar = 100 else: Npar = 0 pos = 1.0 * (numpy.arange(Npar * len(pm.Nmesh))).reshape(-1, len(pm.Nmesh)) * (7, 7) pos %= (pm.Nmesh + 1) layout = pm.decompose(pos) npos = layout.exchange(pos) real = pm.paint(npos) complex = real.r2c() complex2 = real.r2c(out=Ellipsis) assert real._base in complex2._base assert_almost_equal(numpy.asarray(complex), numpy.asarray(complex2), decimal=7) real = complex2.c2r() real2 = complex2.c2r(out=Ellipsis) assert real2._base in complex2._base assert_almost_equal(numpy.asarray(real), numpy.asarray(real2), decimal=7)
def make_galaxy_pk(scale_factor,nc,seed,bs=1536,T=40,B=2,simpath=sc_simpath,outpath=sc_outpath,Rsm=0): aa = scale_factor # since particle IDs are ordered only need to load at one redshift zz = 1/aa-1 dgrow = cosmo.scale_independent_growth_factor(zz) fname = get_filename(nc,seed,T=T,B=B) spath = simpath + fname opath = outpath + fname galpath = opath + '/hod/' try: os.makedirs(opath+'spectra/') except : pass zz = 1/aa-1 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8') rank = pm.comm.rank dyn = BigFileCatalog(spath + '/fastpm_%0.4f/1'%aa) # Load Matter Mesh fpos = dyn['Position'].compute() mlay = pm.decompose(fpos) mmesh = pm.paint(fpos, layout=mlay) mmesh /= mmesh.cmean() # Load halo mesh: HOD parameters fixed for now... mmin = 10**12.5; m1fac = 20 cendir ='cencat-aa-%.04f-Mmin-%.1f-M1f-%.1f-alpha-0p8-subvol'%(aa,np.log10(mmin), m1fac) satdir ='satcat-aa-%.04f-Mmin-%.1f-M1f-%.1f-alpha-0p8-subvol'%(aa,np.log10(mmin), m1fac) cencat = BigFileCatalog(galpath + cendir) satcat = BigFileCatalog(galpath + satdir) hpos = np.concatenate((cencat['Position'],satcat['Position'])) hlay = pm.decompose(hpos) hmesh = pm.paint(hpos, layout=hlay) hmesh /= hmesh.cmean() phm = FFTPower(mmesh, second=hmesh, mode='1d').power k, phm = phm['k'], phm['power'].real phh = FFTPower(hmesh, mode='1d').power k, phh = phh['k'], phh['power'].real np.savetxt(opath+'spectra/pks-%04d-%04d-%04d-gal.txt'%(aa*10000, bs, nc), np.vstack([k, phh, phm]).T.real, header='k, phh, phm/ ', fmt='%0.4e')
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 savecatalogmesh(bs, nc, aa): dmcat = BigFileCatalog(scratchyf + sim + '/fastpm_%0.4f/' % aa, dataset='1') pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc]) pos = dmcat['Position'].compute() layout = pm.decompose(pos) mesh = pm.paint(pos, layout=layout) mesh = FieldMesh(mesh) path = project + sim + '/fastpm_%0.4f/' % aa + '/dmesh_N%04d' % nc mesh.save(path, dataset='1', mode='real')
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 createmesh(self, bs, nc, positions, weights): '''use this to create mesh of HI ''' pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc]) mesh = pm.create(mode='real', value=0) comm = pm.comm rankweight = sum([wt.sum() for wt in weights]) totweight = comm.allreduce(rankweight) for wt in weights: wt /= totweight / float(nc)**3 for i in range(len(positions)): lay = pm.decompose(positions[i]) mesh.paint(positions[i], mass=weights[i], layout=lay, hold=True) return mesh
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_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)
class Subsample(Algorithm): """ Algorithm to create a subsample from a DataSource, and evaluate the density (1 + delta), smoothed at the given scale """ plugin_name = "Subsample" def __init__(self, datasource, Nmesh, seed=12345, ratio=0.01, smoothing=None, format="hdf5"): from pmesh.pm import ParticleMesh self.datasource = datasource self.Nmesh = Nmesh self.seed = seed self.ratio = ratio self.smoothing = smoothing self.format = format self.pm = ParticleMesh(BoxSize=self.datasource.BoxSize, Nmesh=[self.Nmesh] * 3, dtype="f4", comm=self.comm) @classmethod def fill_schema(cls): s = cls.schema s.description = "create a subsample from a DataSource, and evaluate density \n" s.description += "(1 + delta) smoothed at the given scale" s.add_argument( "datasource", type=DataSource.from_config, help="the DataSource to read; run `nbkit.py --list-datasources` for all options", ) s.add_argument("Nmesh", type=int, help="the size of FFT mesh for painting") s.add_argument("seed", type=int, help="the random seed") s.add_argument("ratio", type=float, help="the fraction of particles to keep") s.add_argument( "smoothing", type=float, help="the smoothing length in distance units. " "It has to be greater than the mesh resolution. " "Otherwise the code will die. Default is the mesh resolution.", ) # this is for output.. s.add_argument("format", choices=["hdf5", "mwhite"], help="the format of the output") def run(self): """ Run the Subsample algorithm """ import mpsort from astropy.utils.misc import NumpyRNGContext if self.smoothing is None: self.smoothing = self.datasource.BoxSize[0] / self.Nmesh[0] elif (self.datasource.BoxSize / self.Nmesh > self.smoothing).any(): raise ValueError("smoothing is too small") def Smoothing(pm, complex): k = pm.k k2 = 0 for ki in k: ki2 = ki ** 2 complex[:] *= numpy.exp(-0.5 * ki2 * self.smoothing ** 2) def NormalizeDC(pm, complex): """ removes the DC amplitude. This effectively divides by the mean """ w = pm.w comm = pm.comm ind = [] value = 0.0 found = True for wi in w: if (wi != 0).all(): found = False break ind.append((wi == 0).nonzero()[0][0]) if found: ind = tuple(ind) value = numpy.abs(complex[ind]) value = comm.allreduce(value) complex[:] /= value # open the datasource and keep the cache with self.datasource.keep_cache(): painter = Painter.create("DefaultPainter", paintbrush="cic") real, stats = painter.paint(self.pm, self.datasource) complex = real.r2c() for t in [Smoothing, NormalizeDC]: t(self.pm, complex) complex.c2r(real) columns = ["Position", "ID", "Velocity"] local_seed = utils.local_random_seed(self.seed, self.comm) dtype = numpy.dtype([("Position", ("f4", 3)), ("Velocity", ("f4", 3)), ("ID", "u8"), ("Density", "f4")]) subsample = [numpy.empty(0, dtype=dtype)] with self.datasource.open() as stream: for Position, ID, Velocity in stream.read(columns): with NumpyRNGContext(local_seed): u = numpy.random.uniform(size=len(ID)) keep = u < self.ratio Nkeep = keep.sum() if Nkeep == 0: continue data = numpy.empty(Nkeep, dtype=dtype) data["Position"][:] = Position[keep] data["Velocity"][:] = Velocity[keep] data["ID"][:] = ID[keep] layout = self.pm.decompose(data["Position"]) pos1 = layout.exchange(data["Position"]) density1 = real.readout(pos1) density = layout.gather(density1) # normalize the position after reading out density! data["Position"][:] /= self.datasource.BoxSize data["Velocity"][:] /= self.datasource.BoxSize data["Density"][:] = density subsample.append(data) subsample = numpy.concatenate(subsample) mpsort.sort(subsample, orderby="ID", comm=self.comm) return subsample def save(self, output, data): if self.format == "mwhite": self.write_mwhite_subsample(data, output) else: self.write_hdf5(data, output) def write_hdf5(self, subsample, output): import h5py size = self.comm.allreduce(len(subsample)) offset = sum(self.comm.allgather(len(subsample))[: self.comm.rank]) if self.comm.rank == 0: with h5py.File(output, "w") as ff: dataset = ff.create_dataset(name="Subsample", dtype=subsample.dtype, shape=(size,)) dataset.attrs["Ratio"] = self.ratio dataset.attrs["CommSize"] = self.comm.size dataset.attrs["Seed"] = self.seed dataset.attrs["Smoothing"] = self.smoothing dataset.attrs["Nmesh"] = self.Nmesh dataset.attrs["Original"] = self.datasource.string dataset.attrs["BoxSize"] = self.datasource.BoxSize for i in range(self.comm.size): self.comm.barrier() if i != self.comm.rank: continue with h5py.File(output, "r+") as ff: dataset = ff["Subsample"] dataset[offset : len(subsample) + offset] = subsample def write_mwhite_subsample(self, subsample, output): size = self.comm.allreduce(len(subsample)) offset = sum(self.comm.allgather(len(subsample))[: self.comm.rank]) if self.comm.rank == 0: with open(output, "wb") as ff: dtype = numpy.dtype( [ ("eflag", "int32"), ("hsize", "int32"), ("npart", "int32"), ("nsph", "int32"), ("nstar", "int32"), ("aa", "float"), ("gravsmooth", "float"), ] ) header = numpy.zeros((), dtype=dtype) header["eflag"] = 1 header["hsize"] = 20 header["npart"] = size header.tofile(ff) self.comm.barrier() with open(output, "r+b") as ff: ff.seek(28 + offset * 12) numpy.float32(subsample["Position"]).tofile(ff) ff.seek(28 + offset * 12 + size * 12) numpy.float32(subsample["Velocity"]).tofile(ff) ff.seek(28 + offset * 4 + size * 24) numpy.float32(subsample["Density"]).tofile(ff) ff.seek(28 + offset * 8 + size * 28) numpy.uint64(subsample["ID"]).tofile(ff)
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)
FieldMesh(LDLmap).save(save) elif args.target == 'kSZ': LDLmap *= 1e-5 #The learned LDL map is actually 1e5*map_ne. save = args.save + '/ne_snap' + str(args.snapNum).zfill(3) + '_Nmesh%d_Nstep%d_n%.2f_map' % (args.Nmesh, args.Nstep, args.n) FieldMesh(LDLmap).save(save) if args.FastPMpath: X = File(args.FastPMpath)['Position'] X = np.array(X[start:end]).astype(np.float32) Vz = File(args.FastPMpath)['Velocity'] Vz = np.array(Vz[start:end])[:,2].astype(np.float32) else: from readTNG import scalefactor a = scalefactor(args.TNGDarkpath, args.snapNum) Vz = load_TNG_data(TNG_basepath=args.TNGDarkpath, snapNum=args.snapNum, partType='dm', field='Velocities', mdi=2) * a**0.5 layout = pm.decompose(X) X = layout.exchange(X) Vz = layout.exchange(Vz) map_vz = pm.create(type="real") map_vz.paint(X, mass=Vz, layout=None, hold=False) map_delta = pm.create(type="real") map_delta.paint(X, mass=1., layout=None, hold=False) select = map_delta > 0 map_vz[select] = map_vz[select] / map_delta[select] map_vz[~select] = 0 LDLmap = LDLmap * map_vz save = args.save + '/' + args.target + '_snap' + str(args.snapNum).zfill(3) + '_Nmesh%d_Nstep%d_n%.2f_map' % (args.Nmesh, args.Nstep, args.n) FieldMesh(LDLmap).save(save) elif args.target in ['tSZ_T', 'Xray_T']:
cat = ArrayCatalog({'ID': idd, 'InitPosition': grid}, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh) #cat.save(lpath + 'initpos', tosave) #cat.attrs = dyn.attrs #header = '1,b1,b2,bg,bk' header = 'b1,b2,bg,bk' names = header.split(',') #cat.save(lpath + 'dynamic/1', ('ID', 'InitPosition', 'Position')) tosave = ['ID', 'InitPosition'] + names if rank == 0: print(tosave) for i in range(len(names)): ff = BigFileMesh(lpath+ '/lag', names[i]).paint() pm = ff.pm rank = pm.comm.rank glay, play = pm.decompose(grid), pm.decompose(fpos) wts = ff.readout(grid, layout = glay, resampler='nearest') cat[names[i]] = wts #x = FieldMesh(pm.paint(fpos, mass=wts, layout=play)) #x.save(lpath + 'eul-z%03d-R%d'%(zz*100, Rsm), dataset=names[i], mode='real') del pm, ff, wts if rank == 0: print(rank, ' got weights') cat.save(lpath + 'lagweights/', tosave)
#dpath = '/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100-fixed/'%(bs, nc) aas = [ 0.5, ] for aa in aas: zz = 1 / aa - 1 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8') rank = pm.comm.rank dyn = BigFileCatalog(dpath + '/fastpm_%0.4f/1' % aa) # Load Matter Mesh fpos = dyn['Position'].compute() mlay = pm.decompose(fpos) mmesh = pm.paint(fpos, layout=mlay) mmesh /= mmesh.cmean() # Load halo mesh cencat = BigFileCatalog(galpath + '/fastpm_%0.4f/' % (aa) + cendir) satcat = BigFileCatalog(galpath + '/fastpm_%0.4f/' % (aa) + satdir) hpos = np.concatenate((cencat['Position'], satcat['Position'])) hlay = pm.decompose(hpos) hmesh = pm.paint(hpos, layout=hlay) hmesh /= hmesh.cmean() phm = FFTPower(mmesh, second=hmesh, mode='1d').power k, phm = phm['k'], phm['power'].real
class TidalTensor(Algorithm): """ Compute and save the tidal force tensor """ plugin_name = "TidalTensor" def __init__(self, field, points, Nmesh, smoothing=None): from pmesh.pm import ParticleMesh self.field = field self.points = points self.Nmesh = Nmesh self.smoothing = smoothing self.pm = ParticleMesh(BoxSize=self.field.BoxSize, Nmesh=[self.Nmesh] * 3, dtype='f4', comm=self.comm) @classmethod def fill_schema(cls): s = cls.schema s.description = "compute the tidal force tensor" s.add_argument( "field", type=DataSource.from_config, help="DataSource; run `nbkit.py --list-datasources` for all options" ) s.add_argument( "points", type=DataSource.from_config, help="A small set of points to calculate tidal force on; " "run `nbkit.py --list-datasources` for all options") s.add_argument("Nmesh", type=int, help='Size of FFT mesh for painting') s.add_argument( "smoothing", type=float, help='Smoothing Length in distance units. ' 'It has to be greater than the mesh resolution. ' 'Otherwise the code will die. Default is the mesh resolution.') def Smoothing(self, pm, complex): k = pm.k k2 = 0 for ki in k: ki2 = ki**2 complex *= numpy.exp(-0.5 * ki2 * self.smoothing**2) def NormalizeDC(self, pm, complex): """ removes the DC amplitude. This effectively divides by the mean """ w = pm.w comm = pm.comm ind = [] value = 0.0 found = True for wi in w: if (wi != 0).all(): found = False break ind.append((wi == 0).nonzero()[0][0]) if found: ind = tuple(ind) value = numpy.abs(complex[ind]) value = comm.allreduce(value) complex[:] /= value def TidalTensor(self, u, v): # k_u k_v / k **2 def TidalTensor(pm, complex): k = pm.k for row in range(complex.shape[0]): k2 = k[0][row]**2 for ki in k[1:]: k2 = k2 + ki[0]**2 k2[k2 == 0] = numpy.inf complex[row] /= k2 complex *= k[u] complex *= k[v] return TidalTensor def run(self): """ Run the TidalTensor Algorithm """ from itertools import product if self.smoothing is None: self.smoothing = self.field.BoxSize[0] / self.Nmesh elif (self.field.BoxSize / self.Nmesh > self.smoothing).any(): raise ValueError("smoothing is too small") painter = Painter.create("DefaultPainter", weight="Mass", paintbrush="cic") real, stats = painter.paint(self.pm, self.field) complex = real.r2c() for t in [self.Smoothing, self.NormalizeDC]: t(complex.pm, complex) with self.points.open() as stream: [[Position]] = stream.read(['Position'], full=True) layout = self.pm.decompose(Position) pos1 = layout.exchange(Position) value = numpy.empty((3, 3, len(Position))) for u, v in product(range(3), range(3)): if self.comm.rank == 0: self.logger.info("Working on tensor element (%d, %d)" % (u, v)) c2 = complex.copy() self.TidalTensor(u, v)(c2.pm, c2) c2.c2r(real) v1 = real.readout(pos1) v1 = layout.gather(v1) value[u, v] = v1 return value.transpose((2, 0, 1)) def save(self, output, data): self.write_hdf5(data, output) def write_hdf5(self, data, output): import h5py size = self.comm.allreduce(len(data)) offset = sum(self.comm.allgather(len(data))[:self.comm.rank]) if self.comm.rank == 0: with h5py.File(output, 'w') as ff: dataset = ff.create_dataset(name='TidalTensor', dtype=data.dtype, shape=(size, 3, 3)) dataset.attrs['Smoothing'] = self.smoothing dataset.attrs['Nmesh'] = self.Nmesh dataset.attrs['Original'] = self.field.string dataset.attrs['BoxSize'] = self.field.BoxSize for i in range(self.comm.size): self.comm.barrier() if i != self.comm.rank: continue with h5py.File(output, 'r+') as ff: dataset = ff['TidalTensor'] dataset[offset:len(data) + offset] = data
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, mode='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)
def make_component_spectra(scale_factor, nc, seed, bs=1536, T=40, B=2, simpath=sc_simpath, outpath=sc_outpath, Rsm=0): aa = scale_factor # since particle IDs are ordered only need to load at one redshift zz = 1 / aa - 1 dgrow = cosmo.scale_independent_growth_factor(zz) fname = get_filename(nc, seed, T=T, B=B) spath = simpath + fname opath = outpath + fname try: os.makedirs(opath) except: pass pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8') rank = pm.comm.rank header = '1,b1,b2,bg,bk' names = header.split(',') dfacs = [1, dgrow, dgrow**2, dgrow**2, dgrow] # dyn = BigFileCatalog(spath + '/fastpm_%0.4f/1' % aa) fpos = dyn['Position'].compute() idd = dyn['ID'] attrs = dyn.attrs play = pm.decompose(fpos) # wts = BigFileCatalog(opath + '/lagweights/') grid = wts['InitPosition'] iddg = wts['ID'] print('asserting') np.allclose(idd.compute(), iddg.compute()) print('read fields') disp = grid.compute() - fpos mask = abs(disp) > bs / 2. disp[mask] = (bs - abs(disp[mask])) * -np.sign(disp[mask]) print(rank, ' Max disp: ', disp.max()) print(rank, ' Std disp: ', disp.std(axis=0)) eul_fields = [] for i in range(len(names)): eul_fields.append( pm.paint(fpos, mass=dfacs[i] * wts[names[i]], layout=play)) print(rank, names[i], eul_fields[-1].cmean()) del dyn, fpos, idd, iddg, wts print('got fields') iz = int(100 * zz + 0.01) k, spectra = tools.getspectra(eul_fields) np.savetxt(opath + '/spectra-z%03d-R%d.txt' % (iz, Rsm), np.vstack([k, spectra]).T.real, header='k / ' + header, fmt='%0.4e')
rank = pm.comm.rank lin = BigFileMesh(dpath + '/mesh', 's').paint() dyn = BigFileCatalog(dpath + '/dynamic/1') hcat = BigFileCatalog(hpath + '/FOF/') # hpos = hcat['CMPosition'][1:num] #print('Mass : ', rank, hcat['Mass'].compute())[1:num] hmass = hcat['Mass'].compute()[1:num] if not masswt: hmass = hmass * 0 + 1. suff = suff + '-pos' else: suff = suff + '-mass' hlay = pm.decompose(hpos) hmesh = pm.paint(hpos, mass=hmass, layout=hlay) hmesh /= hmesh.cmean() ph = FFTPower(hmesh, mode='1d').power k, ph = ph['k'], ph['power'] # grid = dyn['InitPosition'].compute() fpos = dyn['Position'].compute() print(rank, (grid - fpos).std(axis=0)) # dgrow = cosmo.scale_independent_growth_factor(zz) if zadisp: fpos = za.doza(lin.r2c(), grid, z=zz, dgrow=dgrow) suff = suff + '-za'
def get_lagweights(nc, seed, bs=1536, T=40, B=2, simpath=sc_simpath, outpath=sc_outpath, dyn=None): # Note that dyn is the particle catalog, which we reload if not given aa = 1.0000 # since particle IDs are ordered only need to load at one redshift zz = 1 / aa - 1 fname = get_filename(nc, seed, T=T, B=B) spath = simpath + fname opath = outpath + fname try: os.makedirs(opath) except: pass pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f4') rank = pm.comm.rank if dyn is None: particle_file = spath + '/fastpm_%0.4f/1' % aa print("Loading particle data from: " + particle_file) dyn = BigFileCatalog(particle_file) # fpos = dyn['Position'].compute() idd = dyn['ID'].compute() attrs = dyn.attrs grid = tools.getqfromid(idd, attrs, nc) if rank == 0: print('grid computed') cat = ArrayCatalog({ 'ID': idd, 'InitPosition': grid }, BoxSize=pm.BoxSize, Nmesh=pm.Nmesh) header = '1,b1,b2,bg,bk' names = header.split(',') tosave = ['ID', 'InitPosition'] + names if rank == 0: print(tosave) for i in range(len(names)): ff = BigFileMesh(opath + '/lag', names[i]).paint() pm = ff.pm rank = pm.comm.rank glay, play = pm.decompose(grid), pm.decompose(fpos) wts = ff.readout(grid, layout=glay, resampler='nearest') cat[names[i]] = wts del pm, ff, wts if rank == 0: print(rank, ' got weights') cat.save(opath + 'lagweights/', tosave)
try: os.makedirs(ofolder) except: pass pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8') rank = pm.comm.rank header = '1,b1,b2,bg,bk' names = header.split(',') # dyn = BigFileCatalog(dpath + '/fastpm_%0.4f/1' % aa) fpos = dyn['Position'].compute() idd = dyn['ID'] attrs = dyn.attrs play = pm.decompose(fpos) # wts = BigFileCatalog(ldpath + '/lagweights/') wts['1'] = wts['b1'] * 0 + 1. grid = wts['InitPosition'] iddg = wts['ID'] print('asserting') np.allclose(idd.compute(), iddg.compute()) print('read fields') disp = grid.compute() - fpos mask = abs(disp) > bs / 2. disp[mask] = (bs - abs(disp[mask])) * -np.sign(disp[mask]) print(rank, ' Max disp: ', disp.max())
dpath = '/global/cscratch1/sd/chmodi/m3127/cm_lowres/20stepT-B1/%d-%d-9100/' % ( bs, nc) #dpath = '/global/cscratch1/sd/chmodi/m3127/cm_lowres/5stepT-B1/%d-%d-9100-fixed/'%(bs, nc) aa = 1.0000 zz = 1 / aa - 1 Rsm = 0 pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc]) rank = pm.comm.rank #grid = pm.mesh_coordinates()*bs/nc hcat = BigFileCatalog(dpath + '/fastpm_%0.4f/LL-0.200/' % aa) hpos = hcat['Position'].compute() #hmass = print('Mass : ', rank, hcat['Mass'][-1].compute()) hlay = pm.decompose(hpos) hmesh = pm.paint(hpos, layout=hlay) hmesh /= hmesh.cmean() ph = FFTPower(hmesh, mode='1d').power k, ph = ph['k'], ph['power'].real sn = (hpos.shape[0] / bs**3)**-1 print(sn) if rank == 0: for Rsm in [0, 2]: for zadisp in [True, False]: header = '1, b1, b2, bg, bk' if zadisp: spectra = np.loadtxt(
pass pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc], dtype='f8') rank = pm.comm.rank hcat = BigFileCatalog(dpath + '/fastpm_%0.4f/LL-0.200/' % aa) print(rank, 'files read') #print('Mass : ', rank, hcat['Length'][-1].compute()*hcat.attrs['M0']*1e10) numd = [1e-2, 5e-3, 1e-3, 5e-4, 1e-4, 5e-5] num = [int(bs**3 * i) for i in numd] for i in range(len(num)): if rank == 0: print(numd[i]) cat = hcat.gslice(start=0, stop=num[i]) hlay = pm.decompose(cat['Position']) hmesh = pm.paint(cat['Position'], layout=hlay) hmesh /= hmesh.cmean() ph = FFTPower(hmesh, mode='1d').power k, ph = ph['k'], ph['power'] np.savetxt(ofolder + '/ph-%05d.txt' % (numd[i] * 1e5), np.vstack([k, ph]).T.real, header='k ph', fmt='%0.4e')
class TidalTensor(Algorithm): """ Compute and save the tidal force tensor """ plugin_name = "TidalTensor" def __init__(self, field, points, Nmesh, smoothing=None): from pmesh.pm import ParticleMesh self.field = field self.points = points self.Nmesh = Nmesh self.smoothing = smoothing self.pm = ParticleMesh(BoxSize=self.field.BoxSize, Nmesh=[self.Nmesh]*3, dtype='f4', comm=self.comm) @classmethod def fill_schema(cls): s = cls.schema s.description = "compute the tidal force tensor" s.add_argument("field", type=DataSource.from_config, help="DataSource; run `nbkit.py --list-datasources` for all options") s.add_argument("points", type=DataSource.from_config, help="A small set of points to calculate tidal force on; " "run `nbkit.py --list-datasources` for all options") s.add_argument("Nmesh", type=int, help='Size of FFT mesh for painting') s.add_argument("smoothing", type=float, help='Smoothing Length in distance units. ' 'It has to be greater than the mesh resolution. ' 'Otherwise the code will die. Default is the mesh resolution.') def Smoothing(self, pm, complex): k = pm.k k2 = 0 for ki in k: ki2 = ki ** 2 complex *= numpy.exp(-0.5 * ki2 * self.smoothing ** 2) def NormalizeDC(self, pm, complex): """ removes the DC amplitude. This effectively divides by the mean """ w = pm.w comm = pm.comm ind = [] value = 0.0 found = True for wi in w: if (wi != 0).all(): found = False break ind.append((wi == 0).nonzero()[0][0]) if found: ind = tuple(ind) value = numpy.abs(complex[ind]) value = comm.allreduce(value) complex[:] /= value def TidalTensor(self, u, v): # k_u k_v / k **2 def TidalTensor(pm, complex): k = pm.k for row in range(complex.shape[0]): k2 = k[0][row] ** 2 for ki in k[1:]: k2 = k2 + ki[0] ** 2 k2[k2 == 0] = numpy.inf complex[row] /= k2 complex *= k[u] complex *= k[v] return TidalTensor def run(self): """ Run the TidalTensor Algorithm """ from itertools import product if self.smoothing is None: self.smoothing = self.field.BoxSize[0] / self.Nmesh elif (self.field.BoxSize / self.Nmesh > self.smoothing).any(): raise ValueError("smoothing is too small") painter = Painter.create("DefaultPainter", weight="Mass", paintbrush="cic") real, stats = painter.paint(self.pm, self.field) complex = real.r2c() for t in [self.Smoothing, self.NormalizeDC]: t(complex.pm, complex) with self.points.open() as stream: [[Position ]] = stream.read(['Position'], full=True) layout = self.pm.decompose(Position) pos1 = layout.exchange(Position) value = numpy.empty((3, 3, len(Position))) for u, v in product(range(3), range(3)): if self.comm.rank == 0: self.logger.info("Working on tensor element (%d, %d)" % (u, v)) c2 = complex.copy() self.TidalTensor(u, v)(c2.pm, c2) c2.c2r(real) v1 = real.readout(pos1) v1 = layout.gather(v1) value[u, v] = v1 return value.transpose((2, 0, 1)) def save(self, output, data): self.write_hdf5(data, output) def write_hdf5(self, data, output): import h5py size = self.comm.allreduce(len(data)) offset = sum(self.comm.allgather(len(data))[:self.comm.rank]) if self.comm.rank == 0: with h5py.File(output, 'w') as ff: dataset = ff.create_dataset(name='TidalTensor', dtype=data.dtype, shape=(size, 3, 3)) dataset.attrs['Smoothing'] = self.smoothing dataset.attrs['Nmesh'] = self.Nmesh dataset.attrs['Original'] = self.field.string dataset.attrs['BoxSize'] = self.field.BoxSize for i in range(self.comm.size): self.comm.barrier() if i != self.comm.rank: continue with h5py.File(output, 'r+') as ff: dataset = ff['TidalTensor'] dataset[offset:len(data) + offset] = data
wsize = pm.comm.size comm = pm.comm newcomm = comm.Split(color=rank//splits, key=rank) 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) pmsmall = ParticleMesh(BoxSize = bs/ncube, Nmesh = [cube_size, cube_size, cube_size], dtype=np.float32, comm=newcomm) gridsmall = pmsmall.generate_uniform_particle_grid(shift=0) layoutsmall = pmsmall.decompose(gridsmall) for zz in [3.5, 4.0, 5.0, 6.0][::-1]: aa = 1/(1+zz) cats, meshes = setupuvmesh(zz, suff=suff, sim=sim, pm=pm, profile=profile, stellar=stellar, stellarfluc=stellarfluc, Rg=R) 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+\ '/dmesh_N%04d/1/'%nc,'').paint() meshes = [dm, h1meshfid, h1mesh, lmesh, uvmesh] names = ['dm', 'h1fid', 'h1new', 'lum', 'uv'] if rank == 0 : print('\nMeshes read\n')
satcat['blackhole'] = mbh(moster(satcat['Mass'].compute(), z=zz, scatter=0.3), alpha, beta, scatter=0.3) cencat['luminosity'] = lq(cencat['blackhole'], fon=switchon, eta=0.1, scatter=0.3) satcat['luminosity'] = lq(satcat['blackhole'], fon=switchon, eta=0.1, scatter=0.3) clayout = pm.decompose(cencat['Position']) slayout = pm.decompose(satcat['Position']) cmesh = pm.paint(cencat['Position'], mass=cencat['luminosity'], layout=clayout) smesh = pm.paint(satcat['Position'], mass=satcat['luminosity'], layout=slayout) lmesh = cmesh + smesh if lumspectra: calc_bias(aa, lmesh / lmesh.cmean(), suff, fname='Lum') mfpath = mfpathz(zz) if rank == 0: print('At redshift {:.2f}, mean free path is {:.2f}'.format( zz, mfpath)) meshc = lmesh.r2c()
def measurepk(nc=nc, dpath=scratchyf): '''plot the power spectrum of halos on 'nc' grid''' pm = ParticleMesh(BoxSize=bs, Nmesh=[nc, nc, nc]) for i, aa in enumerate(aafiles): zz = zzfiles[i] if rank == 0: print('redshift = ', zz) 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+\ '/dmesh_N%04d/1/'%nc,'').paint() #dm = BigFileMesh(project + sim + '/fastpm_%0.4f/'%aa + '/dmesh_N%04d'%nc, '1').paint() halos = BigFileCatalog(scratchyf + sim + '/fastpm_%0.4f/' % aa, dataset='LL-0.200') mp = halos.attrs['MassTable'][1] * 1e10 if rank == 0: print('Mass of particle is = %0.2e' % mp) hmass = halos["Length"].compute() * mp hpos = halos['Position'].compute() layout = pm.decompose(hpos) if rank == 0: print("paint") hpmesh = pm.paint(hpos, layout=layout) hmesh = pm.paint(hpos, mass=hmass, layout=layout) print(rank, dm.cmean(), hmesh.cmean(), hpmesh.cmean()) pkm = FFTPower(dm / dm.cmean(), mode='1d').power pkh = FFTPower(hmesh / hmesh.cmean(), mode='1d').power pkhp = FFTPower(hpmesh / hpmesh.cmean(), mode='1d').power pkhm = FFTPower(hmesh / hmesh.cmean(), second=dm / dm.cmean(), mode='1d').power pkhpm = FFTPower(hpmesh / hpmesh.cmean(), second=dm / dm.cmean(), mode='1d').power def savebinned(path, binstat, header): if halos.comm.rank == 0: k, p, modes = binstat['k'].real, binstat[ 'power'].real, binstat['modes'].real np.savetxt(path, np.stack((k, p, modes), axis=1), header=header) ofolder = "../data/outputs/halos/{}/".format(sim) if rank == 0: print(ofolder) try: os.makedirs(ofolder) except: pass savebinned(ofolder + 'pkhp_%0.4f.txt' % aa, pkhp, header='k, P(k), Nmodes') savebinned(ofolder + 'pkhm_%0.4f.txt' % aa, pkh, header='k, P(k), Nmodes') savebinned(ofolder + 'pkd_%0.4f.txt' % aa, pkm, header='k, P(k), Nmodes') savebinned(ofolder + 'pkhpxd_%0.4f.txt' % aa, pkhpm, header='k, P(k), Nmodes') savebinned(ofolder + 'pkhmxd_%0.4f.txt' % aa, pkhm, header='k, P(k), Nmodes')
indices[bindex][0], indices[bindex][1], indices[bindex][ 2] = bi, bj, bk bi *= cube_length bj *= cube_length bk *= cube_length print('For rank & index : ', rank, bindex, '-- x, y, z : ', bi, bj, bk) poslarge = gridsmall + np.array([bi, bj, bk]) save = True except IndexError: poslarge = np.empty((1, 3)) save = False #print(rank, 'Position created') layout = pm.decompose(poslarge) #if rank == 0 : print(rank, 'Layout decomposed') for i in range(len(meshes)): vals = meshes[i].readout(poslarge, layout=layout, resampler='nearest').astype( np.float32) name = names[i] if save: savemesh = pmsmall.paint(gridsmall, mass=vals, resampler='nearest') totals[bindex, i] = savemesh.csum() #print('In rank {:3d}, sum for mesh {:8} is equal to {:.3e}'.format(rank, name, savemesh.csum()))
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, mode='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)
class Subsample(Algorithm): """ Algorithm to create a subsample from a DataSource, and evaluate the density (1 + delta), smoothed at the given scale """ plugin_name = "Subsample" def __init__(self, datasource, Nmesh, seed=12345, ratio=0.01, smoothing=None, format='hdf5'): from pmesh.pm import ParticleMesh self.datasource = datasource self.Nmesh = Nmesh self.seed = seed self.ratio = ratio self.smoothing = smoothing self.format = format self.pm = ParticleMesh(BoxSize=self.datasource.BoxSize, Nmesh=[self.Nmesh] * 3, dtype='f4', comm=self.comm) @classmethod def fill_schema(cls): s = cls.schema s.description = "create a subsample from a DataSource, and evaluate density \n" s.description += "(1 + delta) smoothed at the given scale" s.add_argument( "datasource", type=DataSource.from_config, help= "the DataSource to read; run `nbkit.py --list-datasources` for all options" ) s.add_argument("Nmesh", type=int, help='the size of FFT mesh for painting') s.add_argument("seed", type=int, help='the random seed') s.add_argument("ratio", type=float, help='the fraction of particles to keep') s.add_argument( "smoothing", type=float, help='the smoothing length in distance units. ' 'It has to be greater than the mesh resolution. ' 'Otherwise the code will die. Default is the mesh resolution.') # this is for output.. s.add_argument("format", choices=['hdf5', 'mwhite'], help='the format of the output') def run(self): """ Run the Subsample algorithm """ import mpsort from astropy.utils.misc import NumpyRNGContext if self.smoothing is None: self.smoothing = self.datasource.BoxSize[0] / self.Nmesh[0] elif (self.datasource.BoxSize / self.Nmesh > self.smoothing).any(): raise ValueError("smoothing is too small") def Smoothing(pm, complex): k = pm.k k2 = 0 for ki in k: ki2 = ki**2 complex[:] *= numpy.exp(-0.5 * ki2 * self.smoothing**2) def NormalizeDC(pm, complex): """ removes the DC amplitude. This effectively divides by the mean """ w = pm.w comm = pm.comm ind = [] value = 0.0 found = True for wi in w: if (wi != 0).all(): found = False break ind.append((wi == 0).nonzero()[0][0]) if found: ind = tuple(ind) value = numpy.abs(complex[ind]) value = comm.allreduce(value) complex[:] /= value # open the datasource and keep the cache with self.datasource.keep_cache(): painter = Painter.create("DefaultPainter", paintbrush='cic') real, stats = painter.paint(self.pm, self.datasource) complex = real.r2c() for t in [Smoothing, NormalizeDC]: t(self.pm, complex) complex.c2r(real) columns = ['Position', 'ID', 'Velocity'] local_seed = utils.local_random_seed(self.seed, self.comm) dtype = numpy.dtype([ ('Position', ('f4', 3)), ('Velocity', ('f4', 3)), ('ID', 'u8'), ('Density', 'f4'), ]) subsample = [numpy.empty(0, dtype=dtype)] with self.datasource.open() as stream: for Position, ID, Velocity in stream.read(columns): with NumpyRNGContext(local_seed): u = numpy.random.uniform(size=len(ID)) keep = u < self.ratio Nkeep = keep.sum() if Nkeep == 0: continue data = numpy.empty(Nkeep, dtype=dtype) data['Position'][:] = Position[keep] data['Velocity'][:] = Velocity[keep] data['ID'][:] = ID[keep] layout = self.pm.decompose(data['Position']) pos1 = layout.exchange(data['Position']) density1 = real.readout(pos1) density = layout.gather(density1) # normalize the position after reading out density! data['Position'][:] /= self.datasource.BoxSize data['Velocity'][:] /= self.datasource.BoxSize data['Density'][:] = density subsample.append(data) subsample = numpy.concatenate(subsample) mpsort.sort(subsample, orderby='ID', comm=self.comm) return subsample def save(self, output, data): if self.format == 'mwhite': self.write_mwhite_subsample(data, output) else: self.write_hdf5(data, output) def write_hdf5(self, subsample, output): import h5py size = self.comm.allreduce(len(subsample)) offset = sum(self.comm.allgather(len(subsample))[:self.comm.rank]) if self.comm.rank == 0: with h5py.File(output, 'w') as ff: dataset = ff.create_dataset(name='Subsample', dtype=subsample.dtype, shape=(size, )) dataset.attrs['Ratio'] = self.ratio dataset.attrs['CommSize'] = self.comm.size dataset.attrs['Seed'] = self.seed dataset.attrs['Smoothing'] = self.smoothing dataset.attrs['Nmesh'] = self.Nmesh dataset.attrs['Original'] = self.datasource.string dataset.attrs['BoxSize'] = self.datasource.BoxSize for i in range(self.comm.size): self.comm.barrier() if i != self.comm.rank: continue with h5py.File(output, 'r+') as ff: dataset = ff['Subsample'] dataset[offset:len(subsample) + offset] = subsample def write_mwhite_subsample(self, subsample, output): size = self.comm.allreduce(len(subsample)) offset = sum(self.comm.allgather(len(subsample))[:self.comm.rank]) if self.comm.rank == 0: with open(output, 'wb') as ff: dtype = numpy.dtype([('eflag', 'int32'), ('hsize', 'int32'), ('npart', 'int32'), ('nsph', 'int32'), ('nstar', 'int32'), ('aa', 'float'), ('gravsmooth', 'float')]) header = numpy.zeros((), dtype=dtype) header['eflag'] = 1 header['hsize'] = 20 header['npart'] = size header.tofile(ff) self.comm.barrier() with open(output, 'r+b') as ff: ff.seek(28 + offset * 12) numpy.float32(subsample['Position']).tofile(ff) ff.seek(28 + offset * 12 + size * 12) numpy.float32(subsample['Velocity']).tofile(ff) ff.seek(28 + offset * 4 + size * 24) numpy.float32(subsample['Density']).tofile(ff) ff.seek(28 + offset * 8 + size * 28) numpy.uint64(subsample['ID']).tofile(ff)
def main(): ns = ap.parse_args() comm = MPI.COMM_WORLD ff = bigfile.BigFileMPI(comm, ns.fastpm) with ff['.'] as bb: BoxSize = bb.attrs['BoxSize'][0] Redshift = 1 / bb.attrs['ScalingFactor'][0] - 1 Nmesh = int(BoxSize / ns.resolution * 2) # round it to 8. Nmesh -= Nmesh % 8 if comm.rank == 0: logger.info("source = %s", ns.fastpm) logger.info("output = %s", ns.output) logger.info("BoxSize = %g", BoxSize) logger.info("Redshift = %g", Redshift) logger.info("Nmesh = %g", Nmesh) pm = ParticleMesh([Nmesh, Nmesh, Nmesh], BoxSize, comm=comm) real = RealField(pm) real[...] = 0 with ff['Position'] as ds: logger.info(ds.size) for i in range(0, ds.size, ns.chunksize): sl = slice(i, i + ns.chunksize) pos = ds[sl] layout = pm.decompose(pos) lpos = layout.exchange(pos) real.paint(lpos, hold=True) mean = real.cmean() if comm.rank == 0: logger.info("mean particle per cell = %s", mean) real[...] /= mean real[...] -= 1 complex = real.r2c() for k, i, slab in zip(complex.slabs.x, complex.slabs.i, complex.slabs): k2 = sum(kd**2 for kd in k) # tophat f = tophat(ns.filtersize, k2**0.5) slab[...] *= f # zreion slab[...] *= Bk(k2**0.5) slab[...] *= (1 + Redshift) real = complex.c2r() real[...] += Redshift mean = real.cmean() if comm.rank == 0: logger.info("zreion.mean = %s", mean) buffer = numpy.empty(real.size, real.dtype) real.sort(out=buffer) if comm.rank == 0: logger.info("sorted for output") with bigfile.BigFileMPI(comm, ns.output, create=True) as ff: with ff.create_from_array(ns.dataset, buffer) as bb: bb.attrs['BoxSize'] = BoxSize bb.attrs['Redshift'] = Redshift bb.attrs['TopHatFilterSize'] = ns.filtersize bb.attrs['Nmesh'] = Nmesh # # hack: compatible with current MPGadget. This is not really needed # we'll remove the bins later, since BoxSize and Nmesh are known. with ff.create("XYZ_bins", dtype='f8', size=Nmesh) as bb: if comm.rank == 0: bins = numpy.linspace(0, BoxSize * 1000., Nmesh, dtype='f8') bb.write(0, bins) if comm.rank == 0: logger.info("done. written at %s", ns.output)