def reader(fname,vel=True,mhd=True): coolftn=pa.coolftn() dir, id, step, ext, mpi = pa.parse_filename(fname) ds=pa.AthenaDataSet(fname) rstfnames=glob.glob('%s/%s.par' % (ds.dir,ds.id)) par,blocks,fields=pa.parse_par(rstfnames[0]) domain=ds.domain domain['qshear']=eval(par['problem']['qshear'][0]) domain['Omega']=eval(par['problem']['Omega'][0]) fields=['density'] if vel: fields.append('velocity1') fields.append('velocity2') fields.append('velocity3') if mhd: fields.append('magnetic_field1') fields.append('magnetic_field2') fields.append('magnetic_field3') data={} for f in fields: data[f] = ds.read_all_data(f) data['temperature']=coolftn.get_temp(ds.read_all_data('T1')) fields.append('temperature') if vel: r3d,x3d,y3d,z3d=pa.pos3d(domain) vy0=-domain['qshear']*domain['Omega']*x3d data['velocity2'] -= vy0 domain['losdir']=dir+'los/' domain['step']=step return data,domain
def setup_domain(fname,vel=True,mhd=True,shear=True): dir, id, step, ext, mpi = pa.parse_filename(fname) ds=pa.AthenaDataSet(fname) rstfnames=glob.glob('%s/%s.par' % (ds.dir,ds.id)) par,blocks,fields=pa.parse_par(rstfnames[0]) domain=ds.domain domain['qshear']=eval(par['problem']['qshear'][0]) domain['Omega']=eval(par['problem']['Omega'][0]) fields=['density'] fields.append('temperature') if vel: fields.append('velocity1') fields.append('velocity2') fields.append('velocity3') if mhd: fields.append('magnetic_field1') fields.append('magnetic_field2') fields.append('magnetic_field3') domain['fields']=fields domain['shear']=shear if shear: domain['losdir']=dir+'los/' else: domain['losdir']=dir+'los-periodic/' domain['step']=step return ds,domain
def run_R8(base,pid): for i in range(700): print("{} of {}...".format(i,700)) ds=pa.AthenaDataSet('{}/{}/id0/{}.{:04d}.vtk'.format(base,pid,pid,i)) dproj,Tproj,Bproj=get_proj_ims(ds,slab_range=[9,21]) Nz,Nx=dproj.shape ratio=Nz/float(Nx) plt.style.use('default') fig=plt.figure(0,figsize=(15,15*ratio)) draw_merged_proj(fig,dproj,Tproj,Bproj) fig.savefig('{}{}/proj_figures/{}.merged_proj.{:04d}.png'.format(base,pid,pid,i),dpi=100) fig.clf()
Ther_e = [] Ther_l = [] Tot_e = [] Tot_l = [] Ext_e = [] Ext_l = [] Sg_e = [] Sg_l = [] #plt.figure(figsize=(6, 10)) crit = 94 for tidx in range(250, stop): # time step 251, 331, 411, 501 # plt.figure(figsize=(6, 5)) vtkfname = '%s%s/id0/%s.%04d.vtk' % (basedir, simid, simid, tidx) # read in domain information ds = pa.AthenaDataSet(vtkfname) #print(ds.field_list) # print(ds.derived_field_list) # full domain information # print ds.domain # rs = ds.read_all_data('reynold_stress') # print rs # information of grid #0 # print ds.grids[0] T1 = ds.read_all_data('T1') coolftn = pa.coolftn() temp = coolftn.get_temp(T1)
def create_all_pickles(force_recal=False, force_redraw=False, verbose=True, **kwargs): dir = kwargs['base_directory'] + kwargs['directory'] fname = glob.glob(dir + 'id0/' + kwargs['id'] + '.????.vtk') fname.sort() if kwargs['range'] != '': sp = kwargs['range'].split(',') start = eval(sp[0]) end = eval(sp[1]) fskip = eval(sp[2]) else: start = 0 end = len(fname) fskip = 1 fname = fname[start:end:fskip] ngrids = len(glob.glob(dir + 'id*/' + kwargs['id'] + '*' + fname[0][-8:])) ds = pa.AthenaDataSet(fname[0]) mhd = 'magnetic_field' in ds.field_list cooling = 'pressure' in ds.field_list Omega = kwargs['rotation'] rotation = kwargs['rotation'] != 0. if verbose: print("MHD:", mhd) print("cooling:", cooling) print("rotation:", rotation, Omega) slc_fields = ['nH', 'pok', 'temperature', 'velocity_z', 'ram_pok_z'] fields_to_draw = [ 'star_particles', 'nH', 'temperature', 'pok', 'velocity_z' ] if mhd: slc_fields.append('magnetic_field_strength') slc_fields.append('mag_pok') fields_to_draw.append('magnetic_field_strength') mul_factors = { 'pok': to_Pok, 'magnetic_field_strength': to_microG, 'mag_pok': to_Pok, 'ram_pok_z': to_Pok } scal_fields = get_scalars(ds) slc_fields += scal_fields if not os.path.isdir(dir + 'slice/'): os.mkdir(dir + 'slice/') if not os.path.isdir(dir + 'surf/'): os.mkdir(dir + 'surf/') for i, f in enumerate(fname): slcfname = dir + 'slice/' + kwargs['id'] + f[-9:-4] + '.slice.p' surfname = dir + 'surf/' + kwargs['id'] + f[-9:-4] + '.surf.p' tasks = { 'slice': (not compare_files(f, slcfname)) or force_recal, 'surf': (not compare_files(f, surfname)) or force_recal, } do_task = (tasks['slice'] or tasks['surf']) if verbose: print('file number: {} -- Tasks to be done ['.format(i), end='') for k in tasks: print('{}:{} '.format(k, tasks[k]), end='') print(']') if do_task: ds = pa.AthenaDataSet(f) if tasks['surf']: create_projection( ds, surfname, conversion={'z': ds.domain['Lx'][2] * to_surf}) if tasks['slice']: create_slices(ds, slcfname, slc_fields, factors=mul_factors, force_recal=force_recal) aux = set_aux(kwargs['id']) for i, f in enumerate(fname): slcfname = dir + 'slice/' + kwargs['id'] + f[-9:-4] + '.slice.p' surfname = dir + 'surf/' + kwargs['id'] + f[-9:-4] + '.surf.p' starpardir = 'id0/' if os.path.isdir(dir + 'starpar/'): starpardir = 'starpar/' starfname = dir + starpardir + kwargs['id'] + f[-9:-4] + '.starpar.vtk' tasks = { 'slice': (not compare_files(f, slcfname + 'ng')) or force_redraw, 'surf': (not compare_files(f, surfname + 'ng')) or force_redraw, } do_task = (tasks['slice'] and tasks['surf']) if verbose: print('file number: {} -- Tasks to be done ['.format(i), end='') for k in tasks: print('{}:{} '.format(k, tasks[k]), end='') print(']') if tasks['surf']: plot_projection(surfname, starfname, runaway=False, aux=aux['surface_density']) if tasks['slice']: plot_slice(slcfname, starfname, fields_to_draw, aux=aux)
def dump_zprof_one(f,dc,icm_field=None,outdir='zprof_icm', plist=['phase1','phase2','phase3','phase4','phase5']): par=dc.par hst=dc.hst flist=set_field_list(par) full_domain=get_full_domain(par) x,y,z,=pa.cc_arr(full_domain) if icm_field is not None: icm=True print('Reading: ',f) ds = pa.AthenaDataSet(f) time = ds.domain['time'] zp_slab=[] if icm: zp_slab_icm=[] print('Calculating zprof...') for islab in range(ds.NGrids[2]): print('{}/{}'.format(islab,ds.NGrids[2]),) grids=ds._get_slab_grid(islab+1,verbose=False) slab_domain=ds._setup_domain(grids) xs,ys,zs,=pa.cc_arr(slab_domain) data=xr.Dataset() for fi in ds.field_list+['T1']: slab_data=ds.read_all_data(fi,slab=islab+1) if slab_data.ndim == 4: for ivec in [1,2,3]: data['{}{}'.format(fi,ivec)]=xr.DataArray(slab_data[...,ivec-1], dims=['z','y','x'], coords=[zs,ys,xs]) else: data[fi]=xr.DataArray(slab_data,dims=['z','y','x'],coords=[zs,ys,xs]) z_part,zpw_part=get_zprof(data,slab_domain,par,hst) dset=xr.Dataset() for field in zpw_part: dset[field]=xr.DataArray(zpw_part[field], dims=['z','y','x'],coords=[zs,ys,xs]) idx=get_phase(dset) dA=float(dset['A'][0,0,0].data) zp_part=[] if icm: ficm=np.clip(dset[icm_field]/dset['d'],0,1) zp_part_icm=[] for phase in plist: zidx_part=pd.Series(z_part,name='z') zp_part.append(get_mean(dset*idx[phase],dA).expand_dims('phase')) if icm: zp_part_icm.append(get_mean(dset*idx[phase]*ficm,dA).expand_dims('phase')) zp_slab.append(xr.concat(zp_part,dim='phase')) if icm: zp_slab_icm.append(xr.concat(zp_part_icm,dim='phase')) for g in grids: g['data'].clear() zpds=xr.concat(zp_slab,dim='z') zpds.coords['phase']=plist if icm: zpds_icm=xr.concat(zp_slab_icm,dim='z') zpds_icm.coords['phase']=plist print('\nWriting at {}: '.format(os.path.dirname(f).replace('id0',outdir))) for phase in plist: print(phase,) zprof_fname=f.replace('vtk','%s.zprof' % phase).replace('id0',outdir) with open(zprof_fname,'w') as fp: fp.write('# Athena vertical profile at t={}\n'.format(time)) zpdf=zpds.sel(phase=phase).drop('phase').to_dataframe() zpdf.to_csv(fp) if icm: zprof_fname_icm=f.replace('vtk','%s-icm.zprof' % phase).replace('id0',outdir) with open(zprof_fname_icm,'w') as fp: fp.write('# Athena vertical profile at t={}\n'.format(time)) zpdf=zpds_icm.sel(phase=phase).drop('phase').to_dataframe() zpdf.to_csv(fp) if icm: return zpds,zpds_icm else: return zpds
def main(**kwargs): base = kwargs['base_directory'] id = kwargs['id'] if len(kwargs['directory']) > 0: dir = kwargs['directory'] else: dir = id join_vtk = kwargs['join_vtk'] if kwargs['var'] != '': var = kwargs['var'] join_vtk=os.path.dirname(join_vtk)+'/'+\ os.path.basename(join_vtk).replace('join_vtk','join_vtk_var') itime = kwargs['itime'] istart = kwargs['istart'] iend = kwargs['iend'] + 1 fname = '{}{}/id0/{}.{:04d}.vtk'.format(base, dir, id, itime) if kwargs['new_id'] != '': newid = kwargs['new_id'] else: newid = id if kwargs['new_base_directory'] != '': newbase = '%s%s/' % (kwargs['new_base_directory'], newid) else: newbase = '%s%s/merged/' % (base, dir) if not os.path.isdir(newbase): os.mkdir(newbase) parfile = '%s%s/%s.par' % (base, dir, id) if not os.path.isfile(parfile): parfile = '%s/%s.par' % (newbase, newid) par = get_params(parfile) NGrids=[int(par['NGrid_x1']),\ int(par['NGrid_x2']),\ int(par['NGrid_x3'])] Nslab = NGrids[2] Nproc = np.prod(NGrids) Nproc_h = NGrids[0] * NGrids[1] gid = np.arange(Nproc) print(Nproc, NGrids) print(fname) fpath, fbase, fstep, fext, mpi = pa.parse_filename(fname) grids = gid[istart:iend] baseid = newid command = [join_vtk] outfile = '%s/%s.%s.vtk' % (newbase, baseid, fstep) if kwargs['var'] != '': command.append('-o %s' % outfile.replace('vtk', kwargs['var'][0] + '.vtk')) command.append('-v %s' % var) else: command.append('-o %s' % outfile) zmin = 1.e10 zmax = -1.e10 for gidx in grids: if gidx == 0: vtkfile = '%s%s/id%d/%s.%s.%s' % (base, dir, gidx, id, fstep, fext) else: vtkfile = '%s%s/id%d/%s-id%d.%s.%s' % (base, dir, gidx, id, gidx, fstep, fext) ds = pa.AthenaDataSet(vtkfile, serial=True) zmin = min(ds.domain['left_edge'][2], zmin) zmax = max(ds.domain['right_edge'][2], zmax) command.append(vtkfile) #print command print(' '.join(command)) print('id={} to {} corresponds to z={} to {}'.format( istart, iend, zmin, zmax))
def create_all_pickles( datadir, problem_id, nums=None, fields_slc=['nH', 'nHI', 'temperature', 'xn', 'ne', 'nesq', 'Erad0', 'Erad1'], fields_proj=['rho', 'xn', 'nesq'], fields_draw=['star_particles', 'rho_proj', 'xn_proj', 'nesq_proj', 'nH', 'temperature', 'xn', 'Erad0', 'Erad1'], force_recal=False, force_redraw=False, no_save=False, savdir=None, verbose=True, **plt_args): """ Function to pickle slices and projections from AthenaDataset and draw snapshots. Set force_recal to True if additional fields need to be extracted. Parameters ---------- datadir: string Base data directory problem_id: string Prefix for vtk files num: array of integers List of vtk output numbers. Search all vtk files in the directory if None. fields_slc: list of strings List of field names to be sliced fields_proj: list of strings List of field names to be projected fields_draw: list of strings List of field names to be drawn force_recal: bool If True, override existing pickles. force_redraw: bool If True, override existing figures. no_save: bool If True, returns a list of matplotlib figure objects instead of saving them. savdir: str Directory to save snapshot. If None, saves in snapshot subdirectory under datadir. Default value is None. verbose: bool Print verbose message Returns ------- fig: figures Returns lists of figure if no_save is True. """ aux = set_aux(problem_id) _plt_args = dict(zoom=1.0) _plt_args.update(**plt_args) fglob = os.path.join(datadir, problem_id + '.????.vtk') fname = sorted(glob.glob(fglob)) if not fname: fglob = os.path.join(datadir, 'vtk', problem_id + '.????.vtk') fname = glob.glob(fglob) if not fname: fglob = os.path.join(datadir, 'vtk', 'id0', problem_id + '.????.vtk') fname = glob.glob(fglob) if not fname: fglob = os.path.join(datadir, 'id0', problem_id + '.????.vtk') fname = glob.glob(fglob) fname.sort() if not fname: print('No vtk files are found in {0:s}'.format(datadir)) raise if nums is None: nums = [int(f[-8:-4]) for f in fname] if nums[0] == 0: # remove the zeroth snapshot start = 1 del nums[0] else: start = 0 end = len(fname) fskip = 1 fname = fname[start:end:fskip] else: nums = np.atleast_1d(nums) fname = [fname[i] for i in nums] #ngrids = len(glob.glob(datadir+'id*/' + id + '*' + fname[0][-8:])) ds = pa.AthenaDataSet(fname[0]) mhd = 'magnetic_field' in ds.field_list cooling = 'pressure' in ds.field_list print('[Create_pickle_all_rad]') print('- basedir:', datadir) print('- problem id:', problem_id) print('- vtk file num:', end=' ') for i in nums: print(i,end=' ') print('') print('slc: {0:s}'.format(' '.join(fields_slc))) print('proj: {0:s}'.format(' '.join(fields_proj))) print('draw: {0:s}'.format(' '.join(fields_draw))) if mhd: slc_fields.append('magnetic_field_strength') slc_fields.append('mag_pok') draw_fields.append('magnetic_field_strength') mul_factors = {'pok':to_Pok, 'magnetic_field_strength':to_microG, 'mag_pok':to_Pok, 'ram_pok_z':to_Pok} if not os.path.isdir(os.path.join(datadir, 'slice')): os.mkdir(os.path.join(datadir, 'slice')) if not os.path.isdir(os.path.join(datadir, 'proj')): os.mkdir(os.path.join(datadir, 'proj')) print('\n*** Extract slices and projections ***') print('- num: ',end='') for i, f in enumerate(fname): print('{}'.format(int(f.split('.')[-2])), end=' ') fname_slc = os.path.join(datadir, 'slice', problem_id + f[-9:-4] + '.slice.p') fname_proj = os.path.join(datadir, 'proj', problem_id + f[-9:-4] + '.proj.p') tasks = dict(slc=(not compare_files(f,fname_slc)) or force_recal, proj=(not compare_files(f,fname_proj)) or force_recal) do_task = (tasks['slc'] or tasks['proj']) if do_task: ds = pa.AthenaDataSet(f) if tasks['slc']: _ = create_slices(ds, fname_slc, fields_slc, factors=mul_factors, force_recal=force_recal, verbose=verbose) if tasks['proj']: _ = create_projections(ds, fname_proj, fields_proj, aux=aux, force_recal=force_recal, verbose=verbose) del ds gc.collect() #print(fname) print('') print('*** Draw snapshots (zoom {0:.1f}) ***'.format(_plt_args['zoom'])) print('num: ',end='') if no_save: force_redraw = True figs = [] if savdir is None: savdir = os.path.join(datadir, 'snapshots') if not os.path.isdir(savdir): os.mkdir(savdir) print('savdir:', savdir) for i,f in enumerate(fname): num = f.split('.')[-2] print('{}'.format(int(num)), end=' ') fname_slc = os.path.join(datadir, 'slice', problem_id + f[-9:-4] + '.slice.p') fname_proj = os.path.join(datadir, 'proj', problem_id + f[-9:-4] + '.proj.p') starpardir = 'id0' if os.path.isdir(os.path.join(datadir, 'starpar')): starpardir = 'starpar' fname_sp = os.path.join(datadir, starpardir, problem_id + f[-9:-4] + '.starpar.vtk') savname = os.path.join(savdir, problem_id + '.' + num + '.slc_proj.png') if _plt_args['zoom'] == 1.0: savname = os.path.join(savdir, problem_id + '.' + num + '.slc_proj.png') else: # append zoom factor savname = os.path.join(savdir, problem_id + '.' + num + '.slc_proj-' + \ 'zoom{0:02d}'.format(int(10.0*_plt_args['zoom'])) + '.png') tasks = dict(slc_proj=(not compare_files(f, savname)) or force_redraw, proj=(not compare_files(f, fname_proj+'ng')) or force_redraw) do_task = (tasks['slc_proj'] and tasks['proj']) if tasks['proj']: plot_projection(fname_proj, fname_sp, 'rho', runaway=True, aux=aux['rho_proj']) if tasks['slc_proj']: if no_save: savname = None fig = plot_slice_proj(fname_slc, fname_proj, fname_sp, fields_draw, savname, aux=aux, **_plt_args) figs.append(fig) else: plot_slice_proj(fname_slc, fname_proj, fname_sp, fields_draw, savname, aux=aux, **_plt_args) print('') print('*** Done! ***') if no_save: return tuple(figs)
def main(**kwargs): base = kwargs['base_directory'] id = kwargs['id'] if len(kwargs['directory']) > 0: dir = kwargs['directory'] else: dir = id join_vtk= kwargs['join_vtk'] files=glob.glob('%s%s/id0/%s.????.vtk' % (base,dir,id)) files.sort() if kwargs['range'] != '': sp=kwargs['range'].split(',') start = int(sp[0]) end = int(sp[1]) fskip = int(sp[2]) else: start = 0 end = len(files) fskip = 1 files=files[start:end:fskip] if kwargs['new_id'] != '': newid=kwargs['new_id'] else: newid=id if kwargs['new_base_directory'] != '': newbase='%s%s/' % (kwargs['new_base_directory'],newid) else: newbase='%s%s/slab/' % (base,dir) if not os.path.isdir(newbase): os.mkdir(newbase) if not os.path.isdir(newbase+'/starpar'): os.mkdir(newbase+'/starpar') if not os.path.isdir(newbase+'/zprof'): os.mkdir(newbase+'/zprof') if not os.path.isdir(newbase+'/hst'): os.mkdir(newbase+'/hst') if not os.path.isdir(newbase+'/rst'): os.mkdir(newbase+'/rst') rstfiles=glob.glob('%s%s/id0/%s.????.rst' % (base,dir,id)) rstfiles+=glob.glob('%s%s/rst/%s.????.rst' % (base,dir,id)) parfile='%s/%s.par' % (newbase,newid) if os.path.isfile(parfile): print(('par file is already there for %s!' % newid)) else: if len(rstfiles): write_par_from_rst(rstfiles[0],parfile) par=get_params(parfile) ds = pa.AthenaDataSet(files[0],serial=True) NGrids=[int(par['Nx1']/ds.domain['Nx'][0]),\ int(par['Nx2']/ds.domain['Nx'][1]),\ int(par['Nx3']/ds.domain['Nx'][2])] Nslab=NGrids[2] Nproc=np.prod(NGrids) Nproc_h=NGrids[0]*NGrids[1] gid=np.arange(Nproc) print((Nproc,NGrids)) # copy history fpath,fbase,fstep,fext,mpi=parse_filename(files[0]) src_hst_name='%s/id0/%s.hst' % (fpath,fbase) dst_name='%s/hst/%s.hst' % (newbase,newid) if os.path.isfile(src_hst_name): shutil.copy(src_hst_name,dst_name) src_hst_name='%s/id0/%s.sn' % (fpath,fbase) dst_name='%s/hst/%s.sn' % (newbase,newid) if os.path.isfile(src_hst_name): shutil.copy(src_hst_name,dst_name) if fpath != newbase: if not os.path.isdir('%s/id0' % (newbase)): os.mkdir('%s/id0' % (newbase)) src_hst_name='%s/id0/%s.hst' % (fpath,fbase) dst_name='%s/id0/%s.hst' % (newbase,newid) if os.path.isfile(src_hst_name): shutil.copy(src_hst_name,dst_name) src_hst_name='%s/id0/%s.sn' % (fpath,fbase) dst_name='%s/id0/%s.sn' % (newbase,newid) if os.path.isfile(src_hst_name): shutil.copy(src_hst_name,dst_name) for f in files: print(f) fpath,fbase,fstep,fext,mpi=parse_filename(f) remove_flag=True ds = pa.AthenaDataSet(f,serial=True) NGrids=[int(par['Nx1']/ds.domain['Nx'][0]),\ int(par['Nx2']/ds.domain['Nx'][1]),\ int(par['Nx3']/ds.domain['Nx'][2])] Nslab=NGrids[2] Nproc=np.prod(NGrids) Nproc_h=NGrids[0]*NGrids[1] gid=np.arange(Nproc) print((f,Nproc,Nproc_h,NGrids)) for islab in range(Nslab): print(('%d of %d' % (islab, Nslab))) grids=gid[(gid/Nproc_h).astype('int') == islab] if islab == 0: baseid=newid else: baseid='%s-id%d' %(newid,islab) if not os.path.isdir('%s/id%d' % (newbase,islab)): os.mkdir('%s/id%d' % (newbase,islab)) outfile='%s/id%d/%s.%s.vtk' % (newbase,islab,baseid,fstep) if len(grids) > 1: command=[join_vtk] command.append('-o %s' % outfile) else: print(('%s is already merged' % (outfile))) command=['mv'] for gidx in grids: if gidx == 0: vtkfile='%s%s/id%d/%s.%s.%s' % (base,dir,gidx,id,fstep,fext) else: vtkfile='%s%s/id%d/%s-id%d.%s.%s' % (base,dir,gidx,id,gidx,fstep,fext) command.append(vtkfile) if len(grids) == 1: command.append(outfile) remove_flag=False #print command if not compare_files(vtkfile,outfile) or kwargs['overwrite']: subprocess.call(' '.join(command),shell=True) else: print(('%s is newer than %s' % (outfile, vtkfile))) remove_flag=False if not os.path.isfile(outfile): print(('join to %s is failed' % (outfile))) remove_flag=False # delete originals file_originals=glob.glob('%s/id*/%s-id*.%s.%s' % (fpath,fbase,fstep,fext)) if (len(file_originals) > 0) and remove_flag: for f in file_originals: os.remove(f) os.remove('%s/id0/%s.%s.%s' % (fpath,fbase,fstep,fext)) # move starpar.vtk src_starpar_name='%s/id0/%s.%s.starpar.vtk' % (fpath,fbase,fstep) dst_name='%s/starpar/%s.%s.starpar.vtk' % (newbase,newid,fstep) if os.path.isfile(src_starpar_name): shutil.move(src_starpar_name,dst_name) # move zprof src_zprof_names=glob.glob('%s/id0/%s.%s.*.zprof' % (fpath,fbase,fstep)) for f in src_zprof_names: dst_name=f.replace(fpath,newbase).replace('id0/','zprof/').replace(fbase,newid) print(dst_name) if os.path.isfile(f): shutil.move(f,dst_name) subprocess.call('find %s/id* -name *.rst -exec mv {} %s/rst/ \;' % (fpath,newbase),shell=True) subprocess.call('rename %s %s %s/rst/*' % (id,newid,newbase),shell=True)
def data_read(fname,vel=True,mhd=True,usefits=True): dir, id, step, ext, mpi = pa.parse_filename(fname) if usefits: from astropy.io import fits fitsname='%s/fits/%s.%s.%s.fits' % (dir,id,step,'dens') hdulist=fits.open(fitsname,memmap=True) hdr=hdulist[0].header domain={} domain['time']=hdr['time'] domain['qshear']=hdr['qshear'] domain['Omega']=hdr['Omega'] domain['left_edge']=np.array([hdr['xmin'],hdr['ymin'],hdr['zmin']]) domain['right_edge']=np.array([hdr['xmax'],hdr['ymax'],hdr['zmax']]) domain['dx']=np.array([hdr['dx'],hdr['dy'],hdr['dz']]) domain['Nx']=np.array([hdr['naxis1'],hdr['naxis2'],hdr['naxis3']]) domain['Lx']=domain['right_edge']-domain['left_edge'] fields=['density','temperature'] if vel: fields.append('velocity') if mhd: fields.append('magnetic_field') hdulist.close() else: ds=pa.AthenaDataSet(fname) rstfnames=glob.glob('%s/id0/%s.????.rst' % (ds.dir,ds.id)) \ +glob.glob('%s/rst/%s.????.rst' % (ds.dir,ds.id)) par,blocks,fields=pa.parse_par(rstfnames[0]) domain=ds.domain domain['qshear']=eval(par['problem']['qshear'][0]) domain['Omega']=eval(par['problem']['Omega'][0]) fields=['density'] if vel: fields.append('velocity1') fields.append('velocity2') fields.append('velocity3') if mhd: fields.append('magnetic_field1') fields.append('magnetic_field2') fields.append('magnetic_field3') print domain['time'] data={} if usefits: for f in fields: fitsname='%s/fits/%s.%s.%s.fits' % (dir,id,step,f[:4]) if f.startswith('velo') or f.startswith('magn'): hdulist=fits.open(fitsname,memmap=False,lazy_load_hdus=False) hdulist.readall() #fitsdata=fits.getdata(fitsname) for iaxis in range(3): data['%s%d' % (f,iaxis+1)]=hdulist[0].data[:,:,:,iaxis] #data['%s%d' % (f,iaxis+1)]=fitsdata[:,:,:,iaxis] hdulist.close() else: hdulist=fits.open(fitsname,memmap=False,lazy_load_hdus=False) hdulist.readall() #fitsdata=fits.getdata(fitsname) data[f]=hdulist[0].data #data[f]=fitsdata else: for f in fields: data[f] = ds.read_all_data(f) data['temperature']=coolftn.get_temp(ds.read_all_data('T1')) fields.append('temperature') if vel: r3d,x3d,y3d,z3d=pa.pos3d(domain) vy0=-domain['qshear']*domain['Omega']*x3d data['velocity2'] -= vy0 print data.keys() return data,domain