def doall(base,problem_id,problem_dir=None,do_pickling=True,use_yt=True,
force_recal=False, force_redraw=False,vtkdir=None):
"""
This function will do following tasks:
(1) reoranizing files
* .hst, .sn --> hst/
* .zprof --> zprof/
* .starpar.vtk --> starpar/
* .rst --> rst/
(2) creating a parameter file from a restart header
(3) creating merged (along the time axis) zprof files for each thermal phase (using xarray)
(4) creating history data using zprof dumps, including new fields
(5) [if do_pickling] creating pickle files for slice and projection (will use yt [if use_yt])
"""
if problem_dir is None: problem_dir = problem_id
print('preparing metadata for {}...'.format(problem_id))
done=cleanup_directory(base,problem_id,problem_dir=problem_dir)
for k in done:
if done[k]: print('{}.{} is moved'.format(problem_id,k))
parfile='{}{}/{}.par'.format(base,problem_dir,problem_id)
params=pa.get_params(parfile)
zpfile='{}{}/zprof_merged/{}.phase5.zprof.nc'.format(base,problem_dir,problem_id)
from pyathena.utils import compare_files
zpfiles=glob.glob('{}{}/zprof/{}.*.whole.zprof'.format(base,problem_dir,problem_id))
zpfiles.sort()
if len(zpfiles) > 0:
if not compare_files(zpfiles[-1],zpfile): zprof_to_xarray(base,problem_dir,problem_id)
#zpfile='{}{}/zprof_merged/{}.zprof.nc'.format(base,problem_id,problem_id)
#if not os.path.isfile(zpfile): merge_xarray(base,problem_id)
hstzpfile='{}{}/hst/{}.hst_zp.p'.format(base,problem_dir,problem_id)
if os.path.isfile(zpfile): recal_history(base,problem_dir,problem_id)
if do_pickling:
kwargs={'base_directory':base,
'directory':'{}/'.format(problem_dir),
'id':problem_id,
'parallel':False,
'phase':False,
'rotation':params['Omega']*1.e3,
'range':''
}
if use_yt:
print('slicing and projecting with yt ...')
from pyathena.yt_analysis import yt_analysis
yt_analysis.main(force_recal=force_recal,force_redraw=force_redraw,verbose=50,**kwargs)
else:
if not (vtkdir is None): kwargs['vtk_directory']=vtkdir
print('slicing and projecting with pyathena ...')
from pyathena.create_pickle import create_all_pickles
create_all_pickles(force_recal=force_recal,force_redraw=force_redraw,verbose=True,**kwargs)
def doall(base, problem_id, problem_dir=None, do_pickling=True, use_yt=True):
"""
do all preprocessing for a model
"""
if problem_dir is None: problem_dir = problem_id
print('preparing metadata for {}...'.format(problem_id))
done = cleanup_directory(base, problem_id, problem_dir=problem_dir)
for k in done:
if done[k]: print('{}.{} is moved'.format(problem_id, k))
parfile = '{}{}/{}.par'.format(base, problem_dir, problem_id)
params = pa.get_params(parfile)
zpfile = '{}{}/zprof_merged/{}.whole.zprof.nc'.format(
base, problem_dir, problem_id)
from pyathena.utils import compare_files
zpfiles = glob.glob('{}{}/zprof/{}.*.whole.zprof'.format(
base, problem_dir, problem_id))
zpfiles.sort()
if not compare_files(zpfiles[-1], zpfile):
zprof_to_xarray(base, problem_dir, problem_id)
#zpfile='{}{}/zprof_merged/{}.zprof.nc'.format(base,problem_id,problem_id)
#if not os.path.isfile(zpfile): merge_xarray(base,problem_id)
hstzpfile = '{}{}/hst/{}.hst_zp.p'.format(base, problem_dir, problem_id)
if os.path.isfile(zpfile): recal_history(base, problem_dir, problem_id)
if do_pickling:
kwargs = {
'base_directory': base,
'directory': '{}/'.format(problem_dir),
'id': problem_id,
'parallel': False,
'phase': False,
'rotation': params['Omega'] * 1.e3,
'range': ''
}
if use_yt:
print('slicing and projecting with yt ...')
from pyathena.yt_analysis import yt_analysis
yt_analysis.main(force_recal=False,
force_redraw=False,
verbose=50,
**kwargs)
else:
print('slicing and projecting with pyathena ...')
from pyathena.create_pickle import create_all_pickles
create_all_pickles(force_recal=False,
force_redraw=False,
verbose=True,
**kwargs)
field_list = [
'star_particles', 'surface_density', 'nH', 'temperature', 'pok',
'velocity_z'
]
slcdata = p.load(open(slc_files[0], 'rb'), encoding='latin1')
if 'magnetic_field_strength' in slcdata['x']:
if system == 'tigress_rps':
field_list += ['mag_pok']
else:
field_list += ['magnetic_field_strength']
for slcname in slc_files:
print(slcname)
starname = slcname.replace('slice.p',
'starpar.vtk').replace('slice', 'starpar')
projname = slcname.replace('slice', 'surf')
if not compare_files(slcname, slcname.replace(
'.p', '_proj.png')) or overwrite:
#plot_slices.slice2(slcname,starname,field_list,aux=aux,vy0=vy0)
plot_slices.slice_proj(slcname,
projname,
starname,
field_list,
aux=aux,
vy0=vy0)
#if not compare_files(projname,projname+'ng') or overwrite:
plot_projection.plot_projection(projname,
starname,
scale_func=np.cbrt,
runaway=False,
aux=aux_surf,
vy0=vy0)
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 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)