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