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
