def make_base_plot(ds, field, width, depth, cm):
width = parse_lengths(ds, width)
depth = parse_lengths(ds, depth)
normal, north, center = plot_basics(ds)
# Make a smaller data object to hopefully speed up computation. We make
# it bigger than the projection region to ensure nothing is missed.
ds_select = ds.disk(center=center,
normal=normal,
radius=2 * width,
height=2 * depth)
plot = yt.OffAxisProjectionPlot(ds,
fields=field,
center=center,
normal=normal,
north_vector=north,
width=width,
depth=depth,
data_source=ds_select)
plot.set_cmap(field, cm)
plot.set_zlim(field, 0.1, 10**3)
plot.annotate_timestamp(corner='upper_left',
redshift=True,
draw_inset_box=True)
plot.set_unit(field, "msun/pc**2")
return plot
def make_off_axis_projection(ds, vec, north_vec, ion_fields, center, width, data_source, radius, weight_field=None, dir=None):
"""
Use OffAxisProjectionPlot to make projection (cannot specify resolution)
"""
p = yt.OffAxisProjectionPlot(ds, vec, ion_fields, center=center, width=width,
data_source=data_source, north_vector=north_vec, weight_field=weight_field)
p.hide_axes()
p.annotate_scale()
p.annotate_timestamp(redshift=True)
r = radius.in_units('kpc')
p.annotate_sphere(center, (r, 'kpc'), circle_args={'color':'white', 'alpha':0.5, 'linestyle':'dashed', 'linewidth':5})
for field in ion_fields:
p.set_cmap(field, 'dusk')
set_image_details(p, field, True)
p.set_background_color(field)
if dir is None:
dir = 'face/'
p.save(os.path.join('images', dir))
return p.frb
def main():
file = sys.argv[1]
part_file = file[:-12] + 'part' + file[-5:]
ds = yt.load(file, particle_filename=part_file)
L = [0.0, 0.0, 1.0]
myf.set_normal(L)
center_pos = np.array([0.0, 0.0, 0.0])
args = parse_inputs()
field = args.field
dd = ds.all_data()
center_vel = dd['Center_Velocity']
center_pos = dd['Center_Position']
part_pos = dd['All_Particle_Positions']
part_mass = dd['All_Particle_Masses']
print("center vel =", myf.get_center_vel())
print("center pos for field calculations", myf.get_center_pos())
print("particle positions are:", myf.get_part_pos())
print("particle masses are:", myf.get_part_mass())
center_pos = np.array([0.0, 0.0, 0.0])
proj = yt.OffAxisProjectionPlot(
ds,
L, [field, 'cell_mass', 'velx_mw', 'vely_mw', 'magx_mw', 'magy_mw'],
center=(center_pos, 'AU'),
width=(args.ax_lim, 'AU'),
depth=(args.slice_thickness, 'AU'))
image = (proj.frb.data[simfo['field']] / thickness.in_units('cm')).value
velx_full = (
proj.frb.data[('gas', 'Projected_Velocity_mw')].in_units('g*cm**2/s') /
thickness.in_units('cm')).value
vely_full = (proj.frb.data[('gas', 'velz_mw')].in_units('g*cm**2/s') /
thickness.in_units('cm')).value
magx = (proj.frb.data[
('gas', 'Projected_Magnetic_Field_mw')].in_units('g*gauss*cm') /
thickness.in_units('cm')).value
magy = (proj.frb.data[('gas', 'magz_mw')].in_units('g*gauss*cm') /
thickness.in_units('cm')).value
mass = (proj.frb.data[('gas', 'cell_mass')].in_units('cm*g') /
thickness.in_units('cm')).value
proj.save()
def xray_3color(proj_angle=30.):
dir = '/home/ychen/d9/FLASH4/stampede/0529_L45_M10_b1_h1/'
fname = dir + 'MHD_Jet_hdf5_plt_cnt_0620'
#fname = '~/d9/yt_testdata/GasSloshing/sloshing_nomag2_hdf5_plt_cnt_0150'
energy_ranges = [(0.3, 1.5), (1.5, 3.5), (3.5, 7.0)]
fields = ['xray_emissivity_%.1f_%.1f_keV' % er for er in energy_ranges]
ds = yt.load(fname)
for energy_range in energy_ranges:
print add_xray_emissivity_field(ds, *energy_range, \
with_metals=False, constant_metallicity=0.02,\
filename='/d/d9/ychen/yt_testdata/apec_emissivity.h5')
normal = (0.0, np.tan(proj_angle / 180. * np.pi), 1.0)
north = (0.0, 0.1, 1.0)
max_lv = 6
plot = yt.OffAxisProjectionPlot(ds, normal, fields, width=((40,'kpc'), (80,'kpc')),\
depth=(80, 'kpc'), max_level=max_lv, north_vector=north)
#plot.zoom(10)
#plot.frb.export_fits(fname+'.fits', fields=fname)
#plot.save()
ext = ds.arr([-20, 20, -40, 40], input_units='kpc')
r = plot['xray_emissivity_0.3_1.5_keV'].image._A
g = plot['xray_emissivity_1.5_3.5_keV'].image._A
b = plot['xray_emissivity_3.5_7.0_keV'].image._A
image = np.array([r, g, b]).transpose([1, 2, 0])
image = image / image.max()
plt.figure(figsize=(10, 20))
plt.imshow(image, origin='lower', extent=ext)
plt.xlabel('(kpc)')
#plt.ylabel('(kpc)')
plt.savefig('xray_apec002_3color%i_lv%i.png' %
(int(proj_angle), int(max_lv)))
velx_full = proj_dict[proj_dict_keys[1]]
vely_full = proj_dict[proj_dict_keys[2]]
velz_full = proj_dict[proj_dict_keys[3]]
magx = proj_dict[proj_dict_keys[4]]
magy = proj_dict[proj_dict_keys[5]]
elif args.use_angular_momentum != 'False':
proj_root_rank = int(rank/7)*7
#proj_dict = {simfo['field'][1]:[]}
proj_dict = {simfo['field'][1]:[], 'Projected_Velocity_x':[], 'Projected_Velocity_y':[], 'Projected_Velocity_z':[], 'Projected_Magnetic_Field_x':[], 'Projected_Magnetic_Field_y':[], 'Projected_Magnetic_Field_z':[]}
proj_dict_keys = str(proj_dict.keys()).split("['")[1].split("']")[0].split("', '")
#proj_field_list =[simfo['field']]
proj_field_list =[simfo['field'], ('gas', 'Projected_Velocity_x'), ('gas', 'Projected_Velocity_y'), ('gas', 'Projected_Velocity_z'), ('gas', 'Projected_Magnetic_Field_x'), ('gas', 'Projected_Magnetic_Field_y'), ('gas', 'Projected_Magnetic_Field_z')]
for field in yt.parallel_objects(proj_field_list):
proj = yt.OffAxisProjectionPlot(ds, L, field, width=(x_width/2, 'AU'), weight_field=weight_field, method='integrate', center=(center_pos, 'AU'), depth=(args.slice_thickness, 'AU'))
if 'mag' in str(field):
if weight_field == None:
proj_array = np.array(proj.frb.data[field].in_units('cm*gauss')/thickness.in_units('cm'))
else:
proj_array = np.array(proj.frb.data[field].in_units('gauss'))
else:
if weight_field == None:
proj_array = np.array(proj.frb.data[field].in_cgs()/thickness.in_units('cm'))
else:
proj_array = np.array(proj.frb.data[field].in_cgs())
if rank == proj_root_rank:
proj_dict[field[1]] = proj_array
else:
file = open(pickle_file.split('.pkl')[0] + '_proj_data_' + str(proj_root_rank)+ str(proj_dict_keys.index(field[1])) + '.pkl', 'wb')
pickle.dump((field[1], proj_array), file)
# To see all the possible field values contained within the hydro (plt) and particle (part) files:
# for e in ds.derived_field_list:
# print(e)
# ds.print_stats()
# print(ds.derived_field_list)
# print("-"*20 + "\n", ad['x'])
# print(len(ad['x']))
# print(len(ad['temperature']))
# Say I want to make a slice plot of the density field.
field = 'density' # dens, temp, and pres are some shorthand strings recognized by yt.
# ax = 'y' # the axis our slice plot will be "looking down on".
L = [0, 0, 1] # vector normal to cutting plane
plot_ = yt.OffAxisProjectionPlot(ds, L, field)
im_name = 'off_axis_yt.png'
plot_.save(im_name)
# plot_ = yt.SlicePlot(ds, ax, field)
# plot_ = yt.ProjectionPlot(ds, ax, field)
# plot_.set_cmap(field, "binary")
#
# plot_.annotate_timestamp()
# plot_.annotate_grids()
#
# plot_.annotate_title('Testplot')
# plot_.annotate_scale()
#plot_.annotate_magnetic_field()
#plot_.show()
width = 2*radial_extent
log("Reading amiga center for halo in %s" % fn)
c = read_amiga_center(amiga_data, fn, ds)
rvir = read_amiga_rvir(amiga_data, fn, ds)
log('Finding Angular Momentum of Galaxy')
sp = ds.sphere(c, (15, 'kpc'))
L = sp.quantities.angular_momentum_vector(use_gas=False, use_particles=True, particle_type='PartType0')
L, E1, E2 = ortho_find(L)
one = ds.arr([.5, .5, .5], 'Mpc')
box = ds.box(c-one, c+one)
log('Generating Plot 1')
p1 = yt.OffAxisProjectionPlot(ds, E1, ('gas', 'H_number_density'), center=c,
width=width, data_source=box, north_vector=L, weight_field=None)
log('Generating Plot 2')
p2 = yt.OffAxisProjectionPlot(ds, E1, ('gas', 'Mg_p1_number_density'), center=c,
width=width, data_source=box, north_vector=L, weight_field=None)
log('Generating Plot 3')
p3 = yt.OffAxisProjectionPlot(ds, E1, ('gas', 'O_p5_number_density'), center=c,
width=width, data_source=box, north_vector=L, weight_field=None)
log('Generating Plot 4')
p4 = yt.OffAxisProjectionPlot(ds, E1, ('gas', 'Si_p3_number_density'), center=c,
width=width, data_source=box, north_vector=L, weight_field=None)
print("Stitching together")
fig = GridFigure(2, 2, top_buffer=0.01, bottom_buffer=0.01, left_buffer=0.01, right_buffer=0.13, vertical_buffer=0.01, horizontal_buffer=0.01, figsize=(9,8))
# Actually plot in the different axes
plot1 = fig[0].imshow(p1.frb['H_number_density'], norm=LogNorm())
def main():
rank = CW.Get_rank()
size = CW.Get_size()
args = parse_inputs()
prev_args = args
print("Starting mosaic_mod_script on rank", rank)
# Read in directories:
input_file = args.input_file
#save_dir = args.save_directory
#if os.path.exists(save_dir) == False:
# os.makedirs(save_dir)
# Read in input file
print("Reading in input mosaic file on rank", rank)
positions = []
paths = []
args_dict = []
with open(input_file, 'rU') as mosaic_file:
reader = csv.reader(mosaic_file)
for row in reader:
if row[0] == 'Grid_inputs:':
glr = float(row[1])
grl = float(row[2])
glw = float(row[3])
ghspace = float(row[4])
elif row[0][0] != '#':
positions.append((int(row[0]), int(row[1])))
paths.append(row[2])
dict = ""
for col in row[3:]:
dict = dict + col
if col != row[-1]:
dict = dict + ','
dict = ast.literal_eval(dict)
args_temp = argparse.Namespace(**vars(args))
for key in list(dict.keys()):
if key in args:
exec("args_temp."+ key + " = " + "str(dict[key])")
args_dict.append(args_temp)
del args_temp
args = prev_args
import pdb
pdb.set_trace()
positions = np.array(positions)
c = define_constants()
mym.set_global_font_size(args.text_font)
files = []
simfo = []
X = []
Y = []
X_vel = []
Y_vel = []
sim_files = []
L = None
for pit in range(len(paths)):
fs = get_files(paths[pit], args_dict[pit])
files.append(fs)
#print "paths =", paths
#print "fs =", fs
#print "args_dict =", args_dict
sfo = sim_info(paths[pit], fs[-1], args_dict[pit])
simfo.append(sfo)
if args_dict[pit].yt_proj == False:
x, y, x_vel, y_vel, cl = mym.initialise_grid(files[pit][-1], zoom_times=args_dict[pit].zoom_times)
X.append(x)
Y.append(y)
X_vel.append(x_vel)
Y_vel.append(y_vel)
else:
x = np.linspace(sfo['xmin'], sfo['xmax'], sfo['dimension'])
y = np.linspace(sfo['ymin'], sfo['ymax'], sfo['dimension'])
x, y = np.meshgrid(x, y)
annotate_space = (simfo[pit]['xmax'] - simfo[pit]['xmin'])/31.
x_ind = []
y_ind = []
counter = 0
while counter < 31:
val = annotate_space*counter + annotate_space/2. + simfo[pit]['xmin']
x_ind.append(int(val))
y_ind.append(int(val))
counter = counter + 1
x_vel, y_vel = np.meshgrid(x_ind, y_ind)
if args_dict[pit].projection_orientation != None:
y_val = 1./np.tan(np.deg2rad(float(args_dict[pit].projection_orientation)))
if np.isinf(y_val):
y_val = 0.0
L = [1.0, y_val, 0.0]
else:
if has_particles == False or len(dd['particle_posx']) == 1:
L = [0.0, 1.0, 0.0]
else:
pos_vec = [np.diff(dd['particle_posx'].value)[0], np.diff(dd['particle_posy'].value)[0]]
L = [-1*pos_vec[-1], pos_vec[0]]
L.append(0.0)
if L[0] > 0.0:
L = [-1.0*L[0], -1.0*L[1], 0.0]
print("SET PROJECTION ORIENTATION L=", L)
L = np.array(L)
X.append(x)
Y.append(y)
X_vel.append(x_vel)
Y_vel.append(y_vel)
if rank == 0:
print("shape of x, y", np.shape(x), np.shape(y))
if args_dict[pit].yt_proj == False and args_dict[pit].image_center != 0:
sim_fs = sorted(glob.glob(paths[pit] + 'WIND_hdf5_plt_cnt*'))
elif args_dict[pit].yt_proj != False and args_dict[pit].image_center != 0:
sim_fs = files
else:
sim_fs = []
sim_files.append(sim_fs)
#myf.set_normal(L)
#print "SET PROJECTION ORIENTATION L=", myf.get_normal()
# Initialise Grid and build lists
if args.plot_time != None:
m_times = [args.plot_time]
else:
m_times = mym.generate_frame_times(files[0], args.time_step, presink_frames=args.presink_frames, end_time=args.end_time)
no_frames = len(m_times)
m_times = m_times[args.start_frame:]
sys.stdout.flush()
CW.Barrier()
usable_files = []
usable_sim_files = []
for pit in range(len(paths)):
usable_fs = mym.find_files(m_times, files[pit])
usable_files.append(usable_fs)
if args_dict[pit].image_center != 0 and args_dict[pit].yt_proj == False:
usable_sfs = mym.find_files(m_times, sim_files[pit])
usable_sim_files.append(usable_fs)
del sim_files[pit]
else:
usable_sim_files.append([])
sys.stdout.flush()
CW.Barrier()
frames = list(range(args.start_frame, no_frames))
sink_form_time = []
for pit in range(len(paths)):
sink_form = mym.find_sink_formation_time(files[pit])
print("sink_form_time", sink_form_time)
sink_form_time.append(sink_form)
del files
# Define colourbar bounds
cbar_max = args.colourbar_max
cbar_min = args.colourbar_min
if L is None:
if args.axis == 'xy':
L = [0.0, 0.0, 1.0]
else:
L = [1.0, 0.0, 0.0]
L = np.array(L)
if args.axis == 'xy':
y_int = 1
else:
y_int = 2
sys.stdout.flush()
CW.Barrier()
rit = args.working_rank
for frame_val in range(len(frames)):
if rank == rit:
time_val = m_times[frame_val]
plt.clf()
columns = np.max(positions[:,0])
rows = np.max(positions[:,1])
width = float(columns)*(14.5/3.)
height = float(rows)*(17./4.)
fig =plt.figure(figsize=(width, height))
gs_left = gridspec.GridSpec(rows, columns-1)
gs_right = gridspec.GridSpec(rows, 1)
gs_left.update(right=glr, wspace=glw, hspace=ghspace)
gs_right.update(left=grl, hspace=ghspace)
axes_dict = {}
counter = 1
for pit in range(len(paths)):
try:
title_parts = args_dict[pit].title
except:
title_parts = args_dict[pit]['title']
title = ''
for part in title_parts:
if part != title_parts[-1]:
title = title + part + ' '
else:
title = title + part
ax_label = 'ax' + str(counter)
yit = np.where(positions[:,1] == positions[pit][1])[0][0]
if positions[pit][0] == 1 and positions[pit][1] == 1:
if columns > 1:
axes_dict.update({ax_label:fig.add_subplot(gs_left[0,0])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
axes_dict.update({ax_label:fig.add_subplot(gs_right[0,0])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif positions[pit][0] != columns:
if args.share_x and args.share_y:
if yit >= len(axes_dict):
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1], sharex=axes_dict['ax1'])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1], sharex=axes_dict['ax1'], sharey=axes_dict[list(axes_dict.keys())[yit]])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_x:
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[it][1]-1,positions[pit][0]-1], sharex=axes_dict['ax1'])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_y and positions[pit][0]!=1:
yit = np.where(positions[:,1] == positions[pit][1])[0][0]
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1], sharey=axes_dict[list(axes_dict.keys())[yit]])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_y:
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
axes_dict.update({ax_label:fig.add_subplot(gs_left[positions[pit][1]-1,positions[pit][0]-1])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
if args.share_x and args.share_y:
yit = np.where(positions[:,1] == positions[pit][1])[0][0]
axes_dict.update({ax_label:fig.add_subplot(gs_right[positions[pit][1]-1,0], sharex=axes_dict['ax1'], sharey=axes_dict[list(axes_dict.keys())[yit]])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_x:
axes_dict.update({ax_label:fig.add_subplot(gs_right[positions[pit][1]-1,0], sharex=axes_dict['ax1'])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
elif args.share_y:
yit = np.where(positions[:,1] == positions[pit][1])[0][0]
axes_dict.update({ax_label:fig.add_subplot(gs_right[positions[pit][1]-1,0], sharey=axes_dict[list(axes_dict.keys())[yit]])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
else:
axes_dict.update({ax_label:fig.add_subplot(gs_right[positions[pit][1]-1,0])})
#print "ADDED SUBPLOT:", counter, "on rank", rank
counter = counter + 1
axes_dict[ax_label].set(adjustable='box-forced', aspect='equal')
if args.yt_proj and args.plot_time==None and os.path.isfile(paths[pit] + "movie_frame_" + ("%06d" % frames[frame_val]) + ".pkl"):
pickle_file = paths[pit] + "movie_frame_" + ("%06d" % frames[frame_val]) + ".pkl"
print("USING PICKLED FILE:", pickle_file)
file = open(pickle_file, 'r')
#weight_fieldstuff = pickle.load(file)
X[pit], Y[pit], image, magx, magy, X_vel[pit], Y_vel[pit], velx, vely, part_info, args_dict[pit], simfo[pit] = pickle.load(file)
#file_time = stuff[17]
file.close()
else:
time_val = m_times[frame_val]
print("FILE =", usable_files[pit][frame_val])
has_particles = has_sinks(usable_files[pit][frame_val])
if has_particles:
part_info = mym.get_particle_data(usable_files[pit][frame_val], args_dict[pit].axis, proj_or=L)
else:
part_info = {}
center_vel = [0.0, 0.0, 0.0]
if args.image_center != 0 and has_particles:
original_positions = [X[pit], Y[pit], X_vel[pit], y_vel[pit]]
x_pos = np.round(part_info['particle_position'][0][args.image_center - 1]/cl)*cl
y_pos = np.round(part_info['particle_position'][1][args.image_center - 1]/cl)*cl
pos = np.array([part_info['particle_position'][0][args.image_center - 1], part_info['particle_position'][1][args.image_center - 1]])
X[pit] = X[pit] + x_pos
Y[pit] = Y[pit] + y_pos
X_vel[pit] = X_vel[pit] + x_pos
Y_vel[pit] = Y_vel[pit] + y_pos
if args.yt_proj == False:
sim_file = usable_sim_files[frame_val][:-12] + 'part' + usable_sim_files[frame_val][-5:]
else:
sim_file = part_file
if len(part_info['particle_mass']) == 1:
part_ind = 0
else:
min_dist = 1000.0
for part in range(len(part_info['particle_mass'])):
f = h5py.File(sim_file, 'r')
temp_pos = np.array([f[list(f.keys())[11]][part][13]/c['au'], f[list(f.keys())[11]][part][13+y_int]/c['au']])
f.close()
dist = np.sqrt(np.abs(np.diff((temp_pos - pos)**2)))[0]
if dist < min_dist:
min_dist = dist
part_ind = part
f = h5py.File(sim_file, 'r')
center_vel = [f[list(f.keys())[11]][part_ind][18], f[list(f.keys())[11]][part_ind][19], f[list(f.keys())[11]][part_ind][20]]
f.close()
xabel, yabel, xlim, ylim = image_properties(X[pit], Y[pit], args_dict[pit], simfo[pit])
if args_dict[pit].axis == 'xy':
center_vel=center_vel[:2]
else:
center_vel=center_vel[::2]
if args_dict[pit].ax_lim != None:
if has_particles and args_dict[pit].image_center != 0:
xlim = [-1*args_dict[pit].ax_lim + part_info['particle_position'][0][args_dict[pit].image_center - 1], args_dict[pit].ax_lim + part_info['particle_position'][0][args_dict[pit].image_center - 1]]
ylim = [-1*args_dict[pit].ax_lim + part_info['particle_position'][1][args_dict[pit].image_center - 1], args_dict[pit].ax_lim + part_info['particle_position'][1][args_dict[pit].image_center - 1]]
else:
xlim = [-1*args_dict[pit].ax_lim, args_dict[pit].ax_lim]
ylim = [-1*args_dict[pit].ax_lim, args_dict[pit].ax_lim]
if args.yt_proj == False:
f = h5py.File(usable_files[pit][frame_val], 'r')
image = get_image_arrays(f, simfo[pit]['field'], simfo[pit], args_dict[pit], X[pit], Y[pit])
magx = get_image_arrays(f, 'mag'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis, simfo[pit], args_dict[pit], X[pit], Y[pit])
magy = get_image_arrays(f, 'mag'+args.axis[1]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis, simfo[pit], args_dict[pit], X[pit], Y[pit])
x_pos_min = int(np.round(np.min(X[pit]) - simfo[pit]['xmin_full'])/simfo[pit]['cell_length'])
y_pos_min = int(np.round(np.min(Y[pit]) - simfo[pit]['xmin_full'])/simfo[pit]['cell_length'])
if np.shape(f['vel'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis]) == (2048, 2048):
velocity_data = [f['vel'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis], f['vel'+args.axis[1]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis]]
elif args.axis == 'xy':
velocity_data = [f['vel'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis][:,:,0], f['vel'+args.axis[1]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis][:,:,0]]
else:
velocity_data = [f['vel'+args.axis[0]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis][:,0,:], f['vel'+args.axis[1]+'_'+simfo[pit]['movie_file_type']+'_'+args.axis][:,0,:]]
velx, vely = mym.get_quiver_arrays(y_pos_min, x_pos_min, X[pit], velocity_data[0], velocity_data[1], center_vel=center_vel)
else:
if args_dict[pit].image_center == 0 or has_particles == False:
center_pos = np.array([0.0, 0.0, 0.0])
else:
dd = f.all_data()
center_pos = np.array([dd['particle_posx'][args.image_center-1].in_units('AU'), dd['particle_posy'][args.image_center-1].in_units('AU'), dd['particle_posz'][args.image_center-1].in_units('AU')])
x_width = (xlim[1] -xlim[0])
y_width = (ylim[1] -ylim[0])
thickness = yt.YTArray(args.slice_thickness, 'AU')
proj = yt.OffAxisProjectionPlot(f, L, [simfo[pit]['field'], 'cell_mass', 'velz_mw', 'magz_mw', 'Projected_Magnetic_Field_mw', 'Projected_Velocity_mw'], center=(center_pos, 'AU'), width=(x_width, 'AU'), depth=(args.slice_thickness, 'AU'))
image = (proj.frb.data[simfo[pit]['field']]/thickness.in_units('cm')).value
velx_full = (proj.frb.data[('gas', 'Projected_Velocity_mw')].in_units('g*cm**2/s')/thickness.in_units('cm')).value
vely_full = (proj.frb.data[('gas', 'velz_mw')].in_units('g*cm**2/s')/thickness.in_units('cm')).value
magx = (proj.frb.data[('gas', 'Projected_Magnetic_Field_mw')].in_units('g*gauss*cm')/thickness.in_units('cm')).value
magy = (proj.frb.data[('gas', 'magz_mw')].in_units('g*gauss*cm')/thickness.in_units('cm')).value
mass = (proj.frb.data[('gas', 'cell_mass')].in_units('cm*g')/thickness.in_units('cm')).value
velx_full = velx_full/mass
vely_full = vely_full/mass
magx = magx/mass
magy = magy/mass
del mass
velx, vely = mym.get_quiver_arrays(0.0, 0.0, X[pit], velx_full, vely_full, center_vel=center_vel)
del velx_full
del vely_full
if len(frames) == 1:
if rank == 0:
pickle_file = paths[pit] + "movie_frame_" + ("%06d" % frames[frame_val]) + ".pkl"
file = open(pickle_file, 'w+')
pickle.dump((X[pit], Y[pit], image, magx, magy, X_vel[pit], Y_vel[pit], velx, vely, xlim, ylim, has_particles, part_info, simfo[pit], time_val,xabel, yabel), file)
file.close()
print("Created Pickle:", pickle_file, "for file:", usable_files[pit][frame_val])
else:
pickle_file = paths[pit] + "movie_frame_" + ("%06d" % frames[frame_val]) + ".pkl"
file = open(pickle_file, 'w+')
pickle.dump((X[pit], Y[pit], image, magx, magy, X_vel[pit], Y_vel[pit], velx, vely, xlim, ylim, has_particles, part_info, simfo[pit], time_val,xabel, yabel), file)
file.close()
print("Created Pickle:", pickle_file, "for file:", usable_files[pit][frame_val])
f.close()
plot = axes_dict[ax_label].pcolormesh(X[pit], Y[pit], image, cmap=plt.cm.gist_heat, norm=LogNorm(vmin=cbar_min, vmax=cbar_max), rasterized=True)
plt.gca().set_aspect('equal')
if frame_val > 0 or time_val > -1.0:
axes_dict[ax_label].streamplot(X[pit], Y[pit], magx, magy, density=4, linewidth=0.25, arrowstyle='-', minlength=0.5)
else:
axes_dict[ax_label].streamplot(X[pit], Y[pit], magx, magy, density=4, linewidth=0.25, minlength=0.5)
xlim = args_dict[pit]['xlim']
ylim = args_dict[pit]['ylim']
mym.my_own_quiver_function(axes_dict[ax_label], X_vel[pit], Y_vel[pit], velx, vely, plot_velocity_legend=bool(args_dict[pit]['annotate_velocity']), limits=[xlim, ylim], standard_vel=args.standard_vel)
if args_dict[pit]['has_particles']:
if args.annotate_particles_mass == True:
mym.annotate_particles(axes_dict[ax_label], part_info['particle_position'], part_info['accretion_rad'], limits=[xlim, ylim], annotate_field=part_info['particle_mass'])
else:
mym.annotate_particles(axes_dict[ax_label], part_info['particle_position'], part_info['accretion_rad'], limits=[xlim, ylim], annotate_field=None)
if args.plot_lref == True:
r_acc = np.round(part_info['accretion_rad'])
axes_dict[ax_label].annotate('$r_{acc}$='+str(r_acc)+'AU', xy=(0.98*simfo[pit]['xmax'], 0.93*simfo[pit]['ymax']), va="center", ha="right", color='w', fontsize=args_dict[pit].text_font)
if args.annotate_time == "True" and pit == 0:
print("ANNONTATING TIME:", str(int(time_val))+'yr')
time_text = axes_dict[ax_label].text((xlim[0]+0.01*(xlim[1]-xlim[0])), (ylim[1]-0.03*(ylim[1]-ylim[0])), '$t$='+str(int(time_val))+'yr', va="center", ha="left", color='w', fontsize=args.text_font)
time_text.set_path_effects([path_effects.Stroke(linewidth=3, foreground='black'), path_effects.Normal()])
#ax.annotate('$t$='+str(int(time_val))+'yr', xy=(xlim[0]+0.01*(xlim[1]-xlim[0]), ylim[1]-0.03*(ylim[1]-ylim[0])), va="center", ha="left", color='w', fontsize=args.text_font)
title_text = axes_dict[ax_label].text((np.mean(xlim)), (ylim[1]-0.03*(ylim[1]-ylim[0])), title, va="center", ha="center", color='w', fontsize=(args.text_font+2))
title_text.set_path_effects([path_effects.Stroke(linewidth=3, foreground='black'), path_effects.Normal()])
if positions[pit][0] == columns:
cbar = plt.colorbar(plot, pad=0.0, ax=axes_dict[ax_label])
cbar.set_label('Density (gcm$^{-3}$)', rotation=270, labelpad=14, size=args.text_font)
axes_dict[ax_label].set_xlabel(args_dict[pit]['xabel'], labelpad=-1, fontsize=args.text_font)
if positions[pit][0] == 1:
axes_dict[ax_label].set_ylabel(args_dict[pit]['yabel'], labelpad=-20, fontsize=args.text_font)
axes_dict[ax_label].set_xlim(xlim)
axes_dict[ax_label].set_ylim(ylim)
for line in axes_dict[ax_label].xaxis.get_ticklines():
line.set_color('white')
for line in axes_dict[ax_label].yaxis.get_ticklines():
line.set_color('white')
plt.tick_params(axis='both', which='major', labelsize=16)
for line in axes_dict[ax_label].xaxis.get_ticklines():
line.set_color('white')
for line in axes_dict[ax_label].yaxis.get_ticklines():
line.set_color('white')
if positions[pit][0] != 1:
yticklabels = axes_dict[ax_label].get_yticklabels()
plt.setp(yticklabels, visible=False)
if positions[pit][0] == 1:
axes_dict[ax_label].tick_params(axis='y', which='major', labelsize=args.text_font)
if positions[pit][1] == rows:
axes_dict[ax_label].tick_params(axis='x', which='major', labelsize=args.text_font)
if positions[pit][0] != 1:
xticklabels = axes_dict[ax_label].get_xticklabels()
plt.setp(xticklabels[0], visible=False)
if len(usable_files[pit]) > 1:
if args.output_filename == None:
import pdb
pdb.set_trace()
file_name = save_dir + "movie_frame_" + ("%06d" % frames[frame_val])
else:
file_name = args.output_filename + "_" + str(int(time_val))
else:
if args.output_filename != None:
file_name = args.output_filename
else:
import pdb
pdb.set_trace()
file_name = save_dir + "time_" + str(args.plot_time)
plt.savefig(file_name + ".eps", format='eps', bbox_inches='tight')
#plt.savefig(file_name + ".pdf", format='pdf', bbox_inches='tight')
#plt.savefig(file_name + ".jpg", format='jpeg', bbox_inches='tight')
call(['convert', '-antialias', '-quality', '100', '-density', '200', '-resize', '100%', '-flatten', file_name+'.eps', file_name+'.jpg'])
os.remove(file_name + '.eps')
del image
del magx
del magy
del velx
del vely
if args.image_center != 0 and has_particles:
X[pit], Y[pit], X_vel[pit], Y_vel[pit] = original_positions
print('Created frame', (frames[frame_val]), 'of', str(frames[-1]), 'on rank', rank, 'at time of', str(time_val), 'to save_dir:', file_name + '.eps')
rit = rit +1
if rit == size:
rit = 0
print("completed making movie frames on rank", rank)
myf.set_north_vector(north_unit)
#These numbers are just creative pickle file management
div_32 = int(rank/32)
rem_32 = np.remainder(rank,32)
div_4 = int(rem_32/4)
proj_root_rank = div_32*32 + div_4*4
#Fields to project
field_list = [field, ('gas', 'Radial_Velocity'), ('gas', 'Proj_x_velocity'), ('gas', 'Proj_y_velocity')]
proj_dict = {field[1]:[], 'Radial_Velocity':[], 'Proj_x_velocity':[], 'Proj_y_velocity':[]}
proj_dict_keys = str(proj_dict.keys()).split("['")[1].split("']")[0].split("', '")
for field in yt.parallel_objects(field_list):
print("Calculating projection with normal", proj_vector_unit, "for field", field, "on rank", rank)
proj = yt.OffAxisProjectionPlot(ds, proj_vector_unit, field, width=(2*args.ax_lim, 'AU'), weight_field=weight_field, method='integrate', center=(center_pos.value, 'AU'), depth=(args.slice_thickness, 'AU'), north_vector=north_unit)
if args.resolution != 800:
proj.set_buff_size([args.resolution, args.resolution])
if args.field in str(field):
if weight_field == None:
if args.divide_by_proj_thickness == "True":
proj_array = np.array((proj.frb.data[field]/thickness.in_units('cm')).in_units(args.field_unit))
else:
proj_array = np.array(proj.frb.data[field].in_units(args.field_unit+"*cm"))
else:
if args.divide_by_proj_thickness == "True":
proj_array = np.array(proj.frb.data[field].in_units(args.field_unit))
else:
proj_array = np.array(proj.frb.data[field].in_units(args.field_unit)*thickness.in_units('cm'))
def make_two_plots(file_name, ds, sph, center, L, Lx, z, sfr, zoom=70):
fig = plt.figure()
# creates grid on which the figure will be plotted
grid = AxesGrid(fig, (.5, .5, 1.5, 1.5),
nrows_ncols=(1, 3),
axes_pad=.8,
label_mode="L",
cbar_location="right",
cbar_mode="each",
cbar_set_cax=False,
cbar_size="3%",
cbar_pad="0%",
direction='column')
sns.set_style("whitegrid", {'axes.grid': False})
# creates density projection plot
pro = yt.OffAxisProjectionPlot(ds,
Lx, ('gas', 'density'),
center=center,
width=(zoom, 'kpc'),
north_vector=L)
pro.set_font({'size': 10})
pro.hide_axes()
pro.set_unit(('gas', 'density'), 'Msun/pc**2')
pro.set_zlim(('gas', 'density'), 2, 500)
cmap = sns.blend_palette(("black", "#984ea3", "#d73027", "darkorange",
"#ffe34d", "#4daf4a", "white"),
n_colors=60,
as_cmap=True)
pro.set_cmap(("gas", "density"), cmap)
pro.annotate_scale(size_bar_args={'color': 'white'})
# draws density axes onto main figure
pro.plots[('gas', 'density')].figure = fig
pro.plots[('gas', 'density')].axes = grid[0].axes
pro.plots[('gas', 'density')].cax = grid.cbar_axes[0]
pro._setup_plots()
# creates metallicity plot
oap = yt.SlicePlot(ds,
Lx, ("gas", "metallicity"),
center=center,
width=(zoom, 'kpc'),
north_vector=L,
data_source=sph)
oap.annotate_cquiver('r_flux', 'cyl_flux', 16)
oap.set_font({'size': 10})
oap.hide_axes()
oap.set_zlim(("gas", "metallicity"), 2e-4, 10.5)
cmap = sns.blend_palette(
("black", "#984ea3", "#4575b4", "#4daf4a", "#ffe34d", "darkorange"),
as_cmap=True)
oap.set_cmap(("gas", "metallicity"), cmap)
# draws metallicity axes onto main figure
oap.plots[("gas", "metallicity")].figure = fig
oap.plots[("gas", "metallicity")].axes = grid[1].axes
oap.plots[("gas", "metallicity")].cax = grid.cbar_axes[1]
oap._setup_plots()
oap.annotate_scale()
# creates temperature plot
tp = yt.SlicePlot(ds,
Lx, ("gas", "temperature"),
center=center,
width=(zoom, 'kpc'),
north_vector=L,
data_source=sph)
tp.set_font({'size': 10})
tp.hide_axes()
tp.set_zlim(("gas", "temperature"), 400, 8e7)
cmap = sns.blend_palette(("black", "#d73027", "darkorange", "#ffe34d"),
n_colors=50,
as_cmap=True)
tp.set_cmap(("gas", "temperature"), cmap)
# draws axes onto main figure
tp.plots[('gas', 'temperature')].figure = fig
tp.plots[('gas', 'temperature')].axes = grid[2].axes
tp.plots[('gas', 'temperature')].cax = grid.cbar_axes[2]
tp._setup_plots()
tp.annotate_scale()
# creates SFH plot below the first three panels on main figure
ax3 = fig.add_subplot(211, position=[.5, .8, 1.5, .17])
ax3.plot(z, sfr)
ax3.set_xlim(1.05, .45)
ax3.set_ylim(-2, 33)
ax3.set_xlabel('Redshift')
ax3.set_ylabel(r'SFR (M$_\odot$/yr)')
# saves the figure
plt.savefig(plot_dir + file_name[-6:] + '_plot.png', bbox_inches='tight')
def ks_law_new(gas_cylinder,
star_cylinder,
ytsnap,
star_center,
image_width=800,
field="stars",
n_bins=100,
r_max=None):
# hard coded things
# really should be some sort of config
width = ytsnap.arr(80, "kpc")
depth = ytsnap.arr(20, "kpc")
age = ytsnap.arr(
3, "Gyr"
) # increasing this increases the number of resolved points (reduces the "artificial lines")
max_age = age.in_units("yr").value
weight_field = None
if r_max == None:
r_max = gas_cylinder.ds.arr(
15,
"kpc") # reducing this reduces the width of the artificial lines
print "Computing KS law"
print "ks 1) Gathering initial projections"
# todo.. make the output a histogram
width = ((width.v, width.units), (width.v, width.units))
image_width = (image_width, image_width)
# slice of the star particle mass
normal = gas_cylinder.get_field_parameter("normal")
data_source = gen_data_source(0, gas_cylinder, ytsnap, width, width[0],
"kpc")
plot = yt.OffAxisProjectionPlot(data_source.ds,
normal, [
("gas", "density"),
],
center=gas_cylinder.center,
width=width,
depth=width[0],
weight_field=weight_field,
north_vector=normal)
# set the units
# plot.set_axes_unit("kpc")
# plot.set_unit(("gas","density"),"Msun/pc**2")
# plot.set_zlim(("gas","density"),10**(-3),10**4)
images = plot.frb
gas_image = images[('gas', 'density')]
# plot = yt.ProjectionPlot(star_cylinder,2,[("deposit","io_cic"),("deposit","io_density")],center=star_center,width=width[0],weight_field=None)
# # set the units
# plot.set_axes_unit("kpc")
# plot.set_unit(("deposit","io_cic"),"Msun/pc**2")
# plot.set_zlim(("deposit","io_cic"),10**(-3),10**8)
## plot.save("test_star_cylinder")
# plot.annotate_particles(width=(20,"kpc"), p_size=2.0, col='k', marker='o', stride=1.0, ptype="io", minimum_mass=None, alpha=1.0)
# plot.save("test_star_cylinder_stars")
## plot.set_unit(("deposit","io_cic"),"Msun/kpc**2")
# plot.set_unit(("deposit","io_density"),"Msun/kpc**2")
images = plot.frb
star_image = images[('deposit', 'io_cic')]
print "ks 2) collecting data"
# filtering things withing the first 15 kpc does the trick usuall
n_bins = gas_image.shape[0]
bin_width = width[0][0] / float(n_bins)
left = -width[0][0] / 2.0
right = width[0][0] / 2.0
length = np.arange(left, right, bin_width)
shift_thing = (length[2] - length[1]) / 2.0
length = length + shift_thing
# to make radial array, we need to calculate the radial distance of each pixel in a 2d array. Namely.
length_squared = np.power(length, 2.0) # power of 2
radius_squared = np.add.outer(length_squared,
length_squared) # adds as an outer product
radius = np.sqrt(radius_squared)
# now we need to generate some radial bins
radius_bin_width = radius.max() / float(n_bins)
radius_bins = np.arange(0, radius.max(), radius_bin_width)
# flatten the arrays
flat_radius = radius.flatten()
# now need to find the indices in which radius is maximum
radius_filter = (flat_radius < r_max.in_units("kpc").v)
# ok, now get the data
gas_density = np.array(gas_image.in_units("Msun/pc**2").value).flatten()
star_density = np.array(star_image.in_units("Msun/kpc**2").value).flatten()
# filter by radius
gas_density = np.log10(gas_density[radius_filter])
star_density = np.log10(star_density[radius_filter] / max_age)
return gas_density, star_density
def plot_face_edge(isnap=28, selected_field='h2density', sizekpc=7., cutLow=1.e-5, f_out=None, save_plot=True,
overplot_clumps=True, incut=6., field_cut='h2density', n_cell_min=8, largeNum=1.e+42,
data=None, ds=None, dd=None, leaf_clumps=None, plotClumpID=False
):
"""
Parameters
----------
isnap: int
number of snaphot select for the plot
selected_field: str
field for image plot
sizekpc: float
size of the extracted region from isnap
cutLow: float
dynamical range for the colorbar
f_out: str
output filename
save_plot: bool
print to file
overplot_clumps: bool
overplot clumps
incut: float
threshold for clump definition
field_cut: str
field out of which clumps are identified
n_cell_min: int
min cell to define a clump
largeNum: float
BS to have YT to collaborate
"""
if f_out is None:
f_out = 'test_' + selected_field + '_out' + \
str(isnap) + '_yt_unit_plot.png'
from mpl_toolkits.axes_grid import AxesGrid
import matplotlib.pyplot as plt
# prepare input data, if not already passed as arguments
#
if data is None:
data = import_fetch_gal(isnap=isnap)
#
if (ds is None) and (dd is None):
ds, dd = prepare_unigrid(data=data,
add_unit=True,
regionsize_kpc=sizekpc,
debug=False)
#
if overplot_clumps and (leaf_clumps is None):
__, leaf_clumps = ytclumpfind_H2(ds, dd, field_cut, incut,
c_max=None, step=1e+6,
N_cell_min=n_cell_min, save=False,
plot=False, saveplot=None, fold_out='./')
if overplot_clumps:
id_sorted = sorted(range(len(leaf_clumps)),
key=lambda x: np.sum(leaf_clumps[x]["density"]))
# compute inertia tensor and it principal axes
e_value, e_vectors = calculate_eigenvektoren(data=data, sizekpc=sizekpc)
# references
# los_vec = e_vectors[0,:] # face on
# los_vec = e_vectors[1,:] # edge on
# los_vec = e_vectors[2,:] # face on (again, perpendicular direction)
#
# set the camera axes along the principal axes of the inertia tensor
vec_list = [e_vectors[2, :], e_vectors[1, :]]
up_list = [e_vectors[0, :], e_vectors[0, :]]
# setup the plot
#
fig = plt.figure()
#
grid = AxesGrid(fig, (0.075, 0.075, 0.85, 0.85),
nrows_ncols=(1, 2),
axes_pad=0.7,
#label_mode = "L",
label_mode="1",
share_all=True,
cbar_location="right",
cbar_mode="single",
cbar_size="3%",
cbar_pad=0.03)
# plot edge on and face on
#
for iplot, los_vec, up_vec in zip(xrange(2), vec_list, up_list):
prj = yt.OffAxisProjectionPlot(ds=ds, center=[0, 0, 0], normal=los_vec, fields=selected_field, width=(4, 'kpc'), north_vector=up_vec, weight_field='density',
)
prj.set_cmap(field=selected_field, cmap='inferno')
if selected_field == 'density':
selected_unit = 'Msun/pc**3'
elif selected_field == 'h2density':
selected_unit = '1/cm**3'
else:
raise TypeError('unit not implemented for the field')
prj.set_unit(selected_field, selected_unit)
prj.set_zlim(selected_field, cutLow * dd[selected_field].max().to(selected_unit),
dd[selected_field].max().to(selected_unit))
# annotate clump if asked for
if overplot_clumps:
prj.annotate_contour(field=field_cut, ncont=1, factor=1,
clim=(incut, largeNum), plot_args={
'colors': 'white'}
) # to deal w/ stupid yt annotate_clump() bug
for ileaf in id_sorted:
_fc = np.mean(leaf_clumps[ileaf].data.fcoords[:], axis=0)
if plotClumpID:
prj.annotate_marker(_fc,
coord_system='data',
plot_args={'color': 'red', 's': 50})
prj.annotate_text(_fc,
ileaf + 1,
coord_system='data',
text_args={'color': 'red', 'size': 25},
inset_box_args={'boxstyle': 'square',
'facecolor': 'white',
'linewidth': 1.0,
'edgecolor': 'white',
'alpha': 0.})
prj.set_ylabel('kpc')
prj.set_xlabel('kpc')
prj.set_font({'family': 'Times', # 'style': 'italic',
# 'weight': 'bold',
'size': 26})
# set the plot into the grid
plot = prj.plots[selected_field]
plot.figure = fig
plot.axes = grid[iplot].axes
plot.cax = grid.cbar_axes[iplot]
if iplot == 1:
plot.axes.set_axis_off()
# plot.axes.set_label('')
# plot.axes.set_visible(False)
plot.axes.set_xlabel('')
plot.axes.set_xticks(())
plot.axes.set_xticklabels('')
# Finally, this actually redraws the plot.
prj._setup_plots()
for cax in grid.cbar_axes:
cax.toggle_label(True)
# cax.axis[cax.orientation].set_label(field_select)
if save_plot:
prj.save(f_out, mpl_kwargs={'bbox_inches': 'tight'})
print 'dump to ', f_out
return prj, plot.axes
idx_start = args.idx_start
idx_end = args.idx_end
didx = args.didx
prefix_in = args.prefix_in
prefix_out = args.prefix_out
yt.enable_parallelism()
ts = yt.load([
prefix_in + '/Data_%06d' % idx
for idx in range(idx_start, idx_end + 1, didx)
])
for ds in ts.piter():
# project data
plt = yt.OffAxisProjectionPlot(ds,
normal=normal,
fields=field,
north_vector=north)
# save image
plt.set_cmap(field, colormap)
plt.annotate_timestamp(time_unit='code_time', corner='upper_right')
plt.save(prefix_out + '_' + ds.basename + '.png', mpl_kwargs={'dpi': dpi})
# access the image buffer
data = plt._frb.data[field]
print("units = %s" % data.units)
print("max = %s" % data.d.max())
# %% density
# ----------
field = "density"
fig = plt.figure()
grid = AxesGrid(fig, (0.075,0.075,0.85,0.85),
nrows_ncols = (2, len(models)),
axes_pad = 0.05,
label_mode = "L",
share_all = True,
cbar_location="right",
cbar_mode="single",
cbar_size="5%",
cbar_pad="1%")
for i in range(len(models)):
px = yt.OffAxisProjectionPlot(ds[i], [0,0,1], ('gas', '{}'.format(field)), center = center[i], width=45, north_vector=[0,1,0])
px.set_xlabel("x [kpc]")
px.set_ylabel("y [kpc]")
px.set_font({'size':20})
px.set_zlim(field="{}".format(field), zmin=5e-6, zmax=1e-1)
px.annotate_text([-20, 18], models[i][1], coord_system="plot")
plot = px.plots['{}'.format(field)]
plot.figure = fig
plot.axes = grid[i].axes
plot.cax = grid.cbar_axes[i]
px.set_background_color(field)
px._setup_plots()
for i in range(len(models)):
px = yt.OffAxisProjectionPlot(ds[i], [0,1,0], ('gas', '{}'.format(field)), center = center[i], width=45, north_vector=[0,0,-1])
px.set_xlabel("x [kpc]")
px.set_ylabel("z [kpc]")
def ProjectionPlot(Param_Dict, worker):
"""Takes a DataSet object loaded with yt and performs a projectionPlot on
it.
Parameters:
Param_Dict: Dict with Parameters
"""
ds = Param_Dict["CurrentDataSet"]
field = Param_Dict["ZAxis"]
if Param_Dict["DomainDiv"]:
# in case the user wants to divide everything by the domain_height,
# we define a new field which is just the old field divided by height
# and then do a projectionPlot for that.
height = Param_Dict["FieldMaxs"]["DomainHeight"] - Param_Dict["FieldMins"]["DomainHeight"]
field = "Normed " + field
unit = yt.units.unit_object.Unit(Param_Dict["ZUnit"] + "/cm")
realHeight = height.to_value("au") # Important! Bugs occur if we just used "height"
def _NormField(field, data):
return data[Param_Dict["ZAxis"]]/realHeight/yt.units.au
if Param_Dict["ParticlePlot"]:
ds.add_field(("io", field), function=_NormField,
units="auto", dimensions=unit.dimensions,
force_override=True, particle_type=True)
else:
ds.add_field(("gas", field), function=_NormField,
units="auto", dimensions=unit.dimensions,
force_override=True)
gridUnit = Param_Dict["GridUnit"]
c0 = yt.YTQuantity(Param_Dict["XCenter"], gridUnit)
c1 = yt.YTQuantity(Param_Dict["YCenter"], gridUnit)
if Param_Dict["Geometry"] == "cartesian":
c2 = yt.YTQuantity(Param_Dict["ZCenter"], gridUnit)
else:
c2 = Param_Dict["ZCenter"]
width = (Param_Dict["HorWidth"], gridUnit)
height = (Param_Dict["VerWidth"], gridUnit)
if Param_Dict["ParticlePlot"]:
plot = yt.ParticleProjectionPlot(ds, Param_Dict["NAxis"], field,
axes_unit=Param_Dict["GridUnit"],
weight_field=Param_Dict["WeightField"],
fontsize=14, center=[c0, c1, c2],
width=(width, height))
else:
if Param_Dict["NormVecMode"] == "Axis-Aligned":
plot = yt.ProjectionPlot(ds, Param_Dict["NAxis"], field,
axes_unit=Param_Dict["GridUnit"],
weight_field=Param_Dict["WeightField"],
fontsize=14, center=[c0, c1, c2],
width=(width, height))
else:
normVec = [Param_Dict[axis + "NormDir"] for axis in ["X", "Y", "Z"]]
northVec = [Param_Dict[axis + "NormNorth"] for axis in ["X", "Y", "Z"]]
plot = yt.OffAxisProjectionPlot(ds, normVec, field,
north_vector=northVec,
weight_field=Param_Dict["WeightField"],
axes_unit=Param_Dict["GridUnit"],
fontsize=14, center=[c0, c1, c2],
width=(width, height))
emitStatus(worker, "Setting projection plot modifications")
# Set min, max, unit log and color scheme:
setAxisSettings(plot, Param_Dict, "Z")
plot.zoom(Param_Dict["Zoom"])
emitStatus(worker, "Annotating the projection plot")
annotatePlot(Param_Dict, plot)
finallyDrawPlot(plot, Param_Dict, worker)
def ks_law(ytdataset,
raw_snapshot,
width,
max_age,
image_width=800,
field="stars",
n_bins=100,
r_max=None):
"""
Use this to compute the ks_law for each individual pixel within a galaxy
"""
weight_field = None
if r_max == None:
r_max = ytdataset.ds.arr(25, "kpc")
print "Computing KS law"
print "ks 1) Gathering initial projections"
width = ((width.v, width.units), (width.v, width.units))
image_width = (image_width, image_width)
# slice of the star particle mass
normal = ytdataset.get_field_parameter("normal")
data_source = gen_data_source(0, ytdataset, raw_snapshot, width, width[0],
"kpc")
plot = yt.OffAxisProjectionPlot(data_source.ds,
normal,
[("gas", "density"),
("deposit", "%s_density" % field)],
center=ytdataset.center,
width=width,
depth=width[0],
weight_field=weight_field,
north_vector=normal)
# set the units
plot.set_axes_unit("kpc")
plot.set_unit(("gas", "density"), "Msun/pc**2")
plot.set_unit(("deposit", "%s_density" % field), "Msun/kpc**2")
images = plot.frb
gas_image = images[('gas', 'density')]
star_image = images[('deposit', '%s_density' % field)]
print "ks 2) collecting data"
# filtering things withing the first 15 kpc does the trick usuall
n_bins = gas_image.shape[0]
bin_width = width[0][0] / float(n_bins)
left = -width[0][0] / 2.0
right = width[0][0] / 2.0
length = np.arange(left, right, bin_width)
shift_thing = (length[2] - length[1]) / 2.0
length = length + shift_thing
# to make radial array, we need to calculate the radial distance of each pixel in a 2d array. Namely.
length_squared = np.power(length, 2.0) # power of 2
radius_squared = np.add.outer(length_squared,
length_squared) # adds as an outer product
radius = np.sqrt(radius_squared)
# now we need to generate some radial bins
radius_bin_width = radius.max() / float(n_bins)
radius_bins = np.arange(0, radius.max(), radius_bin_width)
# flatten the arrays
flat_radius = radius.flatten()
# now need to find the indices in which radius is maximum
radius_filter = (flat_radius < r_max.in_units("kpc").v)
# ok, now get the data
gas_density = np.array(gas_image.in_units("Msun/pc**2").value).flatten()
star_density = np.array(star_image.in_units("Msun/kpc**2").value).flatten()
# filter by radius
gas_density = np.log10(gas_density[radius_filter])
star_density = np.log10(star_density[radius_filter] / max_age)
return gas_density, star_density
nu=nu,
proj_axis=proj_axis,
extend_cells=gc)
ds_sync = ds
yt.mylog.info('Making projection plots from file: %s', sync_fname)
if proj_axis == 'x':
plot = yt.ProjectionPlot(ds_sync,
proj_axis,
stokes.IQU,
center=[0, 0, 0],
width=width)
else:
plot = yt.OffAxisProjectionPlot(ds_sync,
proj_axis,
stokes.IQU,
center=[0, 0, 0],
width=width,
north_vector=[0, 0, 1])
plot.set_buff_size(res)
plot.set_axes_unit('kpc')
# Setting up colormaps
# Use "hot" for intensity plot and seismic for Q and U plots
for field in stokes.IQU:
if 'nn_emissivity_i' in field[1]:
plot.set_zlim(field, 1E-3 / norm, 1E1 / norm)
#cmap = plt.cm.get_cmap("algae")
#cmap.set_bad((80./256., 0.0, 80./256.))
cmap = plt.cm.hot
cmap.set_bad('k')
plot.set_cmap(field, cmap)
gas_ang_mom_norm = gas_ang_mom / gas_ang_mom_tot
sim_dict.update([('gas L vector', gas_ang_mom_norm)])
edge_on_gas = np.random.randn(3)
edge_on_gas -= edge_on_gas.dot(gas_ang_mom_norm) * gas_ang_mom_norm / np.linalg.norm(gas_ang_mom_norm)**2
sim_dict.update([('gas edge-on', edge_on_gas)])
# Set cylinder
cyl_rad = np.ceil(gal_r90.value)
cyl_h = args.cyl_h
cyl = ds.disk(center = cen_cen, normal = star_ang_mom_norm, radius = (gal_r90), height = (cyl_h, 'kpc'))
# Making the projections: currently each one renders in a bit over half an hour
if False:
if True:
prj = yt.OffAxisProjectionPlot(ds, normal = star_ang_mom_norm, fields = ('deposit', 'stars_density'), center = cen_cen, width=(proj_width, 'kpc'), data_source = cyl)
prj.save('{}/StarsFaceOn{}-z={}.png'.format(figdir, sim, redshift))
if True:
prj = yt.OffAxisProjectionPlot(ds, normal = edge_on_dir, fields = ('deposit', 'stars_density'), center = cen_cen, width=(proj_width, 'kpc'), data_source = cyl)
prj.save('{}/StarsEdgeOn{}-z={}.png'.format(figdir, sim, redshift))
if True:
prj = yt.OffAxisProjectionPlot(ds, normal = gas_ang_mom_norm, fields = 'density', center = cen_cen, width=(proj_width, 'kpc'), data_source = cyl)
prj.save('{}/GasFaceOn{}-z={}.png'.format(figdir, sim, redshift))
if True:
prj = yt.OffAxisProjectionPlot(ds, normal = edge_on_gas, fields = 'density', center = cen_cen, width=(proj_width, 'kpc'), data_source = cyl)
prj.save('{}/GasEdgeOn{}-z={}.png'.format(figdir, sim, redshift))
# Making profile plots for velocity dispersion in the cylinder
def worker_fn(file):
ds = yt.load(file.fullpath)
#for nu in nus:
# write_synchrotron_hdf5(ds, ptype, nu, proj_axis)#, extend_cells=None)
if not os.path.isfile(synchrotron_filename(ds, extend_cells=gc)):
return ds.directory, 'sync file not found %s' % synchrotron_filename(
ds, extend_cells=gc)
maindir = os.path.join(file.pathname, 'cos_synchrotron_QU_nn_%s/' % ptype)
if proj_axis != 'x':
maindir = os.path.join(maindir, '%i_%i_%i' % tuple(proj_axis))
histdir = os.path.join(
maindir, 'histogram_gaussian%i_%i%i%i' % (sigma, *proj_axis))
else:
histdir = os.path.join(maindir,
'histogram_gaussian%i_%s' % (sigma, proj_axis))
ds_sync = yt.load(synchrotron_filename(ds, extend_cells=gc))
width = ds_sync.domain_width[1:] / zoom_fac
res = ds_sync.domain_dimensions[
1:] * ds_sync.refine_by**ds_sync.index.max_level // zoom_fac // 2
psi, frac = {}, {}
I_bin = {}
for nu in nus:
stokes = StokesFieldName(ptype, nu, proj_axis, field_type='flash')
if proj_axis == 'x':
proj = yt.ProjectionPlot(ds_sync,
proj_axis,
stokes.IQU,
center=[0, 0, 0],
width=width).set_buff_size(res)
else:
proj = yt.OffAxisProjectionPlot(
ds_sync,
proj_axis,
stokes.IQU,
width=width,
north_vector=north_vector).set_buff_size(res)
frb_I = proj.frb.data[stokes.I].v
frb_Q = proj.frb.data[stokes.Q].v
frb_U = proj.frb.data[stokes.U].v
#null = plt.hist(np.log10(arri.flatten()), range=(-15,3), bins=100)
frb_I = gaussian_filter(frb_I, sigma)
frb_Q = gaussian_filter(frb_Q, sigma)
frb_U = gaussian_filter(frb_U, sigma)
factor = 1
nx = res[0] // factor
ny = res[1] // factor
I_bin[nu] = frb_I.reshape(ny, factor, nx, factor).sum(3).sum(1)
Q_bin = frb_Q.reshape(ny, factor, nx, factor).sum(3).sum(1)
U_bin = frb_U.reshape(ny, factor, nx, factor).sum(3).sum(1)
# angle between the polarization and horizontal axis
# (or angle between the magnetic field and vertical axis
psi[nu] = 0.5 * np.arctan2(U_bin, Q_bin)
frac[nu] = np.sqrt(Q_bin**2 + U_bin**2) / I_bin[nu]
fig = plt.figure(figsize=(8, 16))
i_plot = fig.add_subplot(111)
i_plot.imshow(np.log10(frb_I + 1e-2), vmin=-1, vmax=1, origin='lower')
xx0, xx1 = i_plot.get_xlim()
yy0, yy1 = i_plot.get_ylim()
X, Y = np.meshgrid(np.linspace(xx0, xx1, nx, endpoint=True),
np.linspace(yy0, yy1, ny, endpoint=True))
mask = I_bin[nu] < 0.1
frac[nu][mask] = 0
psi[nu][mask] = 0
pixX = frac[nu] * np.cos(psi[nu]) # X-vector
pixY = frac[nu] * np.sin(psi[nu]) # Y-vector
# keyword arguments for quiverplots
quiveropts = dict(headlength=0, headwidth=1, pivot='middle')
i_plot.quiver(X, Y, pixX, pixY, scale=8, **quiveropts)
nu_str = '%.1f%s' % nu
fig.savefig(histdir + '/' + ds.basename + '_I_%s.png' % nu_str)
fig = plt.figure(figsize=(16, 4))
def plot_polarization_histogram(frac, psi, I_bin, fig=None, label=None):
if not fig:
fig = plt.figure(figsize=(16, 4))
ax1 = fig.axes[0]
null = ax1.hist(frac[I_bin.nonzero()].flatten() * 100,
range=(0, 80),
bins=40,
alpha=0.5,
weights=I_bin[I_bin.nonzero()].flatten(),
normed=True)
ax1.set_xlabel('Polarization fraction (%)')
ax1.set_xlim(0, 80)
ax2 = fig.axes[1]
null = ax2.hist(psi[I_bin.nonzero()].flatten(),
bins=50,
range=(-0.5 * np.pi, 0.5 * np.pi),
alpha=0.5,
weights=I_bin[I_bin.nonzero()].flatten(),
normed=True)
x_tick = np.linspace(-0.5, 0.5, 5, endpoint=True)
x_label = [r"$-\pi/2$", r"$-\pi/4$", r"$0$", r"$+\pi/4$", r"$+\pi/2$"]
ax2.set_xlim(-0.5 * np.pi, 0.5 * np.pi)
ax2.set_xticks(x_tick * np.pi)
ax2.set_xticklabels(x_label)
#ax2.set_title(ds.basename + ' %.1f %s' % nu)
ax3 = fig.axes[2]
null = ax3.hist(np.abs(psi[I_bin.nonzero()].flatten()),
bins=25,
range=(0.0, 0.5 * np.pi),
alpha=0.5,
label=label)
ax3.legend()
ax3.set_xlim(0.0, 0.5 * np.pi)
ax3.set_xticks([x_tick[2:] * np.pi])
ax3.set_xticks(x_tick[2:] * np.pi)
ax3.set_xticklabels(x_label[2:])
return fig
fig = plt.figure(figsize=(16, 4))
ax1 = fig.add_subplot(131)
ax2 = fig.add_subplot(132)
ax3 = fig.add_subplot(133)
for nu in nus:
nu_str = '%.1f%s' % nu
fig = plot_polarization_histogram(frac[nu],
psi[nu],
I_bin[nu],
fig=fig,
label=nu_str)
ax1.set_title(file.pathname)
ax2.set_title(ds.basename + ' %.1f %s' % nu + ' gaussian %i' % sigma)
fig.savefig(histdir + '/' + ds.basename)
return file.pathname, ds.basename[-4:]
import yt
# Load the dataset.
ds = yt.load("IsolatedGalaxy/galaxy0030/galaxy0030")
# Create a 15 kpc radius sphere, centered on the center of the sim volume
sp = ds.sphere("center", (15.0, "kpc"))
# Get the angular momentum vector for the sphere.
L = sp.quantities.angular_momentum_vector()
print(f"Angular momentum vector: {L}")
# Create an OffAxisProjectionPlot of density centered on the object with the L
# vector as its normal and a width of 25 kpc on a side
p = yt.OffAxisProjectionPlot(ds, L, "density", sp.center, (25, "kpc"))
p.save()
Zmetal = data['PartType0', 'Metallicity_00']
ZHe = data['PartType0', 'Metallicity_01']
mass = data['PartType0', 'Masses'].in_cgs()
density = data['PartType0', 'Density'].in_cgs()
protonmass_in_g_yt = ds.quan(protonmass_in_g,'g')
edensity = density/protonmass_in_g_yt*(1.0-Zmetal-ZHe)*ne
print 'edensity', edensity
return edensity
ds.add_field(('PartType0', 'electrondensity'), function=_electrondensity,
units="1/cm**3", particle_type=True)
add_volume_weighted_smoothed_field('PartType0', 'Coordinates', 'Masses',
'SmoothingLength', 'Density',
'electrondensity', ds.field_info)
if mode=='projected':
px = yt.OffAxisProjectionPlot(ds, projectionaxis, ('deposit', 'PartType0_smoothed_electrondensity'), center=center, width=new_box_size)
if mode=='Slice':
if rotface==1:
px = yt.OffAxisSlicePlot(ds, projectionaxis, ('deposit', 'PartType0_smoothed_electrondensity'), center=center, width=new_box_size,north_vector=north_vector)
else:
px = yt.SlicePlot(ds, projectionaxis, ('deposit', 'PartType0_smoothed_electrondensity'), center=center, width=new_box_size)
#px.set_buff_size(128)
px.set_buff_size(64)
# Extract X, Y, U, V from the frb
px_frb = px.frb
px_dens = np.array(px_frb[('deposit', 'PartType0_smoothed_electrondensity')])
cbcolor = 'YlGnBu'
if mode=='projected':
cblabel = r'${\rm \log (n_{\rm e} [cm^{-2}])}$'
if mode=='Slice':
cblabel = r'${\rm \log (n_{\rm e} [cm^{-3}])}$'
def plot_a_region(
ram,
out,
ds,
center,
fields='den',
kind='slc',
axis='z',
width=None,
center_vel=None,
L=None,
direcs='face',
zlims={},
l_max=None,
is_id=False,
bar_length=2e3,
sketch=False,
time_offset='tRelax',
is_set_size=True,
is_time=True,
streamplot_kwargs={},
more_kwargs={},
):
"""Do projection or slice plots in a region.
Args:
ram: a ramtools.Ramses instance
out (int): output frame
ds (yt.ds): yt.load instance
center (list_like): the center of the region like in yt.SlicePlot
fields (str or tuple): the field to plot. The avaialble options are: 'den' or 'density' - density. 'logden' - log of density. 'T' or 'temperature' - temperature. 'T_vel_rela' - temperature overplot with velocity field. 'pressure' - thermal pressure. 'magstream' - stream lines of magnetic fields on top of density slice. The magnetic strength is be indicated by the color of the stream lines. 'beta' - plasma beta parameter. 'AlfMach' - Alfvenic Mach number. 'xHII' - hydrogen ionization fraction. 'mach' - thermal Mach number. 'mag' - magnetic field strength. 'vel' - velocity field on top of density slice. 'vel_rela' - relative velocity field, i.e. the velocity field in the frame with a velocity defined by center_vel. 'mach2d' - thermal Mach number in the plane. 'p_mag' - magnetic pressure.
kind (str): 'slc' or 'prj'. Default: 'slc'
axis (str or int): One of (0, 1, 2, 'x', 'y', 'z').
width (float or tuple): width of the field of view. e.g. 0.01,
(1000, 'AU').
center_vel (list_like or None): the velocity of the center. (Default
None)
L (list_like): the line-of-sight vector. Will overwrite axis.
Default: None
direcs (str): 'face' or 'edge'. Will only be used if L is not None.
zlims (dict): The limits of the fields. e.g. {'density': [1e4, 1e8]},
{'B': [1e-5, 1e-2]} (the default of B field).
l_max (int):
is_id (bool): toggle marking sink ID (the order of creation)
bar_length (float): not using
sketch (bool): If True, will do a faster plot of a simple imshow for
test purpose. Default: False.
is_time (bool): toggle overplotting time tag
time_offset (str or int or float): One of the following cases:
str: 'rRelax'
int: the frame of which the time is used
float: time in Myr
is_set_size (bool): toggle setting figure size to 6 to make texts
bigger. Default: True.
streamplot_kwargs (dict): More kwargs for plt.streamplot that is used when fields='magstream'. Default: {}
more_kwargs (dict): more field-specific kwargs. Default: {}
Returns:
yt figure or plt.figure
"""
from matplotlib import colors
r = ram
ytfast.set_data_dir(os.path.join(r.ramses_dir, 'h5_data'))
assert width is not None
width = to_boxlen(width, r.ds1)
if isinstance(time_offset, int):
time_offset = r.get_time(time_offset)
elif isinstance(time_offset, str):
if time_offset is 'tRelax':
time_offset = r.tRelax
else:
time_offset = time_offset
if l_max is None: l_max = ds.max_level
use_h5 = False # TODO: enable use_h5
# calculated north and right, if L is not None
if L is None:
if axis in [0, 1, 2]:
axis = ['x', 'y', 'z'][axis]
is_onaxis = True
los = axis
axismap = {'x': [1, 0, 0], 'y': [0, 1, 0], 'z': [0, 0, 1]}
los_v = axismap[axis]
north = axismap[{'x': 'z', 'y': 'x', 'z': 'y'}[axis]]
logger.debug('is_onaxis = {}'.format(is_onaxis))
else:
is_onaxis = False
L /= norm(L)
tmpright = np.cross([0, 0, 1], L)
tmpright = tmpright / norm(tmpright)
if direcs == 'face':
los = L
north = -1 * tmpright
elif direcs == 'edge':
los = tmpright
north = L
los_v = los
right = np.cross(north, los_v)
logger.debug('is_onaxis = {}'.format(is_onaxis))
logger.info(f'doing {r.jobPath}, los={los}, north={north}, right={right}')
# ------------------ doing the plot --------------------
if 'density' in zlims:
den_zlim = zlims['density']
elif 'den' in zlims:
den_zlim = zlims['den']
else:
den_zlim = [None, None]
gas_field = fields
is_log = True
zlim = None
cb_label = None
cmap = 'viridis'
if fields in ['T', 'temperature', 'T_vel_rela']:
gas_field = 'temperature'
if fields in ['pressure', 'p_mag']:
cmap = 'gist_heat'
if fields in ['beta', 'magbeta', 'AlfMach']:
cmap = 'seismic'
add_beta_corrected(ds)
zlim = [-3, 3]
is_log = False
if fields == 'magbeta':
gas_field = "logPlasmaBeta"
def _logBeta(field, data):
return np.log10(data['beta'])
ds.add_field(("gas", gas_field), function=_logBeta)
cb_label = "log (Plasma Beta)"
if fields == 'AlfMach':
gas_field = ("gas", "LogMachA")
def _LogMachA(field, data):
""" MachA = v_rms / v_A = gamma / 2 * Mach^2 * beta """
gamma = 5 / 3
return np.log10(gamma / 2 * data["mach"]**2 * data["beta"])
ds.add_field(gas_field, function=_LogMachA)
cb_label = "log $M_A$"
if fields in ['xHII']:
zlim = [1e-6, 1]
if fields in ['mach', 'AlfMach']:
add_mach(ds, center_vel)
if gas_field == 'mach':
gas_field = 'logmach'
add_logmach(ds)
cb_label = r'log ($\mathcal{M}_s$)'
is_log = False
cmap = 'seismic'
zlim = [-2, 2]
if fields in ['den', 'density', 'mag', 'vel', 'vel_rela']:
gas_field = 'density'
if fields in ['logden', ('gas', 'logden')]:
def _logden(_field, data):
mu = 1.4
return data[('gas', 'density')] / (
mu * yt.physical_constants.mass_hydrogen)
gas_field = ('gas', 'logden')
ds.add_field(gas_field,
function=_logden,
take_log=True,
sampling_type='cell',
units="cm**(-3)")
if fields == 'magstream':
def _rel_B_1(_field, data):
Bx = (data["x-Bfield-left"] + data["x-Bfield-right"]) / 2.
By = (data["y-Bfield-left"] + data["y-Bfield-right"]) / 2.
Bz = (data["z-Bfield-left"] + data["z-Bfield-right"]) / 2.
B = yt.YTArray([Bx, By, Bz])
return np.tensordot(right, B, 1)
def _rel_B_2(_field, data):
Bx = (data["x-Bfield-left"] + data["x-Bfield-right"]) / 2.
By = (data["y-Bfield-left"] + data["y-Bfield-right"]) / 2.
Bz = (data["z-Bfield-left"] + data["z-Bfield-right"]) / 2.
B = yt.YTArray([Bx, By, Bz])
return np.tensordot(north, B, 1)
ds.add_field(('gas', 'rel_B_1'), function=_rel_B_1)
ds.add_field(('gas', 'rel_B_2'), function=_rel_B_2)
if fields in ['vel_rela', 'mach2d']:
assert north is not None
def _rel_vel_1(_field, data):
rel_vel = yt.YTArray([
data['velocity_x'] - yt.YTArray(center_vel[0], 'cm/s'),
data['velocity_y'] - yt.YTArray(center_vel[1], 'cm/s'),
data['velocity_z'] - yt.YTArray(center_vel[2], 'cm/s')
])
return yt.YTArray(np.tensordot(right, rel_vel, 1), 'cm/s')
def _rel_vel_2(_field, data):
rel_vel = yt.YTArray([
data['velocity_x'] - yt.YTArray(center_vel[0], 'cm/s'),
data['velocity_y'] - yt.YTArray(center_vel[1], 'cm/s'),
data['velocity_z'] - yt.YTArray(center_vel[2], 'cm/s')
])
return yt.YTArray(np.tensordot(north, rel_vel, 1), 'cm/s')
def _mach2d(_field, data):
rel_speed_sq = data['rel_vel_1'].in_cgs().value ** 2 + \
data['rel_vel_2'].in_cgs().value ** 2 # cm/s
T = data['temperature'].value
# km/s, assuming gamma=5/3. mu is inside T, therefore this
# is precise.
cs = 0.11729 * np.sqrt(T)
return 1e-5 * np.sqrt(rel_speed_sq) / cs # dimensionless
ds.add_field(('gas', 'rel_vel_1'), function=_rel_vel_1, units="cm/s")
ds.add_field(('gas', 'rel_vel_2'), function=_rel_vel_2, units="cm/s")
ds.add_field(
('gas', 'mach2d'),
function=_mach2d,
)
if fields == 'mach2d':
gas_field = 'logmach2d'
def _logmach2d(_field, data):
return np.log10(data["mach2d"])
ds.add_field(
('gas', gas_field),
function=_logmach2d,
)
cb_label = r'log ($\mathcal{M}_{s, 2D}$)'
is_log = False
cmap = "seismic"
zlim = [-2, 2]
if fields != 'magstream':
if kind == 'slc':
if is_onaxis:
p = yt.SlicePlot(ds,
los,
gas_field,
width=width,
center=center)
else:
p = yt.OffAxisSlicePlot(ds,
los,
gas_field,
width=width,
center=center,
north_vector=north)
elif kind == 'prj':
if is_onaxis:
p = ytfast.ProjectionPlot(ds,
los,
gas_field,
width=width,
center=center,
weight_field='density')
else:
print("Doing OffAxis projection plot...")
p = yt.OffAxisProjectionPlot(ds,
los,
gas_field,
width=width,
center=center,
weight_field='density',
max_level=l_max,
north_vector=north)
print("Done")
elif kind == 'colden':
if is_onaxis:
p = ytfast.ProjectionPlot(ds,
los,
gas_field,
width=width,
center=center,
weight_field=None)
else:
p = yt.OffAxisProjectionPlot(ds,
los,
gas_field,
width=width,
center=center,
weight_field=None,
max_level=l_max,
north_vector=north)
# p.set_colorbar_label(field, _label)
if is_set_size:
p.set_figure_size(6)
p.set_axes_unit('AU')
if gas_field in zlims:
if zlims[gas_field] is not None:
p.set_zlim(gas_field, *zlims[gas_field])
if is_time:
p.annotate_timestamp(
time_format='{time:.2f} {units}',
time_offset=time_offset,
text_args={
'fontsize': 8,
'color': 'w'
},
corner='upper_left',
)
# Overplot sink particles
args = {}
if 'sink_colors' in more_kwargs.keys():
args['colors'] = more_kwargs['sink_colors']
if 'mass_lims' in more_kwargs.keys():
args['lims'] = more_kwargs['mass_lims']
if 'colors' in more_kwargs.keys():
args['colors'] = more_kwargs['colors']
r.overplot_sink_with_id(p,
out,
center,
width / 2,
is_id=is_id,
zorder='mass',
withedge=1,
**args)
# set cmap, zlim, and other styles
p.set_log(gas_field, is_log)
p.set_cmap(gas_field, cmap)
if zlim is not None:
p.set_zlim(gas_field, zlim)
p.set_colorbar_label(gas_field, cb_label)
if gas_field == 'density':
den_setup(p, den_zlim, time_offset=time_offset)
if gas_field == 'temperature':
T_setup(p, [3, 1e4], time_offset=time_offset)
if fields == 'vel':
p.annotate_velocity(factor=30,
scale=scale,
scale_units='x',
plot_args={"color": "cyan"})
p.set_colorbar_label(gas_field, r'log ($\mathcal{M}_s$)')
if fields == ['vel_rela', 'T_vel_rela']:
p.annotate_quiver('rel_vel_1',
'rel_vel_2',
factor=30,
scale=scale,
scale_units='x',
plot_args={"color": "cyan"})
if fields in ['vel', 'vel_rela', 'T_vel_rela']:
# scale ruler for velocities
ruler_v_kms = 4 # km/s
width_au = width * r.boxlen * 2e5
bar_length = 2e3 * width_au / 2e4
coeff = bar_length # length of ruler in axis units [AU]
scale = 1e5 * ruler_v_kms / coeff # number of cm/s per axis units
# add a scale ruler
p.annotate_scale(coeff=coeff,
unit='AU',
scale_text_format="{} km/s".format(ruler_v_kms))
return p
else: # streamplot
if not sketch:
sl = ds.cutting(los_v, center, north_vector=north)
hw = width / 2.
bounds = [-hw, hw, -hw, hw]
size = 2**10
frb = yt.FixedResolutionBuffer(sl, bounds, [size, size])
m_h = 1.6735575e-24
mu = 1.4
den = frb['density'].value / (mu * m_h)
unitB = utilities.get_unit_B(ds)
logger.info(f"unitB.value = {unitB.value}")
Bx = frb['rel_B_1'].value * unitB.value # Gauss
By = frb['rel_B_2'].value * unitB.value # Gauss
Bmag = np.sqrt(Bx**2 + By**2)
logger.info(f"Bmag min = {Bmag.min()}, max = {Bmag.max()}")
hw_au = hw * r.unit_l / AU
grid_base = np.linspace(-hw_au, hw_au, size)
bounds_au = [-hw_au, hw_au, -hw_au, hw_au]
else:
den = np.random.random([3, 3])
hw_au = 10
den_zlim = [None, None]
bounds_au = [-hw_au, hw_au, -hw_au, hw_au]
# fig, ax = plt.subplots()
if 'figax' in more_kwargs.keys():
fig, ax = more_kwargs['figax']
else:
fig, ax = plt.subplots()
im = ax.imshow(np.log10(den),
cmap='inferno',
extent=bounds_au,
vmin=den_zlim[0],
vmax=den_zlim[1],
origin='lower')
colornorm_den = colors.Normalize(vmin=den_zlim[0], vmax=den_zlim[1])
ax.set(xlabel="Image x (AU)",
ylabel="Image y (AU)",
xlim=[-hw_au, hw_au],
ylim=[-hw_au, hw_au])
cmap = mpl.cm.get_cmap(
"Greens",
6,
)
Bvmin, Bvmax = -5, -2
if "Blims_log" in more_kwargs.keys():
Bvmin, Bvmax = more_kwargs["Blims_log"]
if "B" in zlims.keys():
Bvmin = np.log10(zlims['B'][0])
Bvmax = np.log10(zlims['B'][1])
colornorm_stream = colors.Normalize(vmin=Bvmin, vmax=Bvmax)
# streamplot_kwargs = {}
# if "streamplot" in more_kwargs.keys():
# streamplot_kwargs = more_kwargs["streamplot"]
# linewidth = 0.2
if "stream_linewidth" in more_kwargs.keys():
linewidth = more_kwargs["stream_linewidth"]
if not sketch:
logmag = np.log10(Bmag)
# scaled_mag = (logmag - Bvmin) / (Bvmax - Bvmin)
strm = ax.streamplot(
grid_base,
grid_base,
Bx,
By,
color=np.log10(Bmag),
density=1.2,
# linewidth=linewidth,
# linewidth=1.5 * scaled_mag,
# cmap='Greens',
cmap=cmap,
norm=colornorm_stream,
arrowsize=0.5,
**streamplot_kwargs,
)
# plt.subplots_adjust(right=0.8)
plot_cb = True
plot_cb2 = True
if 'plot_cb' in more_kwargs.keys():
plot_cb = more_kwargs['plot_cb']
if 'plot_cb2' in more_kwargs.keys():
plot_cb2 = more_kwargs['plot_cb2']
if plot_cb:
if 'cb_axis' in more_kwargs.keys():
ax2 = more_kwargs['cb_axis']
# pos0 = ax.get_position()
# cbaxis = fig.add_axes([pos0.x1, pos0.y0, 0.06, pos0.height])
# cb = plt.colorbar(im, cax=cbaxis)
cb2 = mpl.colorbar.ColorbarBase(
ax2,
orientation='vertical',
cmap=mpl.cm.get_cmap("inferno"),
norm=colornorm_den,
)
else:
cb2 = fig.colorbar(im, ax=ax, pad=0)
cb2.set_label("log $n$ (cm$^{-3}$)")
if plot_cb2:
if 'cb2_axis' in more_kwargs.keys():
ax2 = more_kwargs['cb2_axis']
elif plot_cb:
pos1 = ax.get_position()
ax2 = fig.add_axes([pos1.x1 + .15, pos1.y0, 0.02, pos1.height])
else:
pos1 = ax.get_position()
ax2 = fig.add_axes([pos1.x0, pos1.y0, 0.02, pos1.height])
colornorm2 = colors.Normalize(vmin=den_zlim[0], vmax=den_zlim[1])
cb2 = mpl.colorbar.ColorbarBase(
ax2,
orientation='vertical',
cmap=cmap,
norm=colornorm_stream,
# ticks=[-4, -3, -2]
)
cb2.set_label("log B (Gauss)")
# plt.savefig(fs['magstream'], dpi=300)
return fig
def visualisation(viz_type,
container,
raw_snapshot,
module=config.default_module,
gas=True,
stars=False,
dark=False,
gas_fields=["temperature", "density"],
gas_units=["K", "g/cm**3"],
gas_extrema=[[None, None], [None, None]],
dark_fields=[('deposit', 'dark_density')],
dark_units=["g/cm**3"],
dark_extrema=[[None, None]],
star_fields=[('deposit', 'stars_density')],
star_units=["g/cm**3"],
star_extrema=[[None, None]],
filter=None,
return_objects=False,
return_fields=False,
callbacks=[],
width=width_thing,
extra_width=None,
depth=depth_thing,
name="plot",
format=".png",
prefix="",
normal_vector=[1.0, 0.0, 0.0],
axis=[0, 1, 2],
plot_images=True,
weight_field=None,
image_width=1000,
extra_image_width=None,
suffix=None):
"""
This routine is designed to handle almost all of the visualisation routines that you can think of.
this has a disgusting amount of kwargs, so lets break everything down
arguments
viz_type = slice, projection, off_axis_slice, off_axis_projection. Not available for all modules all these combinations. But it will know what you mean depending on the config that you load
container = the data container or object of interest. This will accept either a YT raw snapshot or a pymses raw snapshot
module = yt or pymses
kwargs
raw_snapshot = the raw snapshot of the simulation (i.e the whole box or a larger portion of the box). This may be used to make an additional cut for the datasource
gas = do we want to plot gas data, default to True
stars = do we want to plot stellar data, default to false
dark = do we want to plot dark matter data, default to false
gas_fields = a list of strings (or turples for some YT magic) of fields that we want to plot
gas_units = the units of those fields
dark_fields = see above
dark_units = see above
star_fields = see above
star_units = see above
filter = if there needs to be any additonal filters used... but really that should be handled before this, but this is more specific for YT
return_objects = whether you wish to return a list of plot objects and images or not for further work
callbacks = extra callbacks? this is YT specific
width = YTArray containing the width of the image and its unit
extra_width = YTArray if you want to do something a bit more fancy with the image dimensions
depth = projection depth if applicable
name = name of the plot if you want to rename it
format = the output format
normal_vector = the normal vector for the off axis plots
axis = axis for on axis plots
data_source = YT data source if applicable
plot_images = if you actually want to plot something
TODO use regex to rename dark_deposit with io_deposit if age does not exist
"""
# customise the width etc
# if not [("all","particle_age")] in container:
# print "WARNING, DM only run for particles. Viz may fail since the " + \
# "age field does not exist in DM only runs"
axis_unit = width.units
if extra_width != None:
width = ((width.v, width.units), (extra_width.v, width.units))
else:
width = ((width.v, width.units), (width.v, width.units))
if extra_image_width != None:
image_width = (image_width, extra_image_width)
else:
image_width = (image_width, image_width)
depth = (depth.v, depth.units)
fields = []
units = []
extrema = []
if gas_fields and gas == True:
fields = fields + gas_fields
units = units + gas_units
extrema = extrema + gas_extrema
if star_fields and stars == True:
fields = fields + star_fields
units = units + star_units
extrema = extrema + star_extrema
if dark_fields and dark == True:
fields = fields + dark_fields
units = units + dark_units
extrema = extrema + dark_extrema
# correct length of the extrema
# usually if this is the case, then we obviously don't care
if len(fields) > len(extrema):
diff = len(fields) - len(extrema)
for i in range(0, diff):
extrema = extrema + [[None, None]]
for i in range(0, len(extrema)):
if extrema[i][0] == None:
extrema[i][0] = "min"
if extrema[i][1] == None:
extrema[i][1] = "max"
print fields
plots = {}
frb = {}
basis_vectors = ortho_find(normal_vector)
north_vectors = ortho_find(normal_vector)
print "basis_vectors", basis_vectors
# check if basis vectors are in right units
basis_vectors = YTArray(basis_vectors, "g*kpc**2/s")
north_vectors = YTArray(north_vectors, "g*kpc**2/s")
print basis_vectors
# select the module
if module == "yt":
# select the type of plot
if viz_type == "slice":
# axis
print "on axis slice plot"
if 0 in axis:
# x
# to note, a slice plot is pretty irrelevent of the data source.
# see e.g https://bpaste.net/show/f08dda811ae0
# which will reproduce the same result
print "plotting axis 0"
print width[0][0], "width"
plot = yt.SlicePlot(container.ds,
0,
fields,
center=container.center,
width=width)
image = plot.data_source.to_frb(
yt.YTQuantity(width[0][0], axis_unit), image_width[0])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plots["0_plot"] = plot
frb["0_frb"] = image
if plot_images:
plot_viz(plot, name, suffix)
# y
if 1 in axis:
# x
print "plotting axis 1"
plot = yt.SlicePlot(container.ds,
1,
fields,
center=container.center,
width=width)
image = plot.data_source.to_frb(
yt.YTQuantity(width[0][0], axis_unit),
[image_width[0], image_width[1]])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plots["1_plot"] = plot
frb["1_frb"] = image
if plot_images:
plot_viz(plot, name, suffix)
# z
if 2 in axis:
# x
print "plotting axis 2"
plot = yt.SlicePlot(container.ds,
2,
fields,
center=container.center,
width=width)
image = plot.data_source.to_frb(
yt.YTQuantity(width[0][0], axis_unit),
[image_width[0], image_width[1]])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plots["2_plot"] = plot
frb["2_frb"] = image
if plot_images:
plot_viz(plot, name, suffix)
if viz_type == "projection":
# axis
print "on axis projection plot"
if 0 in axis:
# x
# in this instance.. the depth is along the axis that you are viewing "down".. so in a way is user defined
data_source = gen_data_source(0, container, raw_snapshot,
width, depth, axis_unit)
print "plotting axis 0"
plot = yt.ProjectionPlot(container.ds,
0,
fields,
center=container.center,
width=width,
weight_field=weight_field,
data_source=data_source)
image = plot.data_source.to_frb(
yt.YTQuantity(width[0][0], axis_unit),
[image_width[0], image_width[1]])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plots["0_plot"] = plot
frb["0_frb"] = image
if plot_images:
plot_viz(plot, name, suffix)
# y
if 1 in axis:
print "plotting axis 1"
data_source = gen_data_source(1, container, raw_snapshot,
width, depth, axis_unit)
plot = yt.ProjectionPlot(container.ds,
1,
fields,
center=container.center,
width=width,
weight_field=weight_field,
data_source=data_source)
image = plot.data_source.to_frb(
yt.YTQuantity(width[0][0], axis_unit),
[image_width[0], image_width[1]])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plots["1_plot"] = plot
frb["1_frb"] = image
if plot_images:
plot_viz(plot, name, suffix)
# z
if 2 in axis:
print "plotting axis 2"
data_source = gen_data_source(2, container, raw_snapshot,
width, depth, axis_unit)
plot = yt.ProjectionPlot(container.ds,
2,
fields,
center=container.center,
width=width,
weight_field=weight_field,
data_source=data_source)
image = plot.data_source.to_frb(
yt.YTQuantity(width[0][0], axis_unit),
[image_width[0], image_width[1]])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plots["2_plot"] = plot
frb["2_frb"] = image
if plot_images:
plot_viz(plot, name, suffix)
if viz_type == "off_axis_slice":
# off axis.. may as well do all 3
# same rules RE the slice apply here too
print "off axis slice plot"
if 0 in axis:
print "plotting axis 0"
plot = yt.OffAxisSlicePlot(container.ds,
basis_vectors[0],
fields,
center=container.center,
width=width,
north_vector=north_vectors[0])
image = plot.data_source.to_frb(
yt.YTQuantity(width[0][0], axis_unit),
[image_width[0], image_width[1]])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plots["0_plot"] = plot
frb["0_frb"] = image
if plot_images:
plot_viz(plot, name + "_axis_0", suffix)
if 1 in axis:
print "plotting axis 1"
plot = yt.OffAxisSlicePlot(container.ds,
basis_vectors[1],
fields,
center=container.center,
width=width,
north_vector=north_vectors[0])
image = plot.data_source.to_frb(
yt.YTQuantity(width[0][0], axis_unit),
[image_width[0], image_width[1]])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plots["1_plot"] = plot
frb["1_frb"] = image
if plot_images:
plot_viz(plot, name + "_axis_1", suffix)
if 2 in axis:
print "plotting axis 2"
plot = yt.OffAxisSlicePlot(container.ds,
basis_vectors[2],
fields,
center=container.center,
width=width,
north_vector=north_vectors[0])
image = plot.data_source.to_frb(
yt.YTQuantity(width[0][0], axis_unit),
[image_width[0], image_width[1]])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plots["2_plot"] = plot
frb["2_frb"] = image
if plot_images:
plot_viz(plot, name + "_axis_2", suffix)
if viz_type == "off_axis_projection":
print "off axis projection plot"
if 0 in axis:
print "plotting axis 0"
# since there is no way of doing an off axis box.. our data source is essentially a cube.. it ignore depth
# oh, and you probably want this rather than your disk dataset
#TODO make sure filtered datasts carry across... but it seems to be the case here.
data_source = gen_data_source(0, container, raw_snapshot,
width, width[0], axis_unit)
plot = yt.OffAxisProjectionPlot(data_source.ds,
basis_vectors[0],
fields,
center=container.center,
width=width,
depth=width[0],
weight_field=weight_field,
north_vector=north_vectors[0])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
plot.annotate_marker((4.0, 6.92820),
coord_system='plot',
plot_args={
'color': 'black',
's': 400
})
image = plot.frb
plots["0_plot"] = plot
frb["0_frb"] = image
if plot_images:
plot_viz(plot, name + "_axis_0", suffix)
if 1 in axis:
print "plotting axis 1"
data_source = gen_data_source(1, container, raw_snapshot,
width, width[0], axis_unit)
plot = yt.OffAxisProjectionPlot(data_source.ds,
basis_vectors[1],
fields,
center=container.center,
width=width,
depth=width[0],
weight_field=weight_field,
north_vector=north_vectors[0])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
image = plot.frb
plots["1_plot"] = plot
frb["1_frb"] = image
if plot_images:
plot_viz(plot, name + "_axis_1", suffix)
if 2 in axis:
print "plotting axis 2"
data_source = gen_data_source(2, container, raw_snapshot,
width, width[0], axis_unit)
plot = yt.OffAxisProjectionPlot(data_source.ds,
basis_vectors[2],
fields,
center=container.center,
width=width,
depth=width[0],
weight_field=weight_field,
north_vector=north_vectors[0])
# set the units
plot.set_axes_unit("kpc")
for i in range(0, len(fields)):
plot.set_unit(fields[i], units[i])
plot.set_zlim(fields[i], extrema[i][0], extrema[i][1])
image = plot.frb
plots["2_plot"] = plot
frb["2_frb"] = image
if plot_images:
plot_viz(plot, name + "_axis_2", suffix)
else:
print "module not defined, please either use YT, pynbody or pymses"
return
if module == "pymses":
from pymses.analysis.visualization import *
from ramses_pp.modules.pymses import PymsesProjection
if return_objects == True and return_fields == True:
print plots, frb, fields
return plots, frb, fields
if return_objects == True:
return plots, frb, None
if return_fields == True:
return None, None, fields
