def pretty_pic(data_hdf5,args):
ds = yt.load(data_hdf5) # load data
ds.add_field("KaysVx", function=_KaysVx, units="cm/s")
ds.add_field("KaysVy", function=_KaysVy, units="cm/s")
ds.add_field("KaysVz", function=_KaysVz, units="cm/s")
ds.add_field("KaysBx", function=_KaysBx, units="gauss")
ds.add_field("KaysBy", function=_KaysBy, units="gauss")
ds.add_field("KaysBz", function=_KaysBz, units="gauss")
msink=13./32
cs=1.
mach=5.
rho0=1.e-2
rBH= msink/cs**2/(mach**2+1)
sink= InitSink(args.data_sink)
sink.coord= dict(x=0.,y=0.,z=0.) #override it
center = [0,0,0]
width = (float(ds.domain_width[0]), 'cm')
Npx= 1000
res = [Npx, Npx] # create an image with 1000x1000 pixels
fig,ax= plt.subplots() #sharey=True,sharex=True)
#fig.set_size_inches(15, 10)
cut_dir='z'
if args.twod == 'slice':
proj = ds.slice(axis=cut_dir,coord=sink.coord[cut_dir])
else:
proj = ds.proj(field="density",axis=cut_dir)
frb = proj.to_frb(width=width, resolution=res, center=center)
img= np.array(frb['density'])[::-1]/rho0
xlab,ylab='x/rBH','y/rBH'
im= ax.imshow(np.log10(img),vmin=np.log10(args.clim[0]),vmax=np.log10(args.clim[1]))
ax.autoscale(False)
#sink marker
ax.scatter((sink.x+1)*Npx/2,(sink.y+1)*Npx/2,s=50,c='black',marker='o') #all sinks
#user defined ticks
xticks=(Npx*np.array([0.,0.25,0.5,0.75,1.])).astype('int')
# xticks_labs= (xticks-500)*rBH/Npx
ax.set_xticks(xticks)
ax.set_xticklabels(((xticks-Npx/2)*float(args.width_rBH)/Npx).astype('int')) #ax.xaxis.get_ticklocs()/100.)
ax.set_yticks(xticks)
ax.set_yticklabels(((xticks[::-1]-Npx/2)*float(args.width_rBH)/Npx).astype('int')) #ax.yaxis.get_ticklocs()/100.)
ax.set_xlabel(xlab)
ax.set_ylabel(ylab)
ax.set_title(cut_dir+' '+args.twod)
cax = fig.add_axes([0.83, 0.15, 0.02, 0.7])
cbar= fig.colorbar(im, cax=cax)
# cbar = fig.colorbar(cax) # ticks=[-1,0,1])
cbar_labels = [item.get_text() for item in cbar.ax.get_yticklabels()]
cbar_labels = ['']*len(cbar_labels)
cbar_labels[0]= args.clim[0]
cbar_labels[-1]= args.clim[1]
cbar.ax.set_yticklabels(cbar_labels)
if args.twod == 'slice': lab= r'$\mathbf{ \rho/\bar{\rho} }$'
else: lab= r'$\mathbf{ \Sigma/\bar{\Sigma} }$'
cbar.set_label(r'$\mathbf{ \rho/\bar{\rho} }$',fontsize='xx-large')
add_name= 'prettypic'
plt.savefig(os.path.join(args.outdir,args.twod+'_'+add_name+'_'+os.path.basename(data_hdf5)+'_.png'))
def get_image(data_hdf5,args):
ds = yt.load(data_hdf5) # load data
ds.add_field("TestDensity", function=_TestDensity, units="g/cm**3")
ds.add_field("KaysVx", function=_KaysVx, units="cm/s")
ds.add_field("KaysVy", function=_KaysVy, units="cm/s")
ds.add_field("KaysVz", function=_KaysVz, units="cm/s")
ds.add_field("KaysBx", function=_KaysBx, units="gauss")
ds.add_field("KaysBy", function=_KaysBy, units="gauss")
ds.add_field("KaysBz", function=_KaysBz, units="gauss")
msink=13./32
cs=1.
mach=5.
rho0=1.e-2
rBH= msink/cs**2/(mach**2+1)
tBH= msink/(mach**2+cs**2)**(3./2)
sink= InitSink(args.data_sink)
if args.sink_id:
sink.set_coord(args.sink_id)
center= sink.xyz_for_sid(args.sink_id)
width= (args.width_rBH*rBH,'cm')
else:
sink.coord= dict(x=0.,y=0.,z=0.) #override it
center = [0,0,0]
width = (float(ds.domain_width[0]), 'cm')
Npx= 1000
res = [Npx, Npx] # create an image with 1000x1000 pixels
cmaps=dict(inferno=cm.inferno,viridis=cm.viridis,default=None)
if args.test_panel: #make quick plot and quit
if args.twod == 'projection': p = yt.ProjectionPlot(ds, 'x', 'TestDensity')
else: p = yt.SlicePlot(ds, 'x', 'TestDensity')
p.set_cmap(field="TestDensity", cmap=cmaps[args.cmap])
p.save(os.path.join(args.outdir,'test_'+args.twod+'_'+os.path.basename(data_hdf5)+'_.png'))
sys.exit(0)
fig,axes= plt.subplots(1,3) #sharey=True,sharex=True)
fig.subplots_adjust(bottom=0.25,wspace=0.2) #,hspace=-0.335) #,hspace=0.05)
fig.set_size_inches(20, 7)
ax=axes.flatten()
for i,cut_dir in zip(range(3),['z','x','y']):
if args.twod == 'slice':
proj = ds.slice(axis=cut_dir,coord=sink.coord[cut_dir])
else:
proj = ds.proj(field="density",axis=cut_dir)
frb = proj.to_frb(width=width, resolution=res, center=center)
img= np.array(frb['density'])[::-1]/rho0
if cut_dir == 'x':
Vhoriz,Vvert= np.array(frb['KaysVy'])[::-1], np.array(frb['KaysVz'])[::-1]
Bhoriz,Bvert= np.array(frb['KaysBy'])[::-1], np.array(frb['KaysBz'])[::-1]
xlab,ylab='y [rBH]','z [rBH]'
elif cut_dir == 'y':
Vhoriz,Vvert= np.array(frb['KaysVx'])[::-1], np.array(frb['KaysVz'])[::-1]
Bhoriz,Bvert= np.array(frb['KaysBx'])[::-1], np.array(frb['KaysBz'])[::-1]
xlab,ylab='x [rBH]','z [rBH]'
else:
Vhoriz,Vvert= np.array(frb['KaysVx'])[::-1], np.array(frb['KaysVy'])[::-1]
Bhoriz,Bvert= np.array(frb['KaysBx'])[::-1], np.array(frb['KaysBy'])[::-1]
xlab,ylab='x [rBH]','y [rBH]'
if args.clim: im= ax[i].imshow(np.log10(img),cmap=cmaps[args.cmap],\
vmin=np.log10(args.clim[0]),vmax=np.log10(args.clim[1]))
else: im= ax[i].imshow(np.log10(img),cmap=cmaps[args.cmap])
ax[i].autoscale(False)
#sink marker
if args.sink_id: ax[i].scatter(Npx/2,Npx/2,s=50,c='white',marker='o') #put sink at center
else: ax[i].scatter((sink.x+1)*Npx/2,(sink.y+1)*Npx/2,s=50,c='white',marker='o') #all sinks
#velocity arrows
X, Y = np.meshgrid(np.arange(0,Npx), np.arange(0,Npx))
skip= int(Npx/10.)
Q = ax[i].quiver(X[::skip, ::skip], Y[::skip, ::skip], Vhoriz[::skip, ::skip], Vvert[::skip, ::skip],
pivot='mid', color='white', units='inches')
#bfield integrateed streamlines
if args.nf:
xmin,xmax= 0.,float(Npx-1)
dx=[1.,1.,1.]
ft_from_center= Npx*args.ft_from_center/args.width_rBH #height of hline of footpoints from top of image
points,xf,yf= my_streamlines(Bhoriz,Bvert, xmax,xmin,dx,args.nf,ft_from_center,\
ft_circle=True)
#ax[i].plot(xf,yf,marker='o',color='white',linestyle='None')
for j in range(0,len(points)):
xdata = (points[j])[:,0]
ydata = (points[j])[:,1]
ind = np.where(xdata>xmin)
xdata=xdata[ind]; ydata=ydata[ind]
ind = np.where(xdata<xmax)
xdata=xdata[ind]; ydata=ydata[ind]
ind = np.where(ydata>xmin)
xdata=xdata[ind]; ydata=ydata[ind]
ind = np.where(ydata<xmax)
xdata=xdata[ind]; ydata=ydata[ind]
print "N pts in streamline= ",len(xdata)
ax[i].plot(xdata,ydata,color='#33CC33',lw=2)
#finish labeling axes
xticks=(Npx*np.array([0.,0.25,0.5,0.75,1.])) #.astype('int')
ax[i].set_xticks(xticks)
ax[i].set_yticks(xticks)
ax[i].set_yticklabels(((xticks[::-1]-Npx/2.)*float(args.width_rBH)/Npx)) #.astype('int'))
ax[i].set_xticklabels(((xticks-Npx/2)*float(args.width_rBH)/Npx)) #.astype('int'))
ax[i].get_xaxis().set_tick_params(direction='out')
ax[i].get_yaxis().set_tick_params(direction='out')
ax[i].set_ylabel(ylab)
ax[i].set_xlabel(xlab)
ax[i].set_title(cut_dir+' '+args.twod)
time= (float(ds.current_time)-0.8)/tBH #0.8 is simulation time when insert sinks for all beta runs
ax[0].text(0.1,1.1,"Time: %.1f tBH" % time,transform=ax[0].transAxes,ha='center',va='center',\
fontsize='xx-large')
#cax = fig.add_axes([0.9, 0.15, 0.02, 0.7])
cax = fig.add_axes([0.25,0.1, 0.5, 0.05])
cbar= fig.colorbar(im, cax=cax,orientation='horizontal')
# cbar = fig.colorbar(cax) # ticks=[-1,0,1])
if args.clim:
#cbar_labels = [item.get_text() for item in cbar.ax.get_xticklabels()]
#cbar_labels = ['']*len(cbar_labels)
#print "cbar_labels= ",cbar_labels,"its len= ",len(cbar_labels)
#cbar_labels[0]= args.clim[0]
#cbar_labels[-1]= args.clim[1]
#myticks= np.array([args.clim[0],1e-1,1e0,1e1,1e2,args.clim[1]])
myticks= np.logspace(np.log10(args.clim[0]),np.log10(args.clim[1]),6)
cbar.set_ticks(np.log10(myticks))
cbar.ax.set_xticklabels(myticks) #[args.clim[0],1e-1,1e0,1e1,1e2,args.clim[1]])
if args.twod == 'slice': lab= r'$\mathbf{ \rho/\bar{\rho} }$'
else: lab= r'$\mathbf{ \Sigma/\bar{\Sigma} }$'
cbar.set_label(r'$\mathbf{ \rho/\bar{\rho} }$',fontsize='xx-large')
if args.sink_id: add_name= 'sid'+str(args.sink_id)
else: add_name= 'fullbox'
plt.savefig(os.path.join(args.outdir,args.twod+'_'+add_name+'_'+os.path.basename(data_hdf5)+'_.png'))