import yt
from yt.visualization.base_plot_types import get_multi_plot
fn = "GasSloshing/sloshing_nomag2_hdf5_plt_cnt_0150" # dataset to load
orient = "horizontal"
ds = yt.load(fn) # load data
# There's a lot in here:
# From this we get a containing figure, a list-of-lists of axes into which we
# can place plots, and some axes that we'll put colorbars.
# We feed it:
# Number of plots on the x-axis, number of plots on the y-axis, and how we
# want our colorbars oriented. (This governs where they will go, too.
# bw is the base-width in inches, but 4 is about right for most cases.
fig, axes, colorbars = get_multi_plot(3, 2, colorbar=orient, bw=4)
slc = yt.SlicePlot(ds,
"z",
fields=["density", "temperature", "velocity_magnitude"])
proj = yt.ProjectionPlot(ds, "z", "density", weight_field="density")
slc_frb = slc.data_source.to_frb((1.0, "Mpc"), 512)
proj_frb = proj.data_source.to_frb((1.0, "Mpc"), 512)
dens_axes = [axes[0][0], axes[1][0]]
temp_axes = [axes[0][1], axes[1][1]]
vels_axes = [axes[0][2], axes[1][2]]
for dax, tax, vax in zip(dens_axes, temp_axes, vels_axes):
import matplotlib.colorbar as cb
from matplotlib.colors import LogNorm
fn = "GasSloshing/sloshing_nomag2_hdf5_plt_cnt_0150" # dataset to load
orient = 'horizontal'
ds = yt.load(fn) # load data
# There's a lot in here:
# From this we get a containing figure, a list-of-lists of axes into which we
# can place plots, and some axes that we'll put colorbars.
# We feed it:
# Number of plots on the x-axis, number of plots on the y-axis, and how we
# want our colorbars oriented. (This governs where they will go, too.
# bw is the base-width in inches, but 4 is about right for most cases.
fig, axes, colorbars = get_multi_plot(3, 2, colorbar=orient, bw = 4)
slc = yt.SlicePlot(ds, 'z', fields=["density","temperature","velocity_magnitude"])
proj = yt.ProjectionPlot(ds, 'z', "density", weight_field="density")
slc_frb = slc.data_source.to_frb((1.0, "Mpc"), 512)
proj_frb = proj.data_source.to_frb((1.0, "Mpc"), 512)
dens_axes = [axes[0][0], axes[1][0]]
temp_axes = [axes[0][1], axes[1][1]]
vels_axes = [axes[0][2], axes[1][2]]
for dax, tax, vax in zip(dens_axes, temp_axes, vels_axes) :
dax.xaxis.set_visible(False)
dax.yaxis.set_visible(False)
def four_panels(kwargs,\
sink_id, bin_vals, ds=None,width=None,dxmin=None, init_sink_file=None, d_params=None,\
name='radial_velocity.png'):
'''kwargs -- dict with all params needed to choose between hydro vs magn plots'''
if ds is None or dxmin is None or bin_vals is None or width is None or d_params is None: raise ValueError
# Choose model
if kwargs['which'] == 'hydro':
fields=[('gas','my_radial_velocity')]
ylab= r'$\mathbf{ V_r/c_s }$'
cbar_lab= r'$\mathbf{ V_r/c_s }$'
elif kwargs['which'] == 'magn':
fields=[('gas','my_radial_velocity'),('gas','my_btot'),'density']
ylab= r'$\mathbf{ V_r/V_{\rm{a,\, F}} }$'
cbar_lab= r'$\mathbf{ V_r/V_{\rm{a,\, F}} }$'
else: raise ValueError
# Initialize plot
fig, axes, colorbars = get_multi_plot(4, 1, colorbar='horizontal', bw = 4,cbar_padding=0.5)
axes= axes[0]
# Radial profile axis[0]
axes[0].plot(bin_vals['binc']/d_params['rABH'], bin_vals['q50']/d_params['cs'],c='b',ls='-',lw=2, label='q50')
axes[0].fill_between(bin_vals['binc']/d_params['rABH'],bin_vals['q25']/d_params['cs'],bin_vals['q75']/d_params['cs'],\
color='b',alpha=0.25,label='q25, q75')
# Line for 4-12dx
#for ndx in [4,6,8,10,12]:
# axes[0].plot([ndx*dxmin/d_params['rBH']]*2,\
# [bin_vals['q25'].min()/d_params['cs'],bin_vals['q75'].max()/d_params['cs']],\
# c='k',ls='--',lw=1,label=r'%d $\mathbf{ \Delta \, x_{\rm{min}} }$' % ndx)
axes[0].plot(list(axes[0].get_xlim()), [-1]*2,\
c='k',ls=':',lw=2)
laba=dict(fontweight='bold',fontsize='x-large')
xlab= axes[0].set_xlabel(r'$\mathbf{ r [r_{\rm{ABH}}] (\rm{Bin Centers})}$',**laba)
ylab= axes[0].set_ylabel(ylab,**laba)
axes[0].legend(loc='lower right')
# Remove far right tick
xticks= np.linspace(axes[0].get_xlim()[0],axes[0].get_xlim()[1],num=7)[:-1]
axes[0].set_xticks(xticks)
# Set ylim max to 0
axes[0].set_ylim(axes[0].get_ylim()[0],0.)
# Imshow panels
sink= InitSink(init_sink_file)
sink.set_coord(sink_id)
center= sink.xyz_for_sid(sink_id)
slices={}
frbs={}
np_arr={}
xmin,xmax= 0,512 #fixed res buffer
for xyz in ['x','y','z']:
slices[xyz] = yt.SlicePlot(ds, xyz, center=center,fields=fields)
frbs[xyz] = slices[xyz].data_source.to_frb(width, xmax)
# Converting our Fixed Resolution Buffers to numpy arrays
if kwargs['which'] == 'hydro':
np_arr[xyz]= np.array(frbs[xyz][('gas','my_radial_velocity')])
else:
vr= np.array(frbs[xyz][('gas','my_radial_velocity')])
btot= np.array(frbs[xyz][('gas','my_btot')])
dens= np.array(frbs[xyz]['density'])
np_arr[xyz]= vr/get_vafast(btot=btot,dens=dens)
plots={}
labs=dict(x=['y','z'],y=['z','x'],z=['x','y']) #x,y labels from yt when slice
for i,xyz in zip(range(1,4),['x','y','z']):
plots[xyz]= axes[i].imshow(np_arr[xyz], origin='lower') #norm=LogNorm()
#axes[i].contour(np_arr[xyz], [0], hold='on', origin='lower',\
# linewidths=2.,linestyles='solid',colors='w')
#axes[i].contour(np_arr[xyz], [1], hold='on', origin='lower',\
# linewidths=1.,linestyles='dashed',colors='w')
# Set ticklabels about 0 and in units rBH
xticks= np.linspace(xmin,xmax,num=5)
axes[i].set_xticks(xticks) # same 0 to 512Left and rightmost ticks in imshow
axes[i].set_yticks(xticks)
labels= (xticks-xticks[2])/xmax # +/-0.5 from center in plot frame units
labels= (labels*width[0]/d_params['rABH']).astype(str)
if i == 1 or i == 2: labels[-1]= '' # Prevent overlap
axes[i].set_xticklabels(labels)
axes[i].set_yticklabels(labels)
# continue
axes[i].tick_params(direction='out')
ti= axes[i].set_title('%s Slice' % xyz.upper(),fontsize='x-large')
# Confirmed by makign individual slices and seeing xy labs
xlab= axes[i].set_xlabel('%s [rABH]' % labs[xyz][0].upper())
ylab= axes[i].set_ylabel('%s [rABH]' % labs[xyz][1].upper())
# only show some
axes[i].yaxis.set_visible(False)
#axes[i].yaxis.set_visible(False)
for xyz in ['x','y','z']:
plots[xyz].set_clim((np_arr[xyz].min(),np_arr[xyz].max()))
plots[xyz].set_cmap("viridis")
titles=[cbar_lab]*3
for p, cax, t in zip([plots['x'],plots['y'],plots['z']], colorbars[1:], titles):
cbar = fig.colorbar(p, cax=cax, orientation='horizontal')
cbar.set_label(t)
fig.savefig(name, bbox_extra_artists=[ti,xlab,ylab], bbox_inches='tight')
plt.close()
zlims = (7e12, 1e16)
# load in simulation data and add ion fields
frb = h5.File(
'/nobackup/ibutsky/data/YalePaper/multipanel_romulusC_%i_sightline_plot_data_fixed'
% (output), 'r')
orient = 'horizontal'
if len(ray_list) <= 4:
nrows = 1
ncols = len(ray_list)
else:
nrows = 2
ncols = int(len(ray_list) / nrows)
fig, axes, colorbars = get_multi_plot(ncols, nrows, colorbar=None, bw=4)
for i, ray_id in enumerate(ray_list):
row = int(i / ncols)
col = i - ncols * row
print(row, col)
img_data = np.array(frb['ray_%i_%s' % (ray_id, field[1])])
# img_data = ipd.crop_imshow(img_data, data_min, data_max, data_min, data_max)
ax = axes[row][col]
# note: the extent call makes the units of the plot 0 to 1200 (physical) instead of 0 to 1600 (number of pixels)
print(img_data.min(), img_data.max())
im = ax.imshow(img_data, origin = 'lower', norm = LogNorm(),vmin = zlims[0], vmax = zlims[1], \
extent = [rmin, rmax, rmin, rmax], zorder = 1)
ax.scatter(600, 600, marker='+', s=200, c='white', zorder=10)
im.set_cmap(cmap)
zlim_list = [(1e14, 1e20), (1e4, 1e6)]
cbar_title_list = [
r'$\mathrm{H\ I\ Column\ Density}\ (\mathrm{cm^{-2}})$',
r'$\mathrm{Temperature\ (K)}$'
]
# load pre-generated plot data
frb = h5.File(
'/nobackup/ibutsky/data/YalePaper/multipanel_romulusC_%i_plot_data' %
(output), 'r')
for field, cmap, zlim, cbar_title in zip(field_list, cmap_list, zlim_list,
cbar_title_list):
# This uses yt to generate a fig/axes object. You can also use plt.subplots (see commented out line below)
fig, axes, colorbars = get_multi_plot(1, 1, colorbar=None, bw=4)
ax = axes[0][0]
# fig, ax = plt.subplots(nrows = 1, ncols = 1, figsize=(7, 7))
dset = field[1]
img_data = np.array(frb[dset])
im = ax.imshow(img_data, origin = 'lower', norm = LogNorm(),\
cmap = cmap, vmin = zlim[0], vmax = zlim[1])
cbax = inset_axes(ax, width="90%", height="3%", loc=9)
cbar = fig.colorbar(im, cax=cbax, orientation='horizontal')
cbar.set_label(cbar_title, color='white')
fig.tight_layout()
fig.savefig("romulusC_%i_%s.png" % (output, field[1]), dpi=300)
