def work(j): input = Snapshot(inputname(j), 'cmugadget', template=first) id = input[None, 'id'] t1 = 0 t2 = 0 t3 = 0 for i in TheBigGuys: masks = [] output = Snapshot(outputname(i, j), 'cmugadget', create=True, template=first) for ptype in range(6): with Timer() as timer: id = input[ptype, 'id'] mask = cat.mask(id, groupid=i) masks.append(mask) output.C['N'][ptype] = mask.sum() t1 += timer.interval for blk in input.schema: if not (None, blk) in input: continue with Timer() as timer: junk = input[None, blk] for ptype in range(6): if not (ptype, blk) in input: continue if output.C['N'][ptype] == 0: continue output[ptype, blk][:] = input[ptype, blk][masks[ptype]] t2 += timer.interval with Timer() as timer: output.save() t3 += timer.interval return j, t1, t2, t3
def main(): parser = argparse.ArgumentParser() parser.add_argument("-f", dest='format', required=True) parser.add_argument("filename") parser.add_argument("-N", dest='N', default=64, required=False, type=int) parser.add_argument("-o", dest='output', default=None, required=True) args = parser.parse_args() if '%d' in args.filename: snap0 = Snapshot(args.filename % 0, args.format) Nfile = snap0.C['Nfiles'] else: snap0 = Snapshot(args.filename, args.format) Nfile = 1 boxsize = snap0.C['boxsize'] m = MeshIndex(N=args.N, Nd=3, boxsize=boxsize) with sharedmem.Pool(use_threads=False) as pool: def work(fid): if '%d' in args.filename: snap = Snapshot(args.filename % fid, args.format) else: snap = Snapshot(args.filename, args.format) print 'start', snap.file pos = snap['pos'] print 'read', snap.file return m.set(fid, pos) l = pool.map(work, range(Nfile)) m.compile(dict(l)) m.tofile(args.output)
def crop_snapshot(center, size, map, format, snapname, output): cut = Cut(center = center, size=size) mesh = Meshmap(map) gas = Field(components={'mass':'f4', 'rho':'f4', 'ie':'f4', 'xHI':'f4', 'ye':'f4', 'id':'u8', 'sfr':'f4', 'met':'f4', 'sml':'f4'}) bh = Field(components={'mass':'f4', 'bhmdot':'f4', 'bhmass':'f4', 'id':'u8'}) star = Field(components={'mass':'f4', 'sft':'f4', 'met':'f4', 'id':'u8'}) fids = mesh.cut2fid(cut) print fids snapshots = [Snapshot(snapname % fid, format) for fid in fids] gas.take_snapshots(snapshots, 0, cut=cut) star.take_snapshots(snapshots, 4, cut=cut) bh.take_snapshots(snapshots, 5, cut=cut) left = center - size * 0.5 gas['locations'][:, :] -= left[newaxis, :] star['locations'][:, :] -= left[newaxis, :] bh['locations'][:, :] -= left[newaxis, :] print bh['id'] out = Snapshot(output, format, create=True) for f in snapshots[0].header.dtype.names: out.header[f] = snapshots[0].header[f] out.header['unused'][0:3] = center out.header['boxsize'] = size[0] gas.dump_snapshots([out], 0) star.dump_snapshots([out], 4) bh.dump_snapshots([out], 5)
def work(fid): if '%d' in args.filename: snap = Snapshot(args.filename % fid, args.format) else: snap = Snapshot(args.filename, args.format) print 'start', snap.file pos = snap['pos'] print 'read', snap.file return m.set(fid, pos)
def main(): parser = argparse.ArgumentParser() parser.add_argument("-f", dest='format', default='genic', required=False) parser.add_argument("filename") # n is bits per chanel parser.add_argument("-n", dest='bitsperaxis', default=6, required=False, type=int) parser.add_argument("-o", dest='output', default=None, required=True) args = parser.parse_args() bits = args.bitsperaxis * 3 if '%d' in args.filename: snap0 = Snapshot(args.filename % 0, args.format) Nfile = snap0.C['Nfiles'] else: snap0 = Snapshot(args.filename, args.format) Nfile = 1 boxsize = snap0.C['boxsize'] C = [] F = [] with sharedmem.Pool(use_threads=True) as pool: def work(fid): if '%d' in args.filename: snap = Snapshot(args.filename % fid, args.format) else: snap = Snapshot(args.filename, args.format) N = snap.C['N'].sum() x,y,z = snap[None, 'pos'].T del snap scale = fillingcurve.scale(0, boxsize) zkey = fillingcurve.encode(x, y, z, scale=scale) del x, y, z dig = numpy.uint32(zkey >> (fillingcurve.bits * 3 - bits)) bincount = numpy.bincount(dig) return (fid, bincount.nonzero()[0]) def reduce(res): fid, active = res F.append(fid) C.append(list(active)) print fid, len(active) pool.map(work, range(Nfile), callback=reduce) F = numpy.array(F, dtype='u4') l = [len(a) for a in C] C = numpy.concatenate(C) fid = numpy.repeat(F, l) arg = C.argsort() fid = fid[arg] count = numpy.bincount(C, minlength=1<<bits) index = packarray(fid, count) map = MeshIndex(0, boxsize, args.bitsperaxis, index) map.tofile(args.output)
def __init__(self, tabfilename, format, count=10, **kwargs): """ tabfilename is like groups_019/group_tab_019.%d. """ g = Snapshot(tabfilename % 0, format + '.GroupTab', **kwargs) if count < 0 or count > g.C['Ntot'][0]: count = g.C['Ntot'][0] i = 0 # decide number of files to open nread = 0 tabs = [] while nread < count: g = Snapshot(tabfilename % i, format + '.GroupTab', **kwargs) nread += g.C['N'][0] i = i + 1 tabs.append(g) #print 'will read', len(tabs), 'files' Field.__init__(self, numpoints=count, components={'offset':'i8', 'length':'i8', 'massbytype':('f8', 6), 'mass':'f8', 'pos':('f8', 3), 'vel':('f8', 3)}) if len(tabs) > 0: self.take_snapshots(tabs, ptype=0) del tabs # fix the offset which may overflow for large halos self['offset'][1:] = self['length'].cumsum()[:-1] nread = 0 nshallread = self['length'].sum() i = 0 idslen = numpy.zeros(g.C['Nfiles'], dtype='i8') while nread < nshallread: idslen[i] = numpy.fromfile(tabfilename.replace('_tab_', '_ids_') % i, dtype='i4', count=3)[2] nread += idslen[i] i = i + 1 idsoffset = numpy.concatenate(([0], idslen.cumsum())) ids = sharedmem.empty(idslen.sum(), dtype=g.C['idtype']) #print 'reading', i, 'id files' with sharedmem.Pool() as pool: def work(i): more = numpy.memmap(tabfilename.replace('_tab_', '_ids_') % i, dtype=g.C['idtype'], mode='r', offset=28) ids[idsoffset[i]:idsoffset[i] + idslen[i]] = more pool.map(work, range(i)) self.ids = packarray(ids, self['length']) for i in range(self.numpoints): self.ids[i].sort()
def work(i): snap = Snapshot(idsfile % i, 'cmugadget.GroupIDs', idtype='u8', floattype='f4') id = snap[0, 'id'] result = numpy.zeros(len(id), dtype='f4') lookedid = numpy.zeros(len(id), dtype='u8') ind = sfgasid.searchsorted(id) ind.clip(0, len(sfgasid) - 1, out=ind) found = sfgasid[ind] == id result[found] = sfgassfr[ind[found]] with pool.ordered: print 'start', i, matchedidfile.tell() / 8., len(id), id[0] matchedfile.seek(0, 2) result.tofile(matchedfile) matchedfile.flush() matchedfile.seek(0, 2) matchedidfile.seek(0, 2) id.tofile(matchedidfile) matchedidfile.flush() matchedidfile.seek(0, 2) print i, 'found', found.sum(), matchedidfile.tell() / 8. return found.sum()
def work(fid): if '%d' in args.filename: snap = Snapshot(args.filename % fid, args.format) else: snap = Snapshot(args.filename, args.format) N = snap.C['N'].sum() x,y,z = snap[None, 'pos'].T del snap scale = fillingcurve.scale(0, boxsize) zkey = fillingcurve.encode(x, y, z, scale=scale) del x, y, z dig = numpy.uint32(zkey >> (fillingcurve.bits * 3 - bits)) bincount = numpy.bincount(dig) return (fid, bincount.nonzero()[0])
def __init__(self, tabfile): self.tabfile = tabfile first = Snapshot(tabfile % 0, 'cmugadget.SubHaloTab', idtype='u8', floattype='f4') self.Nhalo = int(first.C['NtotSubgroups']) Depot.__init__(self, subdtype, first.C['Nfiles'])
def __init__(self, tabfile): self.tabfile = tabfile first = Snapshot(tabfile % 0, 'cmugadget.GroupTab', idtype='u8', floattype='f4') Depot.__init__(self, groupdtype, first.C['Nfiles']) self.Ngroups = first.C['Ntot'][0]
def __init__(self, tabfile): self.tabfile = tabfile self.posvelfile = self.tabfile.replace('_tab_', '_posvel_') self.idsfile = self.tabfile.replace('_tab_', '_ids_') first = Snapshot(tabfile % 0, 'cmugadget.SubHaloTab', idtype='u8', floattype='f4') Depot.__init__(self, pdtype, first.C['Nfiles'])
def work(i): snap = Snapshot(argv[1] % i, 'cmugadget', template=first) print 'in', i sfr = snap[0, 'sfr'] id = snap[0, 'id'] mask = sfr > 0 print 'done', i, mask.sum(), sfr.max(), sfr.min() return mask.sum()
def work(i): snap = Snapshot(argv[1] % i, 'cmugadget', template=first) print 'i', i sfr = snap[0, 'sfr'] id = snap[0, 'id'] print 'i', i, 'read' mask = sfr > 0 idall[offset[i]:offset[i]+size[i]] = id[mask] sfrall[offset[i]:offset[i]+size[i]] = sfr[mask] print 'i', i, 'assign' print 'i', i, mask.sum()
def main(): parser = argparse.ArgumentParser() parser.add_argument("-f", dest='format', default='genic', required=False) parser.add_argument("filename") parser.add_argument("-p", nargs=2, default=[], dest='put', action='append', required=False) parser.add_argument("-g", dest='get', default=[], action='append', required=False) args = parser.parse_args() snap = Snapshot(args.filename, args.format) if len(args.put) == 0 and len(args.get) == 0: for name in snap.C: print name, '=', snap.C[name] return for field, value in args.put: snap.C[field] = value print field, 'set to', snap.C[field] snap.save_header() for field in args.get: print snap.C[field]
def prefetcher(fid): print 'pref', fid try: snap = Snapshot(snapname % fid, 'cmugadget') snap[0, 'pos'] snap[0, 'ie'] snap[0, 'ye'] snap[0, 'sml'] snap[0, 'mass'] except: pass return
def work(i): snap = Snapshot(argv[1] % i, 'cmugadget', template=first) numpy.random.seed(seedtable[i]) # print 'in', i with pool.critical: pos = snap[1, 'pos'] r = numpy.random.uniform(size=len(pos)) mask = r < 100.**3 / first.C['Ntot'][1] pos = pos[mask] # print 'done', i, mask.sum(), sfr.max(), sfr.min() with pool.critical: numpy.savetxt(stdout, pos, fmt='%0.7g') stdout.flush()
def fetch(self, id): f = Snapshot(self.tabfile % id, 'cmugadget.GroupTab', idtype='u8', floattype='f4') rt = numpy.empty(f.C['N'][0], dtype=groupdtype) rt['pos'] = f[0, 'pos'] rt['vel'] = f[0, 'vel'] rt['mass'] = f[0, 'mass'] rt['len'] = f[0, 'length'] rt['massbytype'] = f[0, 'massbytype'] rt['lenbytype'] = f[0, 'lenbytype'] rt['nhalo'] = 0 return rt
def read(self, fids=None, use_gas=True, use_bh=True, use_star=True, use_halo=False, use_disk=False, use_bulge=False, numthreads=None): if fids is not None: snapnames = [self.snapname % i for i in fids] elif '%d' in self.snapname: snapnames = [self.snapname % i for i in range(self.C['Nfiles'])] else: snapnames = [self.snapname] snapshots = [Snapshot(snapname, self.format) for snapname in snapnames] if use_gas: self.gas.take_snapshots(snapshots, ptype=self.ptype['gas'], nthreads=numthreads, cut=self.cut) if use_halo: self.halo.take_snapshots(snapshots, ptype=self.ptype['halo'], nthreads=numthreads, cut=self.cut) if use_disk: self.disk.take_snapshots(snapshots, ptype=self.ptype['disk'], nthreads=numthreads, cut=self.cut) if use_bulge: self.bulge.take_snapshots(snapshots, ptype=self.ptype['bulge'], nthreads=numthreads, cut=self.cut) if use_bh: self.bh.take_snapshots(snapshots, ptype=self.ptype['bh'], nthreads=numthreads, cut=self.cut) if use_star: self.star.take_snapshots(snapshots, ptype=self.ptype['star'], nthreads=numthreads, cut=self.cut) self.invalidate()
def fetch(self, id): f = Snapshot(self.tabfile % id, 'cmugadget.SubHaloTab', idtype='u8', floattype='f4') rt = numpy.empty(f.C['N'][1], dtype=subdtype) rt['pos'] = f[1, 'halopos'] rt['vel'] = f[1, 'halovel'] rt['len'] = f[1, 'halolen'] rt['mass'] = f[1, 'halomass'] rt['vdisp'] = f[1, 'haloveldisp'] rt['vcirc'] = f[1, 'halovmax'] rt['rcirc'] = f[1, 'halovmaxrad'] # i don't understand this; we do not want halo parent # it's always zero rt['parent'] = f[1, 'haloparent'] rt['lenbytype'] = 0 rt['massbytype'] = 0 return rt
def crop_snapshot(center, size, map, format, snapname, output): cut = Cut(center=center, size=size) mesh = Meshmap(map) gas = Field( components={ 'mass': 'f4', 'rho': 'f4', 'ie': 'f4', 'xHI': 'f4', 'ye': 'f4', 'id': 'u8', 'sfr': 'f4', 'met': 'f4', 'sml': 'f4' }) bh = Field(components={ 'mass': 'f4', 'bhmdot': 'f4', 'bhmass': 'f4', 'id': 'u8' }) star = Field(components={ 'mass': 'f4', 'sft': 'f4', 'met': 'f4', 'id': 'u8' }) fids = mesh.cut2fid(cut) print fids snapshots = [Snapshot(snapname % fid, format) for fid in fids] gas.take_snapshots(snapshots, 0, cut=cut) star.take_snapshots(snapshots, 4, cut=cut) bh.take_snapshots(snapshots, 5, cut=cut) left = center - size * 0.5 gas['locations'][:, :] -= left[newaxis, :] star['locations'][:, :] -= left[newaxis, :] bh['locations'][:, :] -= left[newaxis, :] print bh['id'] out = Snapshot(output, format, create=True) for f in snapshots[0].header.dtype.names: out.header[f] = snapshots[0].header[f] out.header['unused'][0:3] = center out.header['boxsize'] = size[0] gas.dump_snapshots([out], 0) star.dump_snapshots([out], 4) bh.dump_snapshots([out], 5)
def __init__(self, snapid, ROOT='../'): SnapDirReadOnly.__init__(self, snapid, ROOT) snapid = self.snapid self.snapname = ROOT + '/snapdir/snapdir_%03d/snapshot_%03d.%%d' %(snapid, snapid) first = Snapshot(self.snapname % 0, 'cmugadget') self.first = first self.C = first.C C = first.C self.cosmology = Cosmology(h=C['h'], M=C['OmegaM'], L=C['OmegaL']) self.U = self.cosmology.U self.schema = first.schema self.collectivedir = ROOT + \ '/isolatedir/isolate_%03d/%%03d' % snapid self.tabfile = ROOT + \ '/groups_and_subgroups/groups_%03d/subhalo_tab_%03d.%%d' %(snapid, snapid) if not os.path.isfile(self.tabfile % 0): self.tabfile = ROOT + \ '/groups_and_subgroups/no_collective_halos/groups_%03d/subhalo_tab_%03d.%%d' \ %(snapid, snapid) self.grouptabfile = self.tabfile.replace('subhalo', 'group')
from gaepsi.snapshot import Snapshot from sys import argv print len(argv[1:]) sfr = 0 for fn in argv[1:]: snap = Snapshot(fn, 'cmugadget') assert snap.C['Nfiles'] == len(argv[1:]) sfr += snap[0, 'sfr'].sum(dtype='f8') print sfr
ses_import.end() # when opt.range == None, there is only one file in the snapshot if opt.range != None: snapfiles = [opt.snapname % i for i in range(opt.range[0], opt.range[1])] else: if opt.rangefile != None: from numpy import loadtxt snapids = loadtxt(opt.rangefile) snapfiles = [opt.snapname % i for i in snapids] else: snapfiles = [opt.snapname] #decide the boxsize if comm.rank == 0: snap = Snapshot(snapfiles[0], opt.format) boxsize = snap.C['boxsize'] del snap else: boxsize = None opt.boxsize = comm.bcast(boxsize) del boxsize # create the stripes and layers stripe = Stripe(comm=comm, imagesize=opt.geometry, xcut=opt.xcut, ycut=opt.ycut, zcut=opt.zcut) if opt.gas != None:
with Timer() as timer: junk = input[None, blk] for ptype in range(6): if not (ptype, blk) in input: continue if output.C['N'][ptype] == 0: continue output[ptype, blk][:] = input[ptype, blk][masks[ptype]] t2 += timer.interval with Timer() as timer: output.save() t3 += timer.interval return j, t1, t2, t3 def reduce(j, t1, t2, t3): print j, Nfiles, 'times', t1, t2, t3 with sharedmem.Pool() as pool: pool.map(work, range(Nfiles), reduce=reduce) for i in TheBigGuys: Ntot = numpy.zeros(6, 'i8') for j in range(Nfiles): output = Snapshot(outputname(i, j), 'cmugadget', template=first) Ntot += output.C['N'] print i, Ntot for j in range(Nfiles): output = Snapshot(outputname(i, j), 'cmugadget', template=first) output.C['Ntot'][:] = Ntot output.save('header')
from gaepsi.snapshot import Snapshot from gaepsi.tools.analyze import HaloCatalog import numpy import time import sharedmem import os.path from sys import argv first = Snapshot(argv[1] % 0, 'cmugadget', idtype='u8') class Timer: def __enter__(self): self.start = time.clock() return self def __exit__(self, *args): self.end = time.clock() self.interval = self.end - self.start Ntot = first.C['Ntot'] Nout = first.C['Nfiles'] for i in range(Nout): output = Snapshot(argv[2] % i, 'cmugadget', template=first, create=True) output.C['N'] = Ntot * (i + 1) // Nout - Ntot * i // Nout output.save() for ptype in range(6): buffer = {} j = 0 for blk in first.schema:
def use(self, snapname, format, components={}, bhcomponents={ 'bhmass': 'f4', 'bhmdot': 'f4', 'id': 'u8' }, starcomponents={ 'sft': 'f4', 'mass': 'f4' }, gas=0, halo=1, disk=2, bulge=3, star=4, bh=5, cut=None, gascomponents=None, periodic=True): if gascomponents is not None: self.components = gascomponents else: self.components = components self.components['mass'] = 'f4' self.components['sml'] = 'f4' self.snapname = snapname self.format = format self.F['gas'] = Field(components=self.components) self.F['bh'] = Field(components=bhcomponents) self.F['star'] = Field(components=starcomponents) self.F['halo'] = Field(components=self.components) self.F['disk'] = Field(components=self.components) self.F['bulge'] = Field(components=self.components) self.ptype = { "gas": gas, "halo": halo, "disk": disk, "bulge": bulge, "star": star, "bh": bh, } try: snapname = self.snapname % 0 except TypeError: snapname = self.snapname snap = Snapshot(snapname, self.format) self.gas.init_from_snapshot(snap) self.bh.init_from_snapshot(snap) self.star.init_from_snapshot(snap) self.halo.init_from_snapshot(snap) self.disk.init_from_snapshot(snap) self.bulge.init_from_snapshot(snap) self.C = snap.C if cut is not None: self.cut.take(cut) else: try: boxsize = snap.C['boxsize'] self.cut.take( Cut(xcut=[0, boxsize], ycut=[0, boxsize], zcut=[0, boxsize])) except: pass self.boxsize = ones(3) * snap.C['boxsize'] self.redshift = snap.C['redshift'] self.invalidate() self.periodic = periodic
#! python from gaepsi.snapshot import Snapshot from sys import argv template = Snapshot(argv[2], argv[1]) for filename in argv[2:]: try: snap = Snapshot(filename, argv[1], template=template) snap.check() print snap.file, 'ok' except IOError as e: print filename, e
import sharedmem import os.path from sys import argv first = Snapshot(argv[1] % 0, 'cmugadget', idtype='u8') class Timer: def __enter__(self): self.start = time.clock() return self def __exit__(self, *args): self.end = time.clock() self.interval = self.end - self.start Ntot = first.C['Ntot'] Nout = first.C['Nfiles'] for i in range(Nout): output = Snapshot(argv[2] % i, 'cmugadget', template=first, create=True) output.C['N'] = Ntot * (i + 1) // Nout - Ntot * i // Nout output.save() for ptype in range(6): buffer = {} j = 0 for blk in first.schema: buffer[blk] = numpy.empty(0, first.schema[blk].dtype) for i in range(Nout): output = Snapshot(argv[2] % i, 'cmugadget', template=first, readonly=False) while output.C['N'][ptype] > len(buffer['pos']): print 'need to read ', j, output.C['N'][ptype], len(buffer['pos']) newinput = Snapshot(argv[1] % j, 'cmugadget', template=first)
def opensnap(self, fid): return Snapshot(self.snapname %fid, 'cmugadget', template=self.first)
from gaepsi.render import * from matplotlib import cm from sys import argv from gaepsi.snapshot import Snapshot from gaepsi.gaplot import create_cosmology import os.path from gaepsi.compiledbase.geometry import Cubenoid from sys import argv execfile(argv[1]) # center = [ 62995.8671875, 40180.12109375, 52267.77734375] # pos of most massive halo from halo catalog # unfold from here. #snapname = '/physics/yfeng1/e5/run/000/output/snapdir_1366/snapshot_1366.%d' snap = Snapshot(snapname % 0, 'cmugadget') cosmology = create_cosmology(snap.C) Nfile = snap.C['Nfiles'] boxsize = snap.C['boxsize'] cub = Cubenoid(M, [0, 0, 0], [boxsize, boxsize, boxsize], center=0, neworigin=0) ratio = cub.newboxsize[0] / cub.newboxsize[1] width = int(width * dpi) length = int(ratio * width) print 'real length is', 1.0 * length / dpi, 'inch' print 'kpc per inch is', cub.newboxsize[0] / length * dpi print 'kpc per dot is', cub.newboxsize[0] / length print 'pixel size is wxl', width, length print 'total pixels = ', width * length / 1024 / 1024., 'M'
from gaepsi.snapshot import Snapshot from gaepsi.tools.analyze import HaloCatalog import numpy import time import sharedmem import os.path from sys import argv firstcat = Snapshot(argv[1] % 0, 'cmugadget.GroupTab', idtype='u8', floattype='f4') Ncat = (firstcat[0, 'length'] > 8500000).nonzero()[0].max() class Timer: def __enter__(self): self.start = time.clock() return self def __exit__(self, *args): self.end = time.clock() self.interval = self.end - self.start TheBigGuys = range(Ncat) def outputname(haloid, j): outputname = '%03d/' % haloid + os.path.basename(argv[2]) return outputname % j
def splitMain(args): snap = args.snap for type in range(6): try: os.makedirs(snap.subhalodir + '/%d' % type) except OSError: pass file(snap.subhalodir + '/header.txt', mode='w').write(str(snap.C)) print snap.tabfile depot_sub = SubHaloDepot(snap.tabfile) depot_group = GroupDepot(snap.grouptabfile) if not args.check: wrong_file_or_die(snap.groupfile, depot_group.Ngroups * groupdtype.itemsize) depot_par = SubHaloParDepot(snap.tabfile) F = {} for type in range(6): F[type] = {} for field in pdtype.fields: if type != 4 and \ (field == 'SEDindex' or field == 'recfrac'): continue F[type][field] = snap.open(type, field, mode='w') halodump = file(snap.subhalofile, 'w') groupdump = file(snap.groupfile, 'w') realgroupid = 0 for i in range(len(depot_sub)): tab = Snapshot(snap.tabfile % i, 'cmugadget.SubHaloTab', idtype='u8', floattype='f4') group = depot_group.consume(tab.C['N'][0]) for groupid in range(tab.C['N'][0]): if snap.hascollective(realgroupid, group[groupid]): fakenhalo = tab[0, 'nhalo'][groupid] depot_subcol = SubHaloDepot( snap.collectivetabfile(realgroupid)) realhalo = depot_subcol.consume() nhalo = len(realhalo) # skip the pars from ids file (it's of wrong order) if not args.check: depot_parcol = SubHaloParDepot( snap.collectivetabfile(realgroupid)) pars = depot_parcol.consume() junk = depot_par.consume(group['len'][groupid]) assert len(pars) == group['len'][groupid] else: nhalo = tab[0, 'nhalo'][groupid] if not args.check: pars = depot_par.consume(group['len'][groupid]) realhalo = depot_sub.consume(nhalo) halo = numpy.empty(nhalo + 1, dtype=realhalo.dtype) halo[:-1] = realhalo # last halo is contamination halo[-1]['mass'] = numpy.nan halo[-1]['pos'] = numpy.nan halo[:]['groupid'] = realgroupid halo[-1]['len'] = group['len'][groupid] - halo['len'][:-1].sum() assert (halo['len'] >= 0).all() assert (halo['len'][0:-2] >= halo['len'][1:-1]).all() parstart = numpy.concatenate(([0], halo['len'].cumsum())) parend = parstart[1:] #make sure group and tab are of the same file assert group[groupid]['len'] == tab[0, 'length'][groupid] if not args.check: for haloid in range(nhalo + 1): parinhalo = pars[parstart[haloid]:parend[haloid]] parinhalo.sort(order='type') start = parinhalo['type'].searchsorted(range(6), side='left') end = parinhalo['type'].searchsorted(range(6), side='right') for type in range(6): thistype = parinhalo[start[type]:end[type]] halo[haloid]['massbytype'][type] = thistype[ 'mass'].sum(dtype='f8') halo[haloid]['lenbytype'][type] = len(thistype) for field in F[type]: thistype[field].tofile(F[type][field]) group[groupid]['nhalo'] = nhalo print i, 'group', groupid, 'nhalo', \ group['nhalo'][groupid], \ 'npar', group['len'][groupid], halo[-1]['len'], \ 'collective', snap.hascollective(realgroupid, group[groupid]) if not args.check: halo.tofile(halodump) halodump.flush() assert numpy.all( group[groupid]['lenbytype'] == halo['lenbytype'].sum( axis=0)) if not numpy.allclose(group[groupid]['massbytype'], halo['massbytype'].sum(axis=0)): print group[groupid]['massbytype'], halo['massbytype'].sum( axis=0) realgroupid += 1 if not args.check: group.tofile(groupdump) groupdump.flush() for type in range(6): for field in pdtype.fields: F[type][field].flush() assert depot_sub.Nhalo == depot_sub.total_read
def main(): camera = Camera(int(length), int(width)) camera.fade = False camera.lookat(pos=[ cub.newboxsize[0] * 0.5, cub.newboxsize[1] * 0.5, cub.newboxsize[2]], target=[ cub.newboxsize[0] * 0.5, cub.newboxsize[1] * 0.5, cub.newboxsize[2] * 0.5], up=[0, 1, 0]) camera.ortho(near=0, far=cub.newboxsize[2], extent=[-cub.newboxsize[0] * 0.5, cub.newboxsize[0] * 0.5, -cub.newboxsize[1] * 0.5, cub.newboxsize[1] * 0.5]) if not os.path.exists(ccdbuffer): CCD = numpy.memmap(ccdbuffer, mode='w+', dtype=('f4', 2), shape=camera.shape) CCD[...] = 0 def prefetcher(fid): print 'pref', fid try: snap = Snapshot(snapname % fid, 'cmugadget') snap[0, 'pos'] snap[0, 'ie'] snap[0, 'ye'] snap[0, 'sml'] snap[0, 'mass'] except: pass return chunksize = 64 def prefetch(fid): loadnext = [ sharedmem.background( lambda fidnext: prefetcher(fidnext), fidnext) for fidnext in range(fid, fid + chunksize) if fidnext < Nfile ] return loadnext loadnext = prefetch(0) for fid in range(0, Nfile, chunksize): print fid, 'start' for thread in loadnext: thread.wait() del loadnext nextsnap = [Snapshot(snapname % f, 'cmugadget') for f in range(fid, fid + chunksize) if f < Nfile] loadnext = prefetch(fid + chunksize) makegigapan(nextsnap, camera, CCD) del nextsnap print fid, 'done' CCD[..., 0] /= CCD[..., 1] else: CCD = numpy.memmap(ccdbuffer, mode='r', dtype=('f4', 2)) CCD = CCD.reshape(-1, 2) print 'CCD done' C, L = CCD[..., 0], CCD[..., 1] C = C.reshape(-1) L = L.reshape(-1) with sharedmem.MapReduce() as pool: chunksize = 1024 * 1024 * 20 def work(i): sl = slice(i, i + chunksize) l = nl_(L[sl]) return l.vmax, l.vmin vmax, vmin = numpy.array(pool.map(work, range(0, len(L), chunksize))).T vmax = numpy.max(vmax) - 0.2 vmin = numpy.min(vmin) print vmax, vmin im = numpy.memmap(imagename, mode='w+', dtype=('u1', 3), shape=L.shape) with sharedmem.MapReduce() as pool: chunksize = 1024 * 1024 * 128 def work(i): sl = slice(i, i + chunksize) im[sl] = image(nl_(C[sl], vmin=3.0, vmax=8.0), nl_(L[sl], vmax=vmax, vmin=vmin) ** 0.8, cmap=cm.RdBu_r)[..., :3] pool.map(work, range(0, len(L), chunksize)) im.flush()
from gaepsi.snapshot import Snapshot import sharedmem import numpy from sys import argv import os.path first = Snapshot(argv[1] % 0, 'cmugadget') def work(i): snap = Snapshot(argv[1] % i, 'cmugadget', template=first) print 'in', i sfr = snap[0, 'sfr'] id = snap[0, 'id'] mask = sfr > 0 print 'done', i, mask.sum(), sfr.max(), sfr.min() return mask.sum() with sharedmem.Pool() as pool: size = pool.map(work, range(first.C['Nfiles']), ordered=True) print 'phase 1 done' size = numpy.array(size, 'i8') total = size.sum() offset = numpy.concatenate(([0], size.cumsum())) print 'allocation total pars', total idall = sharedmem.empty(total, dtype=first.schema['id'].dtype) sfrall = sharedmem.empty(total, dtype='f4') def work(i):