Exemplo n.º 1
0
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'))