示例#1
0
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 
示例#2
0
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()