def get_halo_particles(my_halo, par_file):
pf = yt.load(par_file)
if yt.is_root():
print ("Reading in particles for halo %d." % my_halo['id'])
halo_files = glob.glob(os.path.join(pf.fullpath, 'MergerHalos_*.h5'))
particle_indices = np.array([])
particle_masses = np.array([])
pos_x = np.array([])
pos_y = np.array([])
pos_z = np.array([])
halo_name = 'Halo%08d' % my_halo['id']
for halo_file in halo_files:
input = h5.File(halo_file, 'r')
if halo_name in input.keys():
particle_indices = np.concatenate([particle_indices, input[halo_name]['particle_index'].value])
particle_masses = np.concatenate([particle_masses, input[halo_name]['ParticleMassMsun'].value])
pos_x = np.concatenate([pos_x, input[halo_name]['particle_position_x'].value])
pos_y = np.concatenate([pos_y, input[halo_name]['particle_position_y'].value])
pos_z = np.concatenate([pos_z, input[halo_name]['particle_position_z'].value])
input.close()
particle_positions = np.array([pos_x, pos_y, pos_z])
if particle_indices is None:
if yt.is_root():
print ("Error: could not locate halo %d." % my_halo['id'])
return None
return (particle_indices, particle_masses, particle_positions)
def initialize():
try:
array = pickle.load(open('minarray.p', 'r'))
if yt.is_root():
print 'minarray.p loaded'
except:
if yt.is_root():
print 'minarray.p not found. Run yt_get_mins.py first'
return
flag = array['Eintmin'] < Eint_threshold
global toplot
toplot = array['basename'][flag]
def phase(sq, phfname, bin_fields, aux={}):
pdfs = my_pdf(sq)
for bf in bin_fields:
nbin1, nbin2 = (128, 128)
if bf[0] in aux: nbin1 = aux[bf[0]]['n_bins']
if bf[1] in aux: nbin2 = aux[bf[1]]['n_bins']
n_bins = (nbin1, nbin2)
logs = {}
unit = {}
extrema = {}
for b in bf:
logs[b] = False
if b in aux:
if 'log' in aux[b]: logs[b] = aux[b]['log']
if 'unit' in aux[b]: unit[b] = aux[b]['unit']
if 'limints' in aux[b]: extrema[b] = aux[b]['limits']
pdf = yt.create_profile(sq,
bf,
fields=fields,
n_bins=n_bins,
logs=logs,
extrema=extrema,
units=unit,
weight_field=None,
fractional=True)
pdfs.add_pdf(pdf, bf)
if yt.is_root():
pickle.dump(pdfs, open(phfname, 'wb'), pickle.HIGHEST_PROTOCOL)
def calculate_mass_in_sphere(dd):
data_sub_dir = dd.name
a = dd.a
r_plus = M*(1 + math.sqrt(1 - a**2))
min_radius = r_plus/4
start_time = time.time()
# load dataset time series
dataset_path = data_root_path + "/" + data_sub_dir + "/KerrSFp_*.3d.hdf5"
ds = yt.load(dataset_path) # this loads a dataset time series
print("loaded data from ", dataset_path)
print("time = ", time.time() - start_time)
N = len(ds)
# set centre
center = [512.0, 512.0, 0]
L = 512.0
data_storage = {}
# iterate through datasets (forcing each to go to a different processor)
for sto, dsi in ds.piter(storage=data_storage):
time_0 = time.time()
# store time
current_time = dsi.current_time
output = [current_time]
# make shell
shell = dsi.sphere(center, max_R+0.5*dR) - dsi.sphere(center, max_R-0.5*dR)
# calculate energy inside sphere
meanJr = shell.mean("J_r", weight="cell_volume")
if half_box:
area = 2*np.pi*max_R**2
flux = meanJr*area
output.append(flux)
# store output
sto.result = output
sto.result_id = str(dsi)
dt = 2.5
i = int(current_time/dt)
print("done {:d} of {:d} in {:.1f} s".format(i+1, N, time.time()-time_0), flush=True)
if yt.is_root():
# make data directory if it does not already exist
makedirs(home_path + output_dir, exist_ok=True)
# output to file
dd.filename = "{:s}_mass_flux_at_R={:d}.csv".format(dd.name, max_R)
output_path = home_path + output_dir + "/" + dd.filename
# output header to file
f = open(output_path, "w+")
f.write("# t mass flux at R=" + str(max_R) + " #\n")
# output data
for key in sorted(data_storage.keys()):
data = data_storage[key]
f.write("{:.3f} ".format(data[0]))
f.write("{:.2f}\n".format(data[1]))
f.close()
print("saved data to file " + str(output_path))
def plot_vel_fil(vel_prof,dists,y):
#Generates plot of velocities along a filament
#Generate plot
if yt.is_root():
plot = plt.figure()
#Set up x and y values, and correctly fill and ensure shape is correct
x = np.empty((dists.shape[0],y.shape[0]))
yy = np.empty((vel_prof.shape[0],y.shape[0]))
x = x.T
for i in xrange(vel_prof.shape[0]):
yy[i] = y
for i in xrange(y.shape[0]):
x[i] = dists
x = x.T
#Ensure our colorbar is in correct form
colormax = np.nanmax(vel_prof)
colormin = abs(np.nanmin(vel_prof))
colormax = max([colormax,colormin])
colormin = 0 - colormax
print( x.shape, yy.shape, vel_prof.shape)
#Create a 'probability' map plot
vel_map = plt.pcolormesh(x,yy,vel_prof,cmap='seismic', vmin =colormin, vmax=colormax)
#Label graph
plt.xlabel("Distance along filament (Mpc)")
plt.ylabel("Radius from Filament (Mpc)")
cbar = plt.colorbar(vel_map)
cbar.ax.set_ylabel("Velocity (km/s)")
return plot
def apply_negative_radial_fit(name,Ncores):
KK=np.load(name+'_KK.npy')
KT=KK[0]
KD=KK[1]
Nbins=len(KT)
my_storage = {}
maps=[]
print('computing optimization')
indices=range(Nbins)
for sto, kt in yt.parallel_objects(set(KT), Ncores, storage = my_storage):
output = temporal_turbulence(name,kt)
sto.result_id = output
sto.result = kt
print('applying optimization')
my_storage2 = {}
for sto, ij in yt.parallel_objects(indices, Ncores, storage = my_storage2):
mapa=apply_temp(name,KT,KD,ij,Nbins)
sto.result_id = ij
sto.result = mapa
if yt.is_root():
for fn, vals in sorted(my_storage2.items()):
if ij==0:
mapa=vals
else:
mapa+=vals
np.save(name+'_negative_fit_%02d' %Nbins,mapa)
def Circulation_negative_turbulence(name,Nbins,start,kmin,kmax,Nk,Ncores):
ktarray=np.exp(np.linspace(np.log(kmin),np.log(kmax),Nk))
my_storage = {}
kdmatrix=[]
errmatrix=[]
for sto, kt in yt.parallel_objects(ktarray, Ncores, storage = my_storage):
output = Circulation_negative_fits(name,kt,Nbins,start)
sto.result_id = kt
sto.result = output
if yt.is_root():
for fn, vals in sorted(my_storage.items()):
kdmatrix.append(vals[0])
errmatrix.append(vals[1])
rcen = vals[2]
kdmatrix=np.array(kdmatrix)
errmatrix=np.array(errmatrix)
KT=np.zeros(Nbins)
KD=np.zeros(Nbins)
for j in range(Nbins):
optimo=np.where(errmatrix[:,j]==errmatrix[:,j].min())
KD[j]=kdmatrix[optimo,j]
KT[j]=ktarray[optimo]
KK=np.array([KT,KD])
np.save(name+'_KK',KK)
return KT,KD
def slices(ds, slcfname, slc_fields):
global aux
ds.coordinates.x_axis[1] = 0
ds.coordinates.x_axis['y'] = 0
ds.coordinates.y_axis[1] = 2
ds.coordinates.y_axis['y'] = 2
time = ds.current_time.in_units('Myr').v
c = ds.domain_center
dx = ds.domain_width / ds.domain_dimensions
slc_data = {}
slc_data['time'] = time
ya.check_aux(slc_fields)
for i, axis in enumerate(['x', 'y', 'z']):
slc = yt.SlicePlot(ds, axis, slc_fields)
ix = ds.coordinates.x_axis[axis]
iy = ds.coordinates.y_axis[axis]
res = (slc.width[1] / dx[ix], slc.width[0] / dx[iy])
slc_frb = slc.data_source.to_frb(slc.width[0], res, c, slc.width[1])
extent = np.array(slc_frb.bounds) / 1.e3
slc_data[axis] = {}
slc_data[axis + 'extent'] = extent
for f in slc_fields:
if aux[f].has_key('unit'):
slc_data[axis][f] = np.array(slc_frb[f].in_units(
aux[f]['unit']))
else:
slc_data[axis][f] = np.array(slc_frb[f])
if aux[f].has_key('factor'): slc_data[axis][f] *= aux[f]['factor']
if yt.is_root():
pickle.dump(slc_data, open(slcfname, 'wb'), pickle.HIGHEST_PROTOCOL)
def phase(sq, phfname, bin_fields):
global aux
pdfs = my_pdf(sq)
for bf in bin_fields:
n_bins = (aux[bf[0]]['n_bins'], aux[bf[1]]['n_bins'])
logs = {}
unit = {}
extrema = {}
for b in bf:
logs[b] = aux[b]['log']
if aux[b].has_key('unit'): unit[b] = aux[b]['unit']
if aux[b].has_key('limits'): extrema[b] = aux[b]['limits']
pdf = yt.create_profile(sq,
bf,
fields=fields,
n_bins=n_bins,
logs=logs,
extrema=extrema,
units=unit,
weight_field=None,
fractional=True)
pdfs.add_pdf(pdf, bf)
if yt.is_root():
pickle.dump(pdfs, open(phfname, 'wb'), pickle.HIGHEST_PROTOCOL)
def slices(ds, slcfname, slc_fields, aux={}):
ds.coordinates.x_axis[1] = 0
ds.coordinates.x_axis['y'] = 0
ds.coordinates.y_axis[1] = 2
ds.coordinates.y_axis['y'] = 2
time = ds.current_time.in_units('Myr').v
c = ds.domain_center
dx = ds.domain_width / ds.domain_dimensions
slc_data = {}
slc_data['time'] = time
for i, axis in enumerate(['x', 'y', 'z']):
slc = yt.SlicePlot(ds, axis, slc_fields)
ix = ds.coordinates.x_axis[axis]
iy = ds.coordinates.y_axis[axis]
# res=(int(slc.width[1]/dx[ix]),int(slc.width[0]/dx[iy]))
res = (ds.domain_dimensions[ix], ds.domain_dimensions[iy])
slc_frb = slc.data_source.to_frb(slc.width[0], res, c, slc.width[1])
extent = np.array(slc_frb.bounds) / 1.e3
slc_data[axis] = {}
slc_data[axis + 'extent'] = extent
for f in slc_fields:
slc_data[axis][f] = np.array(slc_frb[f])
if f in aux:
if 'unit' in aux[f]:
slc_data[axis][f] = np.array(slc_frb[f].in_units(
aux[f]['unit']))
#print(f,aux[f]['unit'])
if 'factor' in aux[f]: slc_data[axis][f] *= aux[f]['factor']
if yt.is_root():
pickle.dump(slc_data, open(slcfname, 'wb'), pickle.HIGHEST_PROTOCOL)
def main(**kwargs):
dir = kwargs['base_directory'] + kwargs['directory']
fname = glob.glob(dir + 'id0/' + kwargs['id'] + '.????.vtk')
fname.sort()
ngrids = len(glob.glob(dir + 'id*/' + kwargs['id'] + fname[-9:]))
comm = yt.communication_system.communicators[-1]
nprocs = comm.size
print ngrids, nprocs
if yt.is_root():
if not os.path.isdir(dir + 'phase/'): os.mkdir(dir + 'phase/')
for f in fname:
phfname = dir + 'phase/' + kwargs['id'] + f[-9:-4] + '.phase.p'
if os.path.isfile(phfname):
print '%s is already there' % phfname
else:
if ngrids > nprocs: ds = yt.load(f, units_override=unit_base)
else: ds = yt.load(f, units_override=unit_base, nprocs=nprocs * 8)
le = np.array(ds.domain_left_edge)
re = np.array(ds.domain_right_edge)
sq = ds.box(le, re)
pdfs = my_pdf(sq)
for bf in bin_fields:
n_bins = (aux[bf[0]]['n_bins'], aux[bf[1]]['n_bins'])
logs = {}
unit = {}
extrema = {}
for b in bf:
logs[b] = aux[b]['log']
if aux[b].has_key('unit'): unit[b] = aux[b]['unit']
if aux[b].has_key('limits'): extrema[b] = aux[b]['limits']
pdf = yt.create_profile(sq,
bf,
fields=fields,
n_bins=n_bins,
logs=logs,
extrema=extrema,
units=unit,
weight_field=None,
fractional=True)
pdfs.add_pdf(pdf, bf)
if yt.is_root():
pickle.dump(pdfs, open(phfname, 'wb'), pickle.HIGHEST_PROTOCOL)
def yt_derived_field_demo():
ds = yt.frontends.libyt.libytDataset()
slc1 = yt.SlicePlot(ds, "z", ("gamer", "level_derived_func"))
slc2 = yt.SlicePlot(ds, "z", ("gamer", "level_derived_func_with_name"))
if yt.is_root():
slc1.save()
slc2.save()
def yt_inline_inputArg( fields ):
# Get data
ds = yt.frontends.libyt.libytDataset()
# Do ProjectionPlot to fields input by gamer
sz = yt.ProjectionPlot(ds, 'z', fields, center='c')
if yt.is_root():
sz.save()
def yt_inline_ProjectionPlot( fields ):
# Load the data, just like using yt.load()
ds = yt.frontends.libyt.libytDataset()
# Do yt operation
prjz = yt.ProjectionPlot(ds, 'z', fields)
# Include this line, otherwise yt will save one copy in each rank.
if yt.is_root():
prjz.save()
def do_projection(ds, surfname):
proj = ds.proj('density', axis='z')
w = ds.domain_width[0]
res = (ds.domain_dimensions[0], ds.domain_dimensions[1])
frb = proj.to_frb(width=w, resolution=res)
surf = np.array(frb['density'].in_units('Msun/pc**2'))
if yt.is_root():
pickle.dump({
'data': surf,
'bounds': frb.bounds
}, open(surfname, 'wb'), pickle.HIGHEST_PROTOCOL)
def projection(ds, surfname):
time = ds.current_time.in_units('Myr').v
ds.coordinates.x_axis[1] = 0
ds.coordinates.x_axis['y'] = 0
ds.coordinates.y_axis[1] = 2
ds.coordinates.y_axis['y'] = 2
c = ds.domain_center
dx = ds.domain_width / ds.domain_dimensions
surf_data = {}
surf_data['time'] = time
for i, axis in enumerate(['x', 'y', 'z']):
proj = ds.proj('density', axis=axis)
ix = ds.coordinates.x_axis[axis]
iy = ds.coordinates.y_axis[axis]
res = (ds.domain_dimensions[ix], ds.domain_dimensions[iy])
frb = proj.to_frb(ds.domain_width[ix], res, c, ds.domain_width[iy])
surf = np.array(frb['density'].in_units('Msun/pc**2'))
bounds = np.array(frb.bounds)
surf_data[axis] = {'data': surf, 'bounds': bounds}
if yt.is_root():
pickle.dump(surf_data, open(surfname, 'wb'), pickle.HIGHEST_PROTOCOL)
scal_fields = ya.get_scalars(ds)
for nscal, sf in enumerate(scal_fields):
scal_data = {}
scal_data['time'] = time
for i, axis in enumerate(['x', 'y', 'z']):
proj = ds.proj(sf, axis=axis, weight_field='density')
ix = ds.coordinates.x_axis[axis]
iy = ds.coordinates.y_axis[axis]
res = (ds.domain_dimensions[ix], ds.domain_dimensions[iy])
frb = proj.to_frb(ds.domain_width[ix], res, c, ds.domain_width[iy])
scal = np.array(frb[sf])
bounds = np.array(frb.bounds)
scal_data[axis] = {'data': scal, 'bounds': bounds}
if yt.is_root():
scalfname = surfname.replace('surf.p', 'scal{}.p'.format(nscal))
pickle.dump(scal_data, open(scalfname, 'wb'),
pickle.HIGHEST_PROTOCOL)
def yt_inline_ParticlePlot():
ds = yt.frontends.libyt.libytDataset()
## ParticleProjectionPlot
#==========================
# par = yt.ParticleProjectionPlot(ds, "z")
## ParticlePlot
#==========================
par = yt.ParticlePlot(ds, "particle_position_x", "particle_position_y", "Level", center = 'c')
if yt.is_root():
par.save()
def yt_inline():
# Get data
ds = yt.frontends.libyt.libytDataset()
# Do ProjectionPlot to field Cloud0.
sz = yt.ProjectionPlot(ds, 'z', ('gamer', 'Cloud0'), center='c')
# Do ParticlePlot
par = yt.ParticlePlot(ds, 'particle_position_x', 'particle_position_y', 'particle_mass', center='c')
if yt.is_root():
sz.save()
par.save()
def yt_plot_mins():
initialize()
tseries_toplot = yt.DatasetSeries(toplot, parallel=True)
logging.getLogger().setLevel(logging.INFO)
for ds in tseries_toplot.piter():
min, loc_min = ds.h.find_min('eint')
print 'Plotting %s' % ds.basename
plotSlices(ds, zoom_fac=4, center=loc_min, drawnozzle=False,\
markcenter=True, proj_axes=['x', 'y', 'z'],\
fields=['eint', 'pressure'])
if yt.is_root():
t2 = time.time()
print 'Total time: %.2f s\n--\ninitialization: %.2f s\nploting: %.2f'\
% (t2-t0, t1-t0, t2-t1)
def Circulation_negative_optimum_mpi(name,kmin,kmax,Nk,Ncores):
ktarray=np.exp(np.linspace(np.log(kmin),np.log(kmax),Nk))
my_storage = {}
kdarray=[]
erarray=[]
for sto, kt in yt.parallel_objects(ktarray, Ncores, storage = my_storage):
output = Circulation_negative_optimum_thread(name,kt)
sto.result_id = kt
sto.result = output
if yt.is_root():
for fn, vals in sorted(my_storage.items()):
kdarray.append(vals[0])
erarray.append(vals[1])
kdarray=np.array(kdarray)
erarray=np.array(erarray)
optimo=np.where(erarray==erarray.min())
return ktarray[optimo],kdarray[optimo]
def main(**kwargs):
dir = kwargs['base_directory'] + kwargs['directory']
fname = glob.glob(dir + 'id0/' + kwargs['id'] + '.????.vtk')
fname.sort()
ngrids = len(glob.glob(dir + 'id*/' + kwargs['id'] + fname[-9:]))
comm = yt.communication_system.communicators[-1]
nprocs = comm.size
print ngrids, nprocs
if yt.is_root():
if not os.path.isdir(dir + 'surf/'): os.mkdir(dir + 'surf/')
for f in fname:
surfname = dir + 'surf/' + kwargs['id'] + f[-9:-4] + '.surf.p'
if os.path.isfile(surfname):
print '%s is already there' % surfname
else:
if ngrids > nprocs: ds = yt.load(f, units_override=unit_base)
else: ds = yt.load(f, units_override=unit_base, nprocs=nprocs * 8)
do_projection(ds, surfname)
def get_halo_sphere_particles(my_halo, par_file, radius_factor=5):
pf = yt.load(par_file)
halo_catalog = os.path.join(pf.fullpath, 'MergerHalos.out')
if 'center' in my_halo:
my_halo_data = my_halo
my_halo_data["center"] = pf.arr(my_halo['center'][0], my_halo['center'][1])
else:
my_halo_data = read_from_catalog(my_halo, halo_catalog)
rho_crit = pf.quan(rho_crit_now, "g/cm**3") * pf.hubble_constant**2 * \
(1 + pf.current_redshift)**3
if 'mass' in my_halo:
if isinstance(my_halo_data['mass'], tuple):
units = my_halo_data['mass'][1]
else:
units = "Msun"
hmass = pf.quan(my_halo_data['mass'][0], units)
else:
hmass = None
if 'radius' in my_halo:
r_200 = pf.quan(my_halo['radius'][0], my_halo['radius'][1])
else:
r_200 = (((3. * hmass) /
(4. * np.pi * rho_crit * 200.))**(1./3.)).to("Mpc")
if yt.is_root():
print ("Reading particles for a sphere surrounding halo %d." % my_halo['id'])
print ("Halo %d, pos: %f, %f, %f, mass: %s, r_200: %s." % \
(my_halo_data['id'], my_halo_data['center'][0], my_halo_data['center'][1],
my_halo_data['center'][2], hmass, r_200))
my_sphere = pf.h.sphere(my_halo_data['center'], radius_factor * r_200)
return (my_halo_data['center'],
my_sphere['particle_index'],
my_sphere['particle_mass'].in_units('Msun'),
np.array([my_sphere['particle_position_x'],
my_sphere['particle_position_y'],
my_sphere['particle_position_z']]))
def run():
input_fn, output_fn, selection, group = sys.argv[1:5]
njobs = int(sys.argv[5])
dynamic = bool(int(sys.argv[6]))
a = ytree.load(input_fn)
if "test_field" not in a.field_list:
a.add_analysis_field("test_field", default=-1, units="Msun")
trees = list(a[:8])
for tree in trees:
for node in ytree.parallel_tree_nodes(tree,
group=group,
njobs=njobs,
dynamic=dynamic):
root = node.root
yt.mylog.info(
f"Doing {node.tree_id}/{root.tree_size} of {root._arbor_index}"
)
node["test_field"] = 2 * node["mass"]
if yt.is_root():
a.save_arbor(filename=output_fn, trees=trees)
def get_halo_indices(my_halo, dataset, method='sphere', radius_factor=5.0):
shifted = np.zeros(3, dtype=bool)
if method == 'halo':
halo_indices, particle_masses, particle_positions = \
get_halo_particles(my_halo, dataset)
elif method == 'sphere':
halo_com, halo_indices, particle_masses, particle_positions = \
get_halo_sphere_particles(my_halo, dataset, radius_factor=radius_factor)
unitary_1 = halo_com.to("unitary").uq
for i, axis in enumerate(axes):
if particle_positions[i].max() - particle_positions[i].min() > 0.5:
if yt.is_root():
print ("Halo periodic in %s." % axis)
particle_positions[i] -= 0.5
particle_positions[i][particle_positions[i] < 0.0] += 1.0
halo_com[i] -= 0.5 * unitary_1
if halo_com[i] < 0.0:
halo_com[i] += 1.0 * unitary_1
shifted[i] = True
if method == 'halo':
halo_com = (particle_positions * particle_masses).sum(axis=1) / particle_masses.sum()
return (halo_indices, halo_com, shifted)
def plot_a_vr(path,
header,
cycle,
use_ghost_zones=True,
annotate=True,
rotate=False,
zoom=False):
"""
volume rendering plot
path [string] : the path of data files
header [string]: the header of the data file
cycle [int] : the data cycle
ex. for file: /data/ccsn3d_hdf5_plt_cnt_0100
path = "/data"
header = "ccsn3d"
cycle = 100
"""
fn = get_fn(path, header, cycle)
ds = yt.load(fn)
# register new units for entropy
ds.unit_registry.add('kB',
1.3806488e-16,
dimensions=energy / temperature,
tex_repr='k_{B}')
ds.unit_registry.add('by', 1.674e-24, dimensions=mass, tex_repr='baryon')
# create a new derived field: Entropy
def _entr(field, data):
"""
fixed the entropy units in flash
"""
entr = data["entr"]
kb_by = yt.YTQuantity(1, 'erg') / yt.YTQuantity(
1, 'K') / yt.YTQuantity(1, 'g') * (1.3806488e-16 / 1.674e-24)
return entr * kb_by
ds.add_field("Entr",
function=_entr,
units="kB/by",
display_name="Entropy",
dimensions=energy / temperature / mass)
# create a new derived field: Entropy
def _entrdens(field, data):
"""
create a new derived field to show both entropy and density
if density > PNS_DENSITY:
entropy = PNS_ENTR
else:
entropy = entropy
"""
dens = data["dens"]
entr = data["entr"]
entrdens = entr * (
np.exp(-(dens.in_cgs() / PNS_DENSITY)**5)) + PNS_ENTR
kb_by = yt.YTQuantity(1, 'erg') / yt.YTQuantity(
1, 'K') / yt.YTQuantity(1, 'g') * (1.3806488e-16 / 1.674e-24)
return entrdens * kb_by
ds.add_field("Entropy",
function=_entrdens,
units="kB/by",
display_name="Entropy",
dimensions=energy / temperature / mass)
#debug: also plot a entropy slice
if yt.is_root():
pc = yt.SlicePlot(ds, 'z', 'Entr')
pc.zoom(40)
pc.set_log('Entr', False)
pc.save('fig_slice_z_' + header + '_' + string.zfill(cycle, 4) +
'.png')
# get the entropy range from time
time = ds.current_time.in_cgs().v - TSHIFT
if CUSTOM_EMAX:
emin, emax = get_emin_emax(time)
else:
entropy_max, pe = ds.find_max('Entropy')
emax = entropy_max.v - 1.0
emin = emax - 3.0
if yt.is_root():
print "emin/emax:", emin, emax
# check if emax > emin
if emax <= emin:
print "Error: emax <= emin"
quit()
# only render the region with r < 1.e8 cm
# this is necessary if we want to include ghost zones
sphere = ds.sphere([0, 0, 0], (1.e8, 'cm'))
sc = yt.create_scene(sphere, field='Entropy')
# set the camera resolution and width
sc.camera.north_vector = [0, 0, 1]
sc.camera.resolution = (VR_RESOLUTION, VR_RESOLUTION)
sc.camera.set_width(ds.quan(VR_WIDTH, 'cm'))
#sc.camera.set_width(ds.quan(2.5e7,'cm'))
# define the tranfer function for high entropy region
def linearFunc(values, minv, maxv):
try:
return ((values) - values.min()) / (values.max() - values.min())
#return (na.sqrt(values) - values.min())/(values.max()-values.min())
except:
return 1.0
source = sc.get_source(0)
source.set_use_ghost_zones(use_ghost_zones) # *** SLOW ***
source.set_log(False)
source.grey_opacity = True
# create a transfer function helper
tfh = TransferFunctionHelper(ds)
tfh.set_field('Entropy')
tfh.set_bounds()
tfh.set_log(False)
tfh.build_transfer_function()
# add a thin layer to show the shock front
esh, ew = get_shock_entropy(time, emin, emax)
tfh.tf.add_layers(1,
w=(ew),
mi=(esh),
ma=(esh + 0.5),
col_bounds=[4.0, (esh + 1.0)],
alpha=10.0 * esh * np.ones(1),
colormap='cool_r')
# plot the PNS at entr = PNS_ENTR
tfh.tf.add_layers(1,
w=0.5,
mi=0.49,
ma=0.55,
col_bounds=[0.05, 0.55],
alpha=100.0 * np.ones(1),
colormap='Purples')
# map the high entropy region: version 3
tfh.tf.map_to_colormap(emin,
emax,
scale=100.0,
scale_func=linearFunc,
colormap='autumn')
# version 1: use many layers
#tfh.tf.add_layers(12,w=0.05,mi=emin,ma=emax,col_bounds=[emin,emax],
# alpha=70*np.linspace(0.0,1.5,10),colormap='hot')
tfh.tf.grey_opacity = True
source.set_transfer_function(tfh.tf)
source.grey_opacity = True
# plot the transfer function
#source.tfh.plot('fig_transfer_function_entr.png', profile_field='cell_mass')
# plot volume rendering plot without annotation
if not annotate:
sc.save('fig_vr_' + header + '_' + string.zfill(cycle, 4) + '.png',
sigma_clip=4)
else:
# with annotation
sc.annotate_axes(alpha=0.8)
#sc.annotate_scale() #
# line annotation
#annotate_width = 1.0e7 # [cm]
#annotate_axis = 1.5e7 # [cm]
#annotate_at = (annotate_axis/2.0 - annotate_width/(2.0*np.sqrt(2.0)))
#colors = np.array([[0.45,0.5,0.5,0.7]]) # white
#vertices = np.array([[[annotate_at,0.,0.],[0.,annotate_at,0.]]])*annotate_width*cm
#lines = LineSource(vertices,colors)
#sc.add_source(lines)
#sc.annotate_domain(ds,color=[1,1,1,0.01])
#text_string= "Time = %.1f (ms)" % (float(ds.current_time.to('s')*1.e3))
text_string = "Time = %.1f (ms)" % (float(time * 1.e3))
sc.save_annotated("fig_vr_" + header + "_annotated_" +
string.zfill(cycle, 4) + '.png',
sigma_clip=4.0,
label_fmt="%.1f",
text_annotate=[[(0.05, 0.95), text_string,
dict(color="w",
fontsize="20",
horizontalalignment="left")]])
# rotate the camera
if DO_ROT:
rotate = True
if rotate:
fstart = 1 # modify here for restart
frames = ROT_FRAMES # total number of frames for rotation
cam = sc.camera
i = 0
for _ in cam.iter_rotate(2.0 * np.pi, frames):
i += 1
if i >= fstart:
sc.render()
if not annotate:
sc.save("fig_vr_" + header + "_" + string.zfill(cycle, 4) +
'_rot_' + string.zfill(i, 4) + '.png',
sigma_clip=2)
else:
sc.save_annotated("fig_vr_" + header + "_annotated_" +
string.zfill(cycle, 4) + '_rot_' +
string.zfill(i, 4) + '.png',
sigma_clip=4.0,
text_annotate=[[
(0.05, 0.95), text_string,
dict(color="w",
fontsize="20",
horizontalalignment="left")
]])
# TODO: zoom in or zoom out
if DO_ZOOM:
zoom = True
if zoom:
fstart = 0 # modify here for restart
frames = ZOOM_FRAMES
cam = sc.camera
i = 0
for _ in cam.iter_zoom(ZOOM_FACT, ZOOM_FRAMES):
i += 1
if i >= fstart:
sc.render()
if not annotate:
sc.save("fig_vr_" + header + "_" + string.zfill(cycle, 4) +
'_zoom_' + string.zfill(i, 4) + '.png',
sigma_clip=2)
else:
sc.save_annotated("fig_vr_" + header + "_annotated_" +
string.zfill(cycle, 4) + '_zoom_' +
string.zfill(i, 4) + '.png',
sigma_clip=4.0,
text_annotate=[[
(0.05, 0.95), text_string,
dict(color="w",
fontsize="20",
horizontalalignment="left")
]])
quit()
return
def gen_profs(ds,fils,dsname,keep=True,totalratio=None,mask=False):
x = np.zeros(10)
storage = {}
#Check to see if totalratio has been set, else determine it.
if totalratio == None:
#Determine simulation-wide ratio of baryon to Total matter, done in parallel to speed up
for stor,i in yt.parallel_objects(range(1),storage = storage):
denstotal = ds.all_data().quantities.weighted_average_quantity("density","cell_volume")
dmtotal = ds.all_data().quantities.weighted_average_quantity("dark_matter_density","cell_volume")
totalratio = denstotal / (dmtotal + denstotal)
del dmtotal
stor.result = totalratio,denstotal.in_units('g/cm**3')
stor.result_id = 1
totalratio,denstotal = storage[1]
print totalratio # print to stdout, allows us to check ratio makes sense/debug
del storage
#Print to a file, allows checking of progress when in job queue
if yt.is_root():
with open("Testinfo.dat",'w') as f:
f.write('Ratio determined')
print("All Ratio determined")
x = []
storage = {}
#Gather list of all profiles on disk
filelist = sorted(os.listdir(''.join(['/shome/mackie/data',dsname,'/profiles'])))
#Include only density profiles
filelist = filelist[:len(filelist)/2]
#Load a numpy mask array if required
if keep == True:
keeplist = np.load(''.join(['/shome/mackie/data',dsname,'/filkeep.npy']))
for stor,file_in_dir in yt.parallel_objects(filelist,storage=storage):
filnum = int(file_in_dir[7:10])
segnum = int(file_in_dir[13:16])
prof = yt.load(''.join(['/shome/mackie/data',dsname,'/profiles/',file_in_dir]))
dm = prof.data['dark_matter_density']
dens = prof.data['density'].in_units('g/cm**3')
result = (dens /( dm + dens)) - totalratio
dens = dens/denstotal
result = result.v
if keep == True:
if keeplist[filnum][segnum] == True:
stor.result = result,dens
stor.result_id = file_in_dir
else:
stor.result = result,dens
stor.result_id = file_in_dir
del prof,dm,dens,filnum,segnum
results = []
denresult = []
for keys,values in storage.items():
results.append(values[0])
denresult.append(values[1])
results = np.array(results)
denresult = np.array(denresult)
if mask:
results = np.ma.masked_array(results,mask=~mask,fill_value=np.nan)
#Get x bins
x = yt.load(''.join(['/shome/mackie/data',dsname,'/profiles/densfil000seg000.h5'])).data['x']
if yt.is_root():
with open("Testinfo.dat",'a') as f:
f.write("returning results to plot")
print results
return results,denresult,x
def yt_inline_ProfilePlot():
ds = yt.frontends.libyt.libytDataset()
profile = yt.ProfilePlot(ds, "x", ["density"])
if yt.is_root():
profile.save()
if scale == scale_in_file:
halo_file = this_file
break
return halo_file
if __name__ == "__main__":
args = parse()
import yt
from yt.analysis_modules.sunrise_export import sunrise_octree_exporter
from yt.analysis_modules.halo_finding.halo_objects import RockstarHaloList
yt.enable_parallelism()
if yt.is_root():
print '/nStarting '+ sys.argv[0]
print 'Parsed arguments: '
print args
print
# Get parsed values
sim_dirs, snap_base = args['sim_dirs'], args['snap_base']
print 'Analyzing ', sim_dirs
out_dir = args['out_dir']
modify_outdir = 0
if 'sim_dir' in out_dir:
out_dir = out_dir.replace('sim_dir','')
modify_outdir = 1
def plot_vel(filament,ds,dataset,fil=-1,maskarray=False):
#Routine to plot velocity of particles along a filament
#Set gravitational constant
G = YTQuantity(6.67408E-11,'m**3/(kg * s**2)')
#Gather velocity and density values from disk, done in parallel to speed computation
#We first gather a list of the profiles on the disk, then reshape this into a list of [density,other] profiles
#This is then iterated over in parallel to load the correct data
filelist = sorted(os.listdir(''.join(['/shome/mackie/data/',dataset,'/profiles'])))
profnumbers = len(filelist)/2
files = [ [filelist[i],filelist[i+profnumbers]] for i in range(profnumbers)]
del filelist,profnumbers
storage = {}
for stor,file_in_dir in ytpar.parallel_objects(files,storage=storage):
#Determines correct index to give to segment profiles
filnum = int(file_in_dir[0][7:10])
segnum = int(file_in_dir[0][13:16])
#Calc total density for each segment
densprof = yt.load(''.join(['/shome/mackie/data/',dataset,'/profiles/',file_in_dir[0]]))
dm = densprof.data['dark_matter_density'].in_units('g/cm**3')
dens = densprof.data['density'].in_units('g/cm**3')
totaldens = dm + dens
del densprof,dm,dens
#Get velocity profiles
velprof = yt.load(''.join(['/shome/mackie/data/',dataset,'/profiles/',file_in_dir[1]]))
vel = velprof.data['cylindrical_radial_velocity'].in_units('km/s')
stor.result = (vel,totaldens)
stor.result_id = (filnum,segnum)
#Restruct dict into np array of correct structure.
vel_profs = [ [] for i in range(len(filament))]
densprofs = [ [] for i in range(len(filament))]
x = yt.load(''.join(['/shome/mackie/data/',dataset,'/profiles/',file_in_dir[0]])).data['x']
xarr = [[] for i in range(len(filament))]
for key,values in sorted(storage.items()):
filnum,segnum = key
vel,dens = values
xarr[filnum].append(x.in_units('Mpc'))
vel_profs[filnum].append(vel.in_units('km/s'))
densprofs[filnum].append(dens)
for i in range(len(xarr)):
xarr[i] = YTArray(np.array(xarr[i]),'Mpc')
vel_profs[i] = YTArray(np.array(vel_profs[i]),'km/s')
vel_profs = YTArray(np.array(vel_profs),'km/s')
xarr = YTArray(np.array(xarr),'Mpc')
del storage
#Turn into np arrays for QoL
densprofs = np.array(densprofs)
#Gather x bins from disk
#Determine masses and thus escape velocities
mass = [get_masses(densprofs[i],x,filament[i],ds,accumulate=True) for i in range(len(filament))]
mass = np.array(mass)
mass = YTArray(mass,'g')
print mass[1][1]
del densprofs
print xarr[1][1]
vel_ratio = ( ( (2*G* mass) / xarr) ** (1.0/2.0))
vel_ratio = vel_ratio.in_units('km/s')
if yt.is_root():
print mass[1][1]
print xarr[1][1]
print vel_ratio[1][1]
#vel_ratio is **approx** escape vel, used to ratio later
#Generate ratio of velocity to escape velocity
vel_profs = (vel_profs.in_units('km/s')/vel_ratio.in_units('km/s'))
del vel_ratio
if fil > -1:
length_plot = plot_vel_fil(vel_profs[fil].v,gen_dists(filament[fil],ds),x)
else:
length_plot = None
if maskarray:
print"Masking Data"
vel_profs = np.ma.masked(vel_profs, mask = ~maskarray,fill_value=np.nan)
#Flatten vel profs, ought to bemore elegant solution
vel_prof_flatten = []
for fil in vel_profs:
for seg in fil:
vel_prof_flatten.append(seg)
vel_profs = np.array(vel_prof_flatten)
del vel_prof_flatten
plot = probmap.prob_heat_map(vel_profs,'radial_velocity',x=x)
return plot,length_plot
imin, imax = None, None
if len(sys.argv) > 1:
imin = int(sys.argv[1])
imax = int(sys.argv[2])
if len(sys.argv) > 3:
if str(sys.argv[3]) != "&":
weight_field = str(sys.argv[3])
if imin is None:
dsfiles = [np.sort(glob.glob("DD????/DD????"))[-1]]
else:
dsfiles = [
"DD%0004i/DD%0004i" % (i, i) for i in np.arange(imin, imax + 1, 1)
]
if yt.is_root():
print(dsfiles)
fields = [('gas', 'number_density'), 'Temperature', ('gas', 'photo_gamma'),
('gas', 'metallicity'), ('gas', 'C_over_N'), ('gas', 'C_over_H'),
('gas', 'O_over_H'), ('gas', 'Fe_over_H')]
fields = [('gas', 'number_density'), 'Temperature', ('gas', 'O_over_H'),
('gas', 'photo_gamma')]
#
# fields = [('gas','number_density')]
#
def format_plot(ds, sp, fields):
def prob_heat_map(profiles,var,x=np.empty(10)):
#routine to generate a probability heat map for a variable(e.g. density/temperature) in a filament. Will also plot mean and median to this map.
'''
if 'densdiff' not in var:
#set up x values, and create empty arrays
x = profiles[0].x
data = np.empty([len(profiles),x.shape[0]])
mask = np.ones([len(profiles),x.shape[0]],dtype=bool)
#populate data, and generate a mask for 0 values
for index in xrange(len(profiles)):
data[index,] = profiles[index][(var)].v
else:
'''
data = np.array(profiles)
if not("velocity" in var or "densdiff" in var or "densrat" in var):
data = np.ma.masked_values(data,0.0,atol = 1e-40)
data=np.ma.masked_invalid(data)
#determine min, max and the binsize for variable
if yt.is_root(): print(np.ravel(data))
mind = np.nanmin(data.reshape(-1))
maxd = np.nanmax(data.reshape(-1))
print 'max'
print maxd
#if mind < 0:
# mind = 0 - np.amax([abs(mind),maxd])
# maxd = 0 - mind
bins = 50
length = maxd - mind
#generate bins for the variable
if var == "cylindrical_radial_velocity" or var == "densdiff" or var == "densrat":
y = np.linspace(mind,maxd,bins)
else:
print mind
if mind == 0:
mind = data.min()
print mind
y = np.logspace(np.log10(mind),np.log10(maxd),bins)
#set up prob array
prob = np.zeros([10,y.size])
#loop over all values of variable and add to relevant bin in prob array
for i in range(len(data)):
for j in range(10):
z = data[i][j]
if math.isnan(z):
break
yindex = 0
#calculate correct index to increment
for index in range(y.size):
if z <= y[index]:
yindex = index
prob[j,yindex] += 1
break
mean = np.empty(x.size)
median = np.empty(x.size)
#normalize probability array, calculate mean and median
for index in range(x.size):
yvals = np.array(prob[index,:])
summed = np.sum(yvals)
if summed > 0:
yvals = yvals / summed
prob[index,:] = yvals
del yvals,summed
for index in range(x.size):
maskeddata = data[:,index]
mean[index] = np.nanmean(maskeddata)
median[index] = np.nanmedian(maskeddata)
del maskeddata
#now plot
plot = plt.figure()
plt.plot(x,mean,'g',label='Mean')
plt.plot(x,median,'r',label='Median')
CS = plt.pcolormesh(x,y,prob.T,cmap='seismic')
plt.xlabel("Radius (MPC)")
plt.xscale('log')
#plt.yscale('log')
if var.lower() == 'density':
plt.ylabel("Density (g/cm^3)")
plt.yscale('log')
plt.xlim(-0.1,2.5)
elif var.lower() == 'temperature':
plt.ylabel("Temperature (K)")
plt.fill_between(x, 10E4 , 10E6 ,color='grey',alpha='0.2')
plt.yscale('log')
plt.xlim(-0.1,2.5)
elif var.lower() == 'cylindrical_radial_velocity':
plt.ylabel("Velocity/Escape Velocity")
plt.xscale('linear')
plt.yscale('linear')
elif var.lower() == "densrat":
plt.ylabel("Local Density/Average Density")
plt.yscale('linear')
plt.xlim(-0.1,2.5)
else:
plt.ylabel("Baryon Density/ Total Density - Avg. Baryon Fraction")
plt.yscale('linear')
plt.xlim(-0.1,2.5)
cbar = plt.colorbar(CS)
cbar.ax.set_ylabel("Probability")
plt.legend()
return plot
def main(**kwargs):
global aux
#yt.funcs.mylog.setLevel(50)
dir = kwargs['base_directory'] + kwargs['directory']
fname = glob.glob(dir + 'id0/' + kwargs['id'] + '.????.vtk')
fname.sort()
if yt.is_root(): print fname[0]
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:]))
if kwargs['parallel']:
from mpi4py import MPI
comm = MPI.COMM_WORLD
nprocs = comm.size
rank = comm.rank
else:
nprocs = 1
rank = 0
print ngrids, rank, nprocs, yt.is_root()
do_phase = kwargs['phase']
if ngrids > nprocs: ds = yt.load(fname[0], units_override=ya.unit_base)
else: ds = yt.load(fname[0], units_override=ya.unit_base, nprocs=nprocs)
mhd = ('athena', 'cell_centered_B_x') in ds.field_list
cooling = ('athena', 'pressure') in ds.field_list
rotation = kwargs['rotation'] != 0.
if rotation:
ya.Omega = ya.YTQuantity(kwargs['rotation'], 'km/s/kpc')
if kwargs['rotation'] == 280:
aux = ya.set_aux('starburst')
if rank == 0:
print "phase plot:", do_phase
print "MHD:", mhd
print "cooling:", cooling
print "rotation:", rotation, ya.Omega
bin_fields = []
bin_fields.append(['nH', 'pok'])
bin_fields.append(['nH', 'temperature'])
bin_fields.append(['velocity_z', 'temperature'])
if rotation:
bin_fields.append(['dvelocity_magnitude', 'temperature'])
else:
bin_fields.append(['velocity_magnitude', 'temperature'])
if mhd:
bin_fields.append(['nH', 'mag_pok'])
bin_fields.append(['nH', 'ram_pok_z'])
bin_fields.append(['nH', 'plasma_beta'])
slc_fields = ['nH', 'pok', 'temperature', 'velocity_z', 'ram_pok_z']
fields_to_draw = ['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')
if rank == 0:
if not os.path.isdir(dir + 'slice/'): os.mkdir(dir + 'slice/')
if not os.path.isdir(dir + 'surf/'): os.mkdir(dir + 'surf/')
if not os.path.isdir(dir + 'phase/'): os.mkdir(dir + 'phase/')
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'
if do_phase:
phfname = dir + 'phase/' + kwargs['id'] + f[-9:-4] + '.phase.p'
else:
phfname = f
if rank == 0:
print i, compare_files(f, slcfname), compare_files(
f, surfname), compare_files(f, phfname)
if compare_files(f,surfname) and \
compare_files(f,phfname) and \
compare_files(f,slcfname):
if rank == 0: print 'all data is already there'
else:
if ngrids > nprocs: ds = yt.load(f, units_override=ya.unit_base)
else: ds = yt.load(f, units_override=ya.unit_base, nprocs=nprocs)
ya.add_yt_fields(ds, mhd=mhd, rotation=rotation, cooling=cooling)
if not compare_files(f, surfname):
if rank == 0: print 'projectiong...'
projection(ds, surfname)
if not compare_files(f, slcfname):
if rank == 0: print 'slicing...'
slices(ds, slcfname, slc_fields)
if not compare_files(f, phfname) and do_phase:
if rank == 0: print 'binning...'
le = np.array(ds.domain_left_edge)
re = np.array(ds.domain_right_edge)
sq = ds.box(le, re)
phase(sq, phfname, bin_fields)
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'
if do_phase:
phfname = dir + 'phase/' + kwargs['id'] + f[-9:-4] + '.phase.p'
else:
phfname = f
if i % nprocs == rank:
if not compare_files(surfname, surfname + 'ng'):
print 'drawing for %s on %d' % (surfname, rank)
plot_projection(surfname, f)
if not compare_files(slcfname, slcfname + 'ng'):
print 'drawing for %s on %d' % (slcfname, rank)
plot_slice(slcfname, f, fields_to_draw)
if not compare_files(phfname, phfname + 'ng') and do_phase:
if do_phase:
print 'drawing for %s on %d' % (phfname, rank)
plot_phase(phfname, bin_fields)
def calculate_mass_in_sphere(dd):
data_sub_dir = dd.name
a = dd.a
r_plus = M*(1 + math.sqrt(1 - a**2))
min_radius = 0.0*r_plus/4
start_time = time.time()
# load dataset time series
if (data_Eulerian_rho and not use_Eulerian_rho):
# derived fields
@derived_field(name = "rho_E_eff", units = "")
def _rho_E_eff(field, data, force_override=True):
r_BL = (data["spherical_radius"]/cm)*(1 + r_plus*cm/(4*data["spherical_radius"]))**2
Sigma2 = r_BL**2 + (data["z"]*a*M/(r_BL*cm))**2
Delta = r_BL**2 + (a*M)**2 - 2*M*r_BL
A = (r_BL**2 + (a*M)**2)**2 - ((a*M)**2)*Delta*(data["x"]**2 + data["y"]**2)/(cm*r_BL)**2
alpha = pow(Delta*Sigma2/A, 0.5)
beta = -2*a*(M**2)*r_BL/A
return (data["rho"]*alpha - beta*data["S_azimuth"])*pow(data["chi"],-3/2)
elif ((data_Eulerian_rho and use_Eulerian_rho) or ((not data_Eulerian_rho) and (not use_Eulerian_rho))):
# derived fields
@derived_field(name = "rho_E_eff", units = "")
def _rho_E_eff(field, data):
return data["rho"]
dataset_path = data_root_path + "/" + data_sub_dir + "/KerrSFp_*.3d.hdf5"
ds = yt.load(dataset_path) # this loads a dataset time series
print("loaded data from ", dataset_path)
print("time = ", time.time() - start_time)
N = len(ds)
ds0 = ds[0] # get the first dataset
# set centre
center = [512.0, 512.0, 0]
L = 512.0
data_storage = {}
# iterate through datasets (forcing each to go to a different processor)
for sto, dsi in ds.piter(storage=data_storage):
time_0 = time.time()
# store time
current_time = dsi.current_time
output = [current_time]
# make sphere
sphere = dsi.sphere(center, max_radius) #- dsi.sphere(center, min_radius)
volume = sphere.sum("cell_volume")
if half_box:
volume = 2*volume
# calculate energy inside sphere
meanE = sphere.mean("rho_E_eff", weight="cell_volume")
E = volume*meanE
output.append(E)
# store output
sto.result = output
sto.result_id = str(dsi)
dt = 1.25
i = int(current_time/dt)
print("done {:d} of {:d} in {:.1f} s".format(i+1, N, time.time()-time_0), flush=True)
if yt.is_root():
# make data directory if it does not already exist
makedirs(home_path + output_dir, exist_ok=True)
# output to file
if use_Eulerian_rho:
#dd.filename = "l={:d}_m={:d}_a={:s}_mu={:s}_Al={:s}_mass_in_r={:d}_Eulerian_rho.csv".format(dd.l, dd.m, str(dd.a), dd.mu, dd.Al, max_radius)
dd.filename = "{:s}_mass_in_r={:d}_Eulerian_rho.csv".format(dd.name, max_radius)
else:
#dd.filename = "l={:d}_m={:d}_a={:s}_mu={:s}_Al={:s}_mass_in_r={:d}_conserved_rho.csv".format(dd.l, dd.m, str(dd.a), dd.mu, dd.Al, max_radius)
dd.filename = "{:s}_mass_in_r={:d}_conserved_rho.csv".format(dd.name, max_radius)
output_path = home_path + output_dir + "/" + dd.filename
# output header to file
f = open(output_path, "w+")
f.write("# t mass in r<=" + str(max_radius) + " #\n")
# output data
for key in sorted(data_storage.keys()):
data = data_storage[key]
f.write("{:.3f} ".format(data[0]))
f.write("{:.2f}\n".format(data[1]))
f.close()
print("saved data to file " + str(output_path))
fp = open("%s/parameter_file.txt" % (ic_dir), "a")
fp.write("\n"
"#\n"
"# must-refine particle parameters\n"
"# *** must also include method 8 in CellFlaggingMethod ***\n"
"# *** do NOT include the RefineRegion parameters above ***\n"
"#\n"
"MustRefineParticlesCreateParticles = 3\n"
"MustRefineParticlesRefineToLevel = %d\n"
"CosmologySimulationParticleTypeName = RefinementMask\n" \
% (params["level"]))
fp.close()
# Copy initial conditions directory to the simulation run directory
print("Moving initial conditions to %s" % (params["sim_dir"]))
os.rename(ic_dir, params["sim_dir"])
return
if __name__ == "__main__":
params = {}
if yt.is_root():
params = startup()
if parallel:
params = comm.bcast(params)
params = get_previous_run_params(params)
params = find_lagrangian_region(params)
if yt.is_root():
run_music(params)
def calculate_flux(dd, old_Jr):
data_sub_dir = dd.name
a = dd.a
r_plus = M * (1 + math.sqrt(1 - a**2))
r_minus = M * (1 - math.sqrt(1 - a**2))
min_R = r_plus / 4
start_time = time.time()
# load dataset time series
dataset_path = data_root_path + "/" + data_sub_dir + "/KerrSFp_*.3d.hdf5"
ds = yt.load(dataset_path) # this loads a dataset time series
print("loaded data from ", dataset_path)
print("time = ", time.time() - start_time)
N = len(ds)
# set centre
center = [512.0, 512.0, 0]
L = 512.0
if old_Jr:
# derived fields
@derived_field(name="rho_Jr_eff", units="")
def _rho_Jr_eff(field, data):
r_BL = (data["spherical_radius"] /
cm) * (1 + r_plus * cm / (4 * data["spherical_radius"]))**2
Sigma2 = r_BL**2 + (data["z"] * a * M / (r_BL * cm))
#Delta = r_BL**2 + (a*M)**2 - 2*M*r_BL
return ((data["spherical_radius"]**2 / cm**2) * (r_BL - r_minus) /
(Sigma2 * r_BL)) * data["S_r"] * pow(data["chi"], -3)
elif not old_Jr:
# derived fields
@derived_field(name="rho_Jr_eff", units="")
def _rho_Jr_eff(field, data):
return data["J_r"]
data_storage = {}
# iterate through datasets (forcing each to go to a different processor)
for sto, dsi in ds.piter(storage=data_storage):
time_0 = time.time()
# store time
current_time = dsi.current_time
output = [current_time]
# make inner and outer shells
outer_shell = dsi.sphere(center,
max_R + 0.5 * outer_thickness) - dsi.sphere(
center, max_R - 0.5 * outer_thickness)
inner_shell = dsi.sphere(center, min_R + inner_thickness) - dsi.sphere(
center, min_R)
# calculate inner and outer flux
Jr_outer = outer_shell.mean(
"rho_Jr_eff", weight="cell_volume") * 4 * math.pi * (max_R**2)
Jr_inner = inner_shell.mean("rho_Jr_eff",
weight="cell_volume") * 4 * math.pi * (
(r_plus / 4)**2)
output.append(Jr_outer)
output.append(Jr_inner)
# store output
sto.result = output
sto.result_id = str(dsi)
dt = 1.25
i = int(current_time / dt)
print("done {:d} of {:d} in {:.1f} s".format(i + 1, N,
time.time() - time_0),
flush=True)
if yt.is_root():
# make data directory if it does not already exist
makedirs(home_path + output_dir, exist_ok=True)
# output to file
if old_Jr:
dd.filename = "l={:d}_m={:d}_a={:s}_phase={:s}_flux.csv".format(
dd.l, dd.m, str(dd.a), dd.phase)
elif not old_Jr:
dd.filename = "l={:d}_m={:d}_a={:s}_phase={:s}_flux.csv".format(
dd.l, dd.m, str(dd.a), dd.phase)
output_path = home_path + output_dir + "/" + dd.filename
# output header to file
f = open(output_path, "w+")
f.write("# t flux r={:.0f} flux r={:.2f} \n".format(max_R, min_R))
# output data
for key in sorted(data_storage.keys()):
data = data_storage[key]
f.write("{:.3f} ".format(data[0]))
f.write("{:.3f} ".format(data[1]))
f.write("{:.3f}\n".format(data[2]))
f.close()
print("saved data to file " + str(output_path))
