def process_a_sink(sink_id, ds=None,radius=None,dxmin=None,init_sink_file=None,model=None,outdir=None):
'''radius -- tuple of (float,"cm") where fload is number cm corresponding some # rBH
dxmin -- dx level 5 in simulation units so "cm"
Returns: dict of q25-75 and bin centers for radius and radial velocity relative location of sink in simulation units "cm"'''
assert(ds is not None and radius is not None and model is not None)
# Data within a sphere of radius ~ rBH, centered on sink
sink= InitSink(init_sink_file)
sink.set_coord(sink_id)
center= sink.xyz_for_sid(sink_id)
print('sink_id=',sink_id,'center= ',center)
sp = ds.sphere(center, radius)
# Compute my_radius relative to sink
sp.set_field_parameter("center", yt.YTArray(center, "cm"))
# Hydro calculations (for all models)
# Compute my_radial_velocity subtracting correct Bulk vel
# Bulk vel is avg Vxyz between radius= (6dxmin, outer radius]
r= np.array( sp[('gas','my_radius')] )
b_keep= r > 6.*dxmin
bulk_v=np.zeros(3)
for i,comp in enumerate(['x','y','z']):
bulk_v[i]= np.average( np.array(sp['my_v'+comp])[b_keep] )
sp.set_field_parameter("bulk_velocity", yt.YTArray(bulk_v, "cm/s"))
# Bin up Vr/cs and Vr/Vafast if later exists
bin_vals={}
# Bin radius, get 25-75 quartiles of my_radial_vel in each bin
rad= np.array( sp[('gas','my_radius')] )
vr= np.array( sp[('gas','my_radial_velocity')] )
bin_vals['hydro']= bin_up(rad,vr, bin_minmax=(6*dxmin,rad.max()),nbins=100)
if model != 'bInf':
rad= np.array( sp[('gas','my_radius')] )
vr= np.array( sp[('gas','my_radial_velocity')] )
btot= np.array( sp[('gas','my_btot')] )
dens= np.array( sp['density'] )
bin_vals['magn']= bin_up(rad,vr/get_vafast(btot=btot,dens=dens), \
bin_minmax=(6*dxmin,rad.max()),nbins=100)
return bin_vals
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()
