コード例 #1
0
def subcloud_entropy_curve(data, lonin=[-1., 361.], rot_fac=1.):

    # declare constants
    cp = 287.04 / 2 * 7
    L = 2.500e6
    omega = 7.2921150e-5 * rot_fac
    a = 6376.0e3
    surf_trop_tdiff = 100.

    # pick longitudes to look at
    lons = pick_lons(data, lonin)

    # Evaluate equivalent potential temperature
    ept = ((data.temp + L / cp * data.sphum) *
           (1000. / data.pfull)**(2. / 7.)).sel(
               pfull=850., lon=lons).mean('lon').drop('pfull')
    data['ept'] = ept

    # moist entropy, s = cp ln(ept)
    s = cp * np.log(ept)

    # entropy curvature d/dy(cos(lat)^3/sin(lat) (Ts-Tt) ds/dy)
    dsdy = gr.ddy(s, vector=False)
    dsdy = dsdy * surf_trop_tdiff * np.cos(
        data.lat * np.pi / 180.)**3. / np.sin(data.lat * np.pi / 180.)
    data['s_curve'] = gr.ddy(dsdy, vector=False)

    # critical curvature = -4 * omega^2 * a^2 *cos(lat)^3 * sin(lat)
    crit_curve = -4. * omega**2 * a**2 * np.cos(
        data.lat * np.pi / 180.)**3. * np.sin(data.lat * np.pi / 180.)

    # Find max ept and location of max
    ept_max = [ept.max('lat'), ept.lat[ept.argmax('lat')]]

    # critical ept = ept_max * e^ (-chi * (cos(ept_max_lat)^2 - cos(lat)^2)^2 / cos(lat)^2)
    crit_ept_prefac = -1. * omega**2 * a**2 / cp / surf_trop_tdiff

    n = len(data.xofyear.values)
    crit_ept_fn = np.zeros([n, 64])
    crit_ept = np.zeros([n, 64])

    for i in range(n):
        crit_ept_fn[i, :] = (np.cos(ept_max[1][i] * np.pi / 180.)**2. -
                             np.cos(data.lat * np.pi / 180.)**2.)**2. / np.cos(
                                 data.lat * np.pi / 180.)**2.
        crit_ept[i, :] = ept_max[0].values[i] * np.exp(
            crit_ept_prefac * crit_ept_fn[i, :])

    data['crit_ept'] = xr.DataArray(crit_ept,
                                    coords=[data.xofyear, data.lat],
                                    dims=['xofyear', 'lat'])

    # moist entropy, s = cp ln(ept)
    crit_s = cp * np.log(data.crit_ept)

    # entropy curvature d/dy(cos(lat)^3/sin(lat) (Ts-Tt) ds/dy)
    dcrit_sdy = gr.ddy(crit_s, vector=False)
    dcrit_sdy = dcrit_sdy * surf_trop_tdiff * np.cos(
        data.lat * np.pi / 180.)**3. / np.sin(data.lat * np.pi / 180.)
    data['crit_s_curve'] = gr.ddy(dcrit_sdy, vector=False)
コード例 #2
0
ファイル: energy_eq.py プロジェクト: subond/python_scripts
def energy_eq(run, month, filename='plev_daily', lev=150.):

    #Load in dataset
    name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc'
    name = name_temp % month
    #read data into xarray
    data = xr.open_dataset(name, decode_times=False)

    # Calculate planetary vorticity
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)

    uchi = data.ucomp - data.upsi
    vchi = data.vcomp - data.vpsi

    pe_to_ke = -1. * (uchi * gr.ddx(data.height * 9.8) +
                      vchi * gr.ddy(data.height * 9.8, vector=False))

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f

    conv_1 = vor * (data.upsi * vchi - uchi * data.vpsi)

    conv_2 = -1. * omega * (data.upsi * gr.ddp(uchi) +
                            data.vpsi * gr.ddp(vchi))

    conv_3 = -1. * (gr.ddx(uchi) + gr.ddy(vchi)) * (data.upsi**2. +
                                                    data.vpsi**2.) / 2.

    ds = xr.Dataset(
        {
            'pe_to_ke': (['time', 'pfull', 'lat', 'lon'], pe_to_ke),
            'conv_1': (['time', 'pfull', 'lat', 'lon'], conv_1),
            'conv_2': (['time', 'pfull', 'lat', 'lon'], conv_2),
            'conv_3': (['time', 'pfull', 'lat', 'lon'], conv_3),
            'upsi': (['time', 'pfull', 'lat', 'lon'], data.upsi),
            'vpsi': (['time', 'pfull', 'lat', 'lon'], data.vpsi),
            'uchi': (['time', 'pfull', 'lat', 'lon'], uchi),
            'vchi': (['time', 'pfull', 'lat', 'lon'], vchi)
        },
        coords={
            'time': ('time', data.time),
            'pfull': ('pfull', data.pfull),
            'lat': ('lat', data.lat),
            'lon': ('lon', data.lon)
        })

    ds.coords['xofyear'] = np.mod(ds.time - 1., 360.) // 5 + 1.
    dsout = ds.groupby('xofyear').mean(('time'))

    fileout = '/scratch/rg419/Data_moist/' + run + '/run%03d/energy_eq.nc'
    fileout = fileout % month
    dsout.to_netcdf(path=fileout)

    print 'data written to ', fileout

    return dsout
コード例 #3
0
def partition_advection(data, lons, lev=150):

    #First do uu terms
    uu_trans_dx = -86400. * gr.ddx(
        (data.ucomp_sq -
         data.ucomp**2).sel(pfull=lev))  # <u'u'> = <uu> - <u><u>

    u = data.ucomp.sel(pfull=lev)  # u
    u_dx = -86400. * gr.ddx(u)  # dudx

    u_dudx_zav = u.sel(lon=lons).mean('lon') * u_dx.sel(lon=lons).mean(
        'lon')  # [u][dudx]

    u_dudx_stat = (u * u_dx).sel(
        lon=lons).mean('lon') - u_dudx_zav  # u*dudx* = [ududx] - [u][dudx]

    data['uu_trans_dx'] = (('xofyear', 'lat'),
                           uu_trans_dx.sel(lon=lons).mean('lon'))
    data['u_dudx_stat'] = (('xofyear', 'lat'), u_dudx_stat)
    data['u_dudx_zav'] = (('xofyear', 'lat'), u_dudx_zav)

    print 'uu terms done'

    #Next do uv terms
    uv_trans_dy = -86400. * gr.ddy(
        (data.ucomp_vcomp - data.ucomp * data.vcomp).sel(pfull=lev), uv=True)

    v = data.vcomp.sel(pfull=lev).load()  # v
    u_dy = -86400. * gr.ddy(u)  # dudy

    v_dudy_zav = v.sel(lon=lons).mean('lon') * u_dy.sel(lon=lons).mean(
        'lon')  # [v][dudy]

    v_dudy_stat = (v * u_dy).sel(
        lon=lons).mean('lon') - v_dudy_zav  # v*dudy* = [vdudy] - [v][dudy]

    data['uv_trans_dy'] = (('xofyear', 'lat'),
                           uv_trans_dy.sel(lon=lons).mean('lon'))
    data['v_dudy_stat'] = (('xofyear', 'lat'), v_dudy_stat)
    data['v_dudy_zav'] = (('xofyear', 'lat'), v_dudy_zav)

    print 'uv terms done'

    #Finally do uw terms
    uw_trans_dp = -86400. * gr.ddp(
        (data.ucomp_omega - data.ucomp * data.omega).sel(lon=lons).mean('lon'))

    w = data.omega.sel(pfull=lev).load()  # w
    u_dp = -86400. * (gr.ddp(data.ucomp)).sel(pfull=lev)  # dudp

    w_dudp_zav = w.sel(lon=lons).mean('lon') * u_dp.sel(lon=lons).mean('lon')
    w_dudp_stat = (w * u_dp).sel(lon=lons).mean('lon') - w_dudp_zav

    data['uw_trans_dp'] = (('xofyear', 'lat'), uw_trans_dp.sel(pfull=lev))
    data['w_dudp_stat'] = (('xofyear', 'lat'), w_dudp_stat)
    data['w_dudp_zav'] = (('xofyear', 'lat'), w_dudp_zav)

    print 'uw terms done'
コード例 #4
0
def transient_plot(run, ax_in, pentad, plot_land=True):

    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
                           '.nc')
    data_vort = xr.open_dataset(
        '/scratch/rg419/Data_moist/climatologies/vort_eq_' + run + '.nc')

    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp.sel(pfull=150))  # dvdx
    u_dy = gr.ddy(data.ucomp.sel(pfull=150))  # dudy
    vor = v_dx - u_dy + f
    div = gr.ddx(data.ucomp.sel(pfull=150)) + gr.ddy(data.vcomp.sel(pfull=150))
    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)

    horiz_md_mean = -86400.**2. * (data.ucomp.sel(pfull=150) * dvordx +
                                   data.vcomp.sel(pfull=150) * dvordy)
    stretching_mean = -86400.**2. * vor * div

    transient = (data_vort.horiz_md.sel(pfull=150).values +
                 data_vort.stretching.sel(pfull=150).values
                 ) * 86400.**2. - horiz_md_mean - stretching_mean

    land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
    land = xr.open_dataset(land_file)
    data['land'] = (('lat', 'lon'), land.land_mask)

    f1 = transient[pentad, :, :].plot.contourf(x='lon',
                                               y='lat',
                                               levels=np.arange(-5., 5.1, 1.),
                                               ax=ax_in,
                                               add_labels=False,
                                               extend='both',
                                               add_colorbar=False)

    if plot_land:
        data.land.plot.contour(x='lon',
                               y='lat',
                               levels=np.arange(0., 2., 1.),
                               ax=ax_in,
                               colors='k',
                               add_colorbar=False,
                               add_labels=False)
    (vor * 86400.)[pentad, :, :].plot.contour(x='lon',
                                              y='lat',
                                              levels=np.arange(-14., 15., 2.),
                                              ax=ax_in,
                                              colors='0.7',
                                              add_labels=False)

    ax_in.set_xlim(0, 180)
    ax_in.set_ylim(-30, 60)
    ax_in.set_xticks(np.arange(0., 185., 30.))
    ax_in.set_yticks(np.arange(-30., 65., 30.))

    return f1
コード例 #5
0
def vort_eq(run, month, filename='plev_daily', period_fac=1.):

    #Load in dataset
    name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/'+filename+'.nc'
    name = name_temp % month 
    #read data into xarray 
    data = xr.open_dataset( name, decode_times=False)
        
    # Calculate planetary vorticity
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat *np.pi/180)
    
    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f
    #vor_rel = v_dx - u_dy
    
    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)
    dvordp = gr.ddp(vor)
    
    # Calculate horizontal material derivative
    horiz_md = -1. * (data.ucomp * dvordx + data.vcomp * dvordy)
    
    # Calculate vertical material derivative
    vert_md = -1. * data.omega * dvordp
    
    # Now do the terms involving gradients of windspeed
    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
    stretching = -1. * vor * div
    tilting = gr.ddp(data.ucomp) * gr.ddy(data.omega, vector=False) - gr.ddp(data.vcomp) * gr.ddx(data.omega)
    
    total = horiz_md + vert_md + stretching + tilting
    
    ds = xr.Dataset({'horiz_md': (['time', 'pfull', 'lat', 'lon'], horiz_md),
                     'vert_md':  (['time', 'pfull', 'lat', 'lon'], vert_md),
                     'stretching':  (['time', 'pfull', 'lat', 'lon'], stretching),
                     'tilting':  (['time', 'pfull', 'lat', 'lon'], tilting),
                     'ucomp': (['time', 'pfull', 'lat', 'lon'], data.ucomp),
                     'vcomp': (['time', 'pfull', 'lat', 'lon'], data.vcomp),
                     'total':  (['time', 'pfull', 'lat', 'lon'], total)},
                     coords={'time': ('time', data.time),
                             'pfull': ('pfull', data.pfull),
                               'lat': ('lat', data.lat),
                               'lon': ('lon', data.lon)})
                              
    dsout = ds#.mean(('time'))
    
    fileout = '/scratch/rg419/Data_moist/' + run + '/run%03d/vort_eq.nc'  
    fileout = fileout % month
    dsout.to_netcdf(path=fileout)
    
    print 'data written to ', fileout

    return dsout
コード例 #6
0
def vort_eq_hm(run, lev=150, rotfac=1.0, period_fac=1.):

    rcParams['figure.figsize'] = 15, 6.25
    rcParams['font.size'] = 18
    rcParams['text.usetex'] = True

    plot_dir = '/scratch/rg419/plots/crit_lat_test/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    #Load in vorticity budget term means
    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
                           run + '.nc')

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    omega = 7.2921150e-5 * rotfac
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f

    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
    stretching_mean = -86400.**2. * vor * div
    stretching_hm = stretching_mean.mean('lon')

    mn_dic = month_dic(1)
    tickspace = np.arange(13, 72, 18) * period_fac
    labels = [mn_dic[(k + 5) / 6] for k in range(13, 72, 18)]
    levels = np.arange(-3., 3.1, 0.25)

    stretching_max = np.unravel_index(
        stretching_hm.sel(pfull=lev).isel(lat=range(32, 64)).argmin(),
        (data.pentad.size, 32))
    print stretching_max
    print data.lat[stretching_max[1] + 32]

    stretching_hm.sel(pfull=lev).plot.contourf(x='pentad',
                                               y='lat',
                                               levels=levels,
                                               extend='both',
                                               add_labels=False)
    #ax4.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
    plt.plot(data.pentad[stretching_max[0]], data.lat[stretching_max[1] + 32],
             'kx')
    plt.ylabel('Latitude')
    plt.xlabel('')
    plt.ylim(-60, 60)
    plt.yticks(np.arange(-60, 61, 30))
    plt.xticks(tickspace, labels, rotation=25)
    plt.title('Vortex stretching', fontsize=17)
    plt.grid(True, linestyle=':')
    plt.tight_layout()

    figname = 'vort_stretching_' + run + '.pdf'
    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()
コード例 #7
0
def get_stretching(run):

    #Load data
    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
                           run + '.nc')

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    f = 2 * mc.omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f - v_dx + u_dy
    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
    stretching_mean = -86400.**2. * (vor * div).sel(pfull=150.)

    return stretching_mean
コード例 #8
0
ファイル: aht_regional.py プロジェクト: subond/python_scripts
def aht_regional(run, lonin=[-1.,361.], period_fac=1.):
    
    area = mc.a*mc.a*cell_area(42, '/scratch/rg419/GFDL_model/GFDLmoistModel/')   # Area of grid cells
    
    #Load in data, add area to dataset
    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run + '.nc')
    data['area'] = (('lat','lon'), area)
    
    if lonin[1]>lonin[0]:
        lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]]
    else:
        lons = [data.lon[i] for i in range(len(data.lon)) if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]]
    
    vh = data.vcomp_temp * mc.cp_air + data.sphum_v * mc.L + data.height*data.vcomp * mc.grav
    
    dvhdy = gr.ddy(vh)
    
    aht = ((dvhdy.sum('pfull')*5000./9.8) * data.area).sum(('lon')).cumsum('lat')
    
#    vh = ((gr.ddy(data.vcomp * h).sum('pfull')*5000./9.8)) #.sel(lon=lons).sum('lon')
    
   # aht_div = gr.ddy(vh)*data.area.
    
    #aht = aht_div.cumsum('lat')
    
    aht_rm = rolling_mean(aht, int(5*period_fac))
    
    aht_rm.plot.contourf(x='xofyear', y='lat', levels=np.arange(-1.5e16,1.6e16,1.e15))
    plt.show()
    
    return aht_rm
コード例 #9
0
ファイル: horiz_adv_ll.py プロジェクト: subond/python_scripts
def get_horiz_adv(run):

    #Load data
    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
                           run + '.nc')

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    f = 2 * mc.omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f - v_dx + u_dy
    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)
    horiz_adv_mean = -86400.**2. * (data.ucomp * dvordx +
                                    data.vcomp * dvordy).sel(pfull=150.)

    return horiz_adv_mean
コード例 #10
0
def get_abs_vort(data):
    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f
    data['abs_vort'] = vor.mean('lon') * 86400.
コード例 #11
0
 def calc_vort(data):
     # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
     f = 2 * mc.omega * np.sin(data.lat * np.pi / 180)
     v_dx = gr.ddx(data.vcomp)  # dvdx
     u_dy = gr.ddy(data.ucomp)  # dudy
     vor = v_dx - u_dy + f
     abs_vort = vor.sel(pfull=150) * 86400.
     return abs_vort
コード例 #12
0
def uv_partitioning(data,lons):
        
    u_ztav = data.ucomp.sel(lon=lons).mean('lon') 
    v_ztav = data.vcomp.sel(lon=lons).mean('lon')
    
    uv_eddy = data.ucomp_vcomp.sel(lon=lons).mean('lon') - u_ztav * v_ztav 
                     
    uv_conv = gr.ddy(uv_eddy, uv=True) * -86400.

    return uv_conv, u_ztav
コード例 #13
0
ファイル: vort_line.py プロジェクト: subond/python_scripts
def vort_eq_hm(run, bef_aft, lev=150, lonin=[-1., 361.]):

    rcParams['figure.figsize'] = 10, 6.25
    rcParams['font.size'] = 18
    rcParams['text.usetex'] = True

    plot_dir = '/scratch/rg419/plots/paper_1_figs/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    #Also load climatological data so that transient eddies can be calculated
    data = xr.open_dataset(
        '/scratch/rg419/Data_moist/climatologies/vort_eq_uv' + run + '.nc')
    print 'climatology loaded'

    if lonin[1] > lonin[0]:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
        ]
    else:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
        ]

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f

    abs_vort = vor.sel(lon=lons).mean('lon') * 86400.

    print 'starting plotting'

    np.abs(
        vor.sel(lon=lons, xofyear=range(bef_aft[0], bef_aft[0] + 4)).mean(
            ('lon', 'xofyear')) * 86400.).plot(color='k')
    np.abs(
        vor.sel(lon=lons, xofyear=range(bef_aft[1], bef_aft[1] + 4)).mean(
            ('lon', 'xofyear')) * 86400.).plot(color='r')
    plt.xlim([-45, 45])
    plt.ylim([-12, 12])
    plt.grid(True, linestyle=':')

    if lonin == [-1., 361.]:
        figname = 'vort_line_' + run + '.pdf'
    else:
        figname = 'vort_line_' + run + '_' + str(int(lonin[0])) + '_' + str(
            int(lonin[1])) + '.pdf'

    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()
コード例 #14
0
def get_transient(run):

    #Load data
    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
                           run + '.nc')

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    f = 2 * mc.omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f
    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)
    horiz_adv_mean = -86400.**2. * (data.ucomp * dvordx +
                                    data.vcomp * dvordy).sel(pfull=150.)
    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
    stretching_mean = -86400.**2. * (vor * div).sel(pfull=150.)

    transient = data.horiz_md.sel(pfull=150.) + data.stretching.sel(
        pfull=150.) - horiz_adv_mean - stretching_mean

    return transient
コード例 #15
0
def partition_advection(data, lons, lev=150):

    #Do vv terms
    vv_trans_dy = -86400. * gr.ddy(
        (data.vcomp_sq - data.vcomp * data.vcomp).sel(pfull=lev), uv=True)

    v = data.vcomp.sel(pfull=lev).load()  # v
    v_dy = -86400. * gr.ddy(v)  # dudy

    v_dvdy_zav = v.sel(lon=lons).mean('lon') * v_dy.sel(lon=lons).mean(
        'lon')  # [v][dudy]

    v_dvdy_stat = (v * v_dy).sel(
        lon=lons).mean('lon') - v_dvdy_zav  # v*dudy* = [vdudy] - [v][dudy]

    data['vv_trans_dy'] = (('xofyear', 'lat'),
                           vv_trans_dy.sel(lon=lons).mean('lon'))
    data['v_dvdy_stat'] = (('xofyear', 'lat'), v_dvdy_stat)
    data['v_dvdy_zav'] = (('xofyear', 'lat'), v_dvdy_zav)

    print 'vv terms done'

    #Do vw terms
    vw_trans_dp = -86400. * gr.ddp(
        (data.vcomp_omega - data.vcomp * data.omega).sel(lon=lons).mean('lon'))

    w = data.omega.sel(pfull=lev).load()  # w
    v_dp = -86400. * (gr.ddp(data.vcomp)).sel(pfull=lev)  # dudp

    w_dvdp_zav = w.sel(lon=lons).mean('lon') * v_dp.sel(lon=lons).mean('lon')
    w_dvdp_stat = (w * v_dp).sel(lon=lons).mean('lon') - w_dvdp_zav

    data['vw_trans_dp'] = (('xofyear', 'lat'), vw_trans_dp.sel(pfull=lev))
    data['w_dvdp_stat'] = (('xofyear', 'lat'), w_dvdp_stat)
    data['w_dvdp_zav'] = (('xofyear', 'lat'), w_dvdp_zav)

    print 'vw terms done'
コード例 #16
0
ファイル: vort.py プロジェクト: subond/python_scripts
def get_abs_vort(run):
        
    #Load data
    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_'+run+'.nc')

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    f = 2 * mc.omega * np.sin(data.lat *np.pi/180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f
    
    abs_vort = vor*86400.
    
    abs_vort = abs_vort.sel(pfull=150) 
    
    return abs_vort
コード例 #17
0
ファイル: psi_u_both.py プロジェクト: subond/python_scripts
def uv_partitioning(data, start_index, lons):

    a = 6376.0e3  #radius used in model
    coslat = np.cos(data.lat * np.pi / 180)

    u_ztav = data.ucomp[start_index:start_index +
                        4, :, :, :].sel(lon=lons).mean(('xofyear', 'lon'))
    v_ztav = data.vcomp[start_index:start_index +
                        4, :, :, :].sel(lon=lons).mean(('xofyear', 'lon'))

    uv_eddy = data.ucomp_vcomp[start_index:start_index +
                               4, :, :, :].sel(lon=lons).mean(
                                   ('xofyear', 'lon')) - u_ztav * v_ztav

    uv_conv = -86400. * gr.ddy(uv_eddy, uv=True)

    #uv_conv = xr.DataArray( cfd( (uv_eddy*coslat*coslat).values, data.lat*np.pi/180, 1 ), [('pfull', data.pfull ), ('lat', data.lat)])
    #uv_conv = -86400.*uv_conv/coslat/coslat/a

    return uv_conv, u_ztav
コード例 #18
0
def subcloud_entropy_gradient(data, lonin=[-1.,361.]):
    
    # declare constants
    cp = 287.04/2*7
    L = 2.500e6
    
    # pick longitudes to look at
    lons = pick_lons(data, lonin)
    
    # Evaluate equivalent potential temperature
    ept = ((data.temp + L/cp*data.sphum)*(1000./data.pfull)**(2./7.)).sel(pfull=850., lon=lons).mean('lon').drop('pfull')    
    data['ept'] = ept
    
    # moist entropy, s = cp ln(ept)
    s = cp*np.log(ept)
    
    # entropy gradient
    dsdy = gr.ddy(s, vector=False)
    
    return dsdy
コード例 #19
0
def vort_eq_ss(run, lonin=[-1., 361.]):

    rcParams['figure.figsize'] = 15, 6.25
    rcParams['font.size'] = 18
    rcParams['text.usetex'] = True

    plot_dir = '/scratch/rg419/plots/steady_state_runs/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    #Load in vorticity budget term means
    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
                           run + '.nc')
    print 'vorticity budget data loaded'

    if lonin[1] > lonin[0]:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
        ]
    else:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
        ]

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    f = 2 * mc.omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f

    abs_vort = vor.sel(lon=lons).mean('lon') * 86400.

    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)

    horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy)

    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
    stretching_mean = -86400.**2. * vor * div

    horiz_md_hm = horiz_md_mean.sel(lon=lons).mean('lon')
    stretching_hm = stretching_mean.sel(lon=lons).mean('lon')

    print 'starting plotting'

    levels = np.arange(-1.5, 1.6, 0.25)

    f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
                                                         3,
                                                         sharex='col',
                                                         sharey='row')
    plt.set_cmap('RdBu_r')
    #First plot
    f2 = horiz_md_hm.plot.contourf(x='lat',
                                   y='pfull',
                                   levels=levels,
                                   ax=ax1,
                                   extend='both',
                                   add_labels=False,
                                   yincrease=False)
    #ax1.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
    ax1.set_xlabel('Latitude')
    ax1.set_ylabel('Pressure, hPa')
    ax1.set_xlim(-60, 60)
    ax1.set_xticks(np.arange(-60, 61, 30))
    ax1.grid(True, linestyle=':')
    #ax1.text(-15, 60, 'a)')

    #Second plot
    dvordy.sel(lon=lons).mean('lon').plot.contourf(x='lat',
                                                   y='pfull',
                                                   levels=np.arange(
                                                       -1.e-10, 1.05e-10,
                                                       0.1e-10),
                                                   ax=ax2,
                                                   extend='both',
                                                   add_labels=False,
                                                   yincrease=False)
    #ax2.set_ylim(-60,60)
    ax2.set_xticks(np.arange(-60, 61, 30))
    ax2.grid(True, linestyle=':')
    #ax2.text(-7, 60, 'b)')

    #Third plot
    data.vcomp.sel(lon=lons).mean('lon').plot.contourf(x='lat',
                                                       y='pfull',
                                                       levels=np.arange(
                                                           -12., 13, 2.),
                                                       ax=ax3,
                                                       extend='both',
                                                       add_labels=False,
                                                       yincrease=False)
    #ax3.set_ylim(-60,60)
    ax3.set_xticks(np.arange(-60, 61, 30))
    ax3.grid(True, linestyle=':')
    #ax3.text(-7, 60, 'c)')

    #Fourth plot
    stretching_hm.plot.contourf(x='lat',
                                y='pfull',
                                levels=levels,
                                ax=ax4,
                                extend='both',
                                add_labels=False,
                                yincrease=False)
    #ax4.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
    ax4.set_xlabel('Latitude')
    ax4.set_ylabel('Pressure, hPa')
    ax4.set_xlim(-60, 60)
    ax4.set_xticks(np.arange(-60, 61, 30))
    #ax4.set_xticks(tickspace)
    #ax4.set_xticklabels(labels,rotation=25)
    ax4.grid(True, linestyle=':')
    #ax4.text(-15, 60, 'd)')

    #Fifth plot
    (div.sel(lon=lons).mean('lon') * 86400.).plot.contourf(x='lat',
                                                           y='pfull',
                                                           levels=np.arange(
                                                               -1., 1.1, 0.1),
                                                           ax=ax5,
                                                           extend='both',
                                                           add_labels=False,
                                                           yincrease=False)
    ax5.set_xlim(-60, 60)
    ax4.set_xlabel('Latitude')
    ax5.set_xticks(np.arange(-60, 61, 30))
    # ax5.set_xticks(tickspace)
    #ax5.set_xticklabels(labels,rotation=25)
    ax5.grid(True, linestyle=':')
    #ax5.text(-7, 60, 'e)')

    #Sixth plot
    (vor.sel(lon=lons).mean('lon') * 86400.).plot.contourf(x='lat',
                                                           y='pfull',
                                                           levels=np.arange(
                                                               -12., 13., 2.),
                                                           ax=ax6,
                                                           extend='both',
                                                           add_labels=False,
                                                           yincrease=False)
    #ax6.set_ylim(-60,60)
    ax4.set_xlabel('Latitude')
    ax6.set_xticks(np.arange(-60, 61, 30))
    #ax6.set_xticks(tickspace)
    #ax6.set_xticklabels(labels,rotation=25)
    ax6.grid(True, linestyle=':')
    #ax6.text(-7, 60, 'f)')

    plt.subplots_adjust(right=0.97,
                        left=0.08,
                        top=0.93,
                        bottom=0.1,
                        hspace=0.25,
                        wspace=0.15)

    if lonin == [-1., 361.]:
        figname = 'vort_breakdown_' + run + '.pdf'
    else:
        figname = 'vort_breakdown_' + run + '_' + str(int(
            lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'

    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()
コード例 #20
0
def max_div(run, lats=np.arange(-90, 90.1, 0.1), rotfac=1.0):

    #Open data
    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/vort_eq_' +
                           run + '.nc')

    #Calculate vorticity, divergence, and pentad mean vortex stretching

    f = 2 * mc.omega * rotfac * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f

    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
    stretching_mean = -86400.**2. * vor * div

    # Select upper level and take zonal mean
    # Take a 20 day window rolling mean, and interpolate onto a 0.1 deg lat grid.
    # Find max values of divergence and stretching, and locate lat of max divergence
    div_hm = div.mean('lon').sel(pfull=150.)
    div_rm = rolling_mean(div_hm, 4, tcoord='pentad')
    div_int = interp_field(div_rm)
    div_max = div_int.max('lat') * 86400.
    div_argmax = div_int.lat.values[div_int.argmax('lat').values]
    div_argmax = xr.DataArray(div_argmax,
                              coords=[div_int.xofyear],
                              dims=['xofyear'])

    stretching_hm = stretching_mean.mean('lon').sel(pfull=150.)
    stretching_hm[:, 32:] = stretching_hm[:, 32:] * -1.
    stretching_rm = rolling_mean(stretching_hm, 4, tcoord='pentad')
    stretching_int = interp_field(stretching_rm)
    stretching_max = stretching_int.max('lat')

    # Plot up max divergence and vortex stretching tendency in top panel, latitude of max divergence below
    #fig, ax1 = plt.subplots()
    fig, (ax1, ax3) = plt.subplots(2, sharex=True)
    div_max.plot(ax=ax1, color='k')
    ax1.set_xlabel('')
    ax1.set_title('')
    # Make the y-axis label, ticks and tick labels match the line color.
    ax1.set_ylabel('Divergence, day$^{-1}$')
    ax1.tick_params('y')
    ax1.grid(True, linestyle=':')
    ax1.set_xlim(0, max(div_max.xofyear) + 1)

    ax2 = ax1.twinx()
    stretching_max.plot(color='b', ax=ax2)
    ax2.set_title('')
    ax2.set_ylabel('Vortex stretching, day$^{-2}$', color='b')
    ax2.tick_params('y', colors='b')
    ax2.set_xlim(0, max(div_max.xofyear) + 1)

    div_argmax.plot(ax=ax3, color='k')
    ax3.set_title('')
    ax3.set_xlabel('Pentad')
    ax3.set_ylabel('Lat of max divergence')
    ax3.grid(True, linestyle=':')
    ax3.set_xlim(0, max(div_max.xofyear) + 1)

    fig.tight_layout()
    plt.savefig(plot_dir + 'div_tseries_' + run + '.pdf', format='pdf')
    plt.close()
コード例 #21
0
def vort_eq_hm(run, lev=150, lonin=[-1., 361.]):

    rcParams['font.size'] = 18
    rcParams['text.usetex'] = True

    plot_dir = '/scratch/rg419/plots/paper_1_figs/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    #Also load climatological data so that transient eddies can be calculated
    data = xr.open_dataset(
        '/scratch/rg419/Data_moist/climatologies/vort_eq_uv' + run + '.nc')
    print 'climatology loaded'

    if lonin[1] > lonin[0]:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
        ]
    else:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
        ]

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f

    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)

    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)

    mn_dic = month_dic(1)
    tickspace = range(13, 72, 18)
    labels = [mn_dic[(k + 5) / 6] for k in tickspace]

    print 'starting plotting'

    # Easier to produce individual plots as scales are different.

    vor.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                y='lat',
                                                extend='both',
                                                add_labels=False,
                                                levels=np.arange(
                                                    -0.00024, 0.00025,
                                                    0.00002))
    plt.ylabel('Latitude')
    plt.ylim(-60, 60)
    plt.yticks(np.arange(-60, 61, 30))
    plt.xlabel('')
    plt.xticks(tickspace, labels, rotation=25)
    plt.grid(True, linestyle=':')
    if lonin == [-1., 361.]:
        figname = 'vor_hm_' + run + '.pdf'
    else:
        figname = 'vor_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
            int(lonin[1])) + '.pdf'
    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()

    div.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                y='lat',
                                                extend='both',
                                                add_labels=False,
                                                levels=np.arange(
                                                    -0.000012, 0.000012,
                                                    0.000001))
    plt.ylabel('Latitude')
    plt.ylim(-60, 60)
    plt.yticks(np.arange(-60, 61, 30))
    plt.xlabel('')
    plt.xticks(tickspace, labels, rotation=25)
    plt.grid(True, linestyle=':')
    if lonin == [-1., 361.]:
        figname = 'div_hm_' + run + '.pdf'
    else:
        figname = 'div_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
            int(lonin[1])) + '.pdf'
    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()

    dvordx.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                   y='lat',
                                                   extend='both',
                                                   add_labels=False)
    plt.ylabel('Latitude')
    plt.ylim(-60, 60)
    plt.yticks(np.arange(-60, 61, 30))
    plt.xlabel('')
    plt.xticks(tickspace, labels, rotation=25)
    plt.grid(True, linestyle=':')
    if lonin == [-1., 361.]:
        figname = 'dvordx_hm_' + run + '.pdf'
    else:
        figname = 'dvordx_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
            int(lonin[1])) + '.pdf'
    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()

    dvordy.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                   y='lat',
                                                   extend='both',
                                                   add_labels=False,
                                                   levels=np.arange(
                                                       -1.2e-10, 1.25e-10,
                                                       0.1e-10))
    plt.ylabel('Latitude')
    plt.ylim(-60, 60)
    plt.yticks(np.arange(-60, 61, 30))
    plt.xlabel('')
    plt.xticks(tickspace, labels, rotation=25)
    plt.grid(True, linestyle=':')
    if lonin == [-1., 361.]:
        figname = 'dvordy_hm_' + run + '.pdf'
    else:
        figname = 'dvordy_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
            int(lonin[1])) + '.pdf'
    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()

    data.ucomp.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                       y='lat',
                                                       extend='both',
                                                       add_labels=False)
    plt.ylabel('Latitude')
    plt.ylim(-60, 60)
    plt.yticks(np.arange(-60, 61, 30))
    plt.xlabel('')
    plt.xticks(tickspace, labels, rotation=25)
    plt.grid(True, linestyle=':')
    if lonin == [-1., 361.]:
        figname = 'u_hm_' + run + '.pdf'
    else:
        figname = 'u_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
            int(lonin[1])) + '.pdf'
    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()

    data.vcomp.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                       y='lat',
                                                       extend='both',
                                                       add_labels=False,
                                                       levels=np.arange(
                                                           -15., 15.5, 1.))
    plt.ylabel('Latitude')
    plt.ylim(-60, 60)
    plt.yticks(np.arange(-60, 61, 30))
    plt.xlabel('')
    plt.xticks(tickspace, labels, rotation=25)
    plt.grid(True, linestyle=':')
    if lonin == [-1., 361.]:
        figname = 'v_hm_' + run + '.pdf'
    else:
        figname = 'v_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
            int(lonin[1])) + '.pdf'
    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()
コード例 #22
0
def mom_budg(run, lev=150):

    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
                           '.nc')

    #First do uu terms
    uu_trans_dx = -86400. * gr.ddx(
        (data.ucomp_sq -
         data.ucomp**2).sel(pfull=lev))  # <u'u'> = <uu> - <u><u>

    u = data.ucomp.sel(pfull=lev)  # u
    u_dx = -86400. * gr.ddx(u)  # dudx

    u_ed = u - u.mean('lon')
    u_dx_ed = u_dx - u_dx.mean('lon')

    u_dudx_zav = u.mean('lon') * u_dx.mean(
        'lon')  # [u][dudx], where brackets denote mean over all longitudes

    u_dudx_cross1 = u.mean('lon') * u_dx_ed  # [u]dudx*

    u_dudx_cross2 = u_ed * u_dx.mean('lon')  # u*[dudx]

    u_dudx_stat = u_ed * u_dx_ed  # u*dudx*

    data['uu_trans_dx'] = (('xofyear', 'lat', 'lon'), uu_trans_dx)
    data['u_dudx_cross1'] = (('xofyear', 'lat', 'lon'), u_dudx_cross1)
    data['u_dudx_cross2'] = (('xofyear', 'lat', 'lon'), u_dudx_cross2)
    data['u_dudx_stat'] = (('xofyear', 'lat', 'lon'), u_dudx_stat)
    data['u_dudx_zav'] = (('xofyear', 'lat'), u_dudx_zav)

    #Next do uv terms
    uv_trans_dy = -86400. * gr.ddy(
        (data.ucomp_vcomp - data.ucomp * data.vcomp).sel(pfull=lev), uv=True)

    v = data.vcomp.sel(pfull=lev).load()  # v
    u_dy = -86400. * gr.ddy(u)  # dudy

    v_ed = v - v.mean('lon')
    u_dy_ed = u_dy - u_dy.mean('lon')

    v_dudy_zav = v.mean('lon') * u_dy.mean('lon')  # [v][dudy]

    v_dudy_cross1 = v.mean('lon') * u_dy_ed  # [v]dudy*

    v_dudy_cross2 = v_ed * u_dy.mean('lon')  # v*[dudy]

    v_dudy_stat = v_ed * u_dy_ed  # v*dudy*

    data['uv_trans_dy'] = (('xofyear', 'lat', 'lon'), uv_trans_dy)
    data['v_dudy_cross1'] = (('xofyear', 'lat', 'lon'), v_dudy_cross1)
    data['v_dudy_cross2'] = (('xofyear', 'lat', 'lon'), v_dudy_cross2)
    data['v_dudy_stat'] = (('xofyear', 'lat', 'lon'), v_dudy_stat)
    data['v_dudy_zav'] = (('xofyear', 'lat'), v_dudy_zav)

    #Finally do uw terms
    uw_trans_dp = -86400. * gr.ddp(
        (data.ucomp_omega - data.ucomp * data.omega))

    w = data.omega.sel(pfull=lev).load()  # w
    u_dp = -86400. * (gr.ddp(data.ucomp)).sel(pfull=lev)  # dudp

    w_ed = w - w.mean('lon')
    u_dp_ed = u_dp - u_dp.mean('lon')

    w_dudp_zav = w.mean('lon') * u_dp.mean('lon')  # [w][dudp]

    w_dudp_cross1 = w.mean('lon') * u_dp_ed  # [w]dudp*

    w_dudp_cross2 = w_ed * u_dp.mean('lon')  # w*[dudp]

    w_dudp_stat = w_ed * u_dp_ed  # w*dudp*

    data['uw_trans_dp'] = (('xofyear', 'lat', 'lon'),
                           uw_trans_dp.sel(pfull=lev))
    data['w_dudp_cross1'] = (('xofyear', 'lat', 'lon'), w_dudp_cross1)
    data['w_dudp_cross2'] = (('xofyear', 'lat', 'lon'), w_dudp_cross2)
    data['w_dudp_stat'] = (('xofyear', 'lat', 'lon'), w_dudp_stat)
    data['w_dudp_zav'] = (('xofyear', 'lat'), w_dudp_zav)

    #Coriolis
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    fv = data.vcomp.sel(pfull=lev) * f * 86400.
    fv_mean = fv.mean('lon')
    fv_local = fv - fv_mean

    totp = (data.convection_rain + data.condensation_rain) * 86400.

    #Geopotential gradient
    dphidx = gr.ddx(data.height.sel(pfull=lev))
    dphidx = -86400. * 9.8 * dphidx
    fv_ageo = fv_local + dphidx

    mom_mean = data.u_dudx_zav + data.v_dudy_zav + data.w_dudp_zav
    mom_cross = data.u_dudx_cross1 + data.v_dudy_cross1 + data.w_dudp_cross1 + data.u_dudx_cross2 + data.v_dudy_cross2 + data.w_dudp_cross2
    mom_trans = data.uu_trans_dx + data.uv_trans_dy + data.uw_trans_dp
    mom_stat = data.u_dudx_stat + data.v_dudy_stat + data.w_dudp_stat

    mom_sum = fv_local + fv_mean + dphidx + mom_mean + mom_trans + mom_stat + mom_cross

    data['mom_mean'] = (('xofyear', 'lat'), mom_mean)
    data['mom_cross'] = (('xofyear', 'lat', 'lon'), mom_cross)
    data['mom_trans'] = (('xofyear', 'lat', 'lon'), mom_trans)
    data['mom_stat'] = (('xofyear', 'lat', 'lon'), mom_stat)
    data['fv_local'] = (('xofyear', 'lat', 'lon'), fv_local)
    data['fv_ageo'] = (('xofyear', 'lat', 'lon'), fv_ageo)
    data['fv_mean'] = (('xofyear', 'lat'), fv_mean)
    data['dphidx'] = (('xofyear', 'lat', 'lon'), dphidx)
    data['mom_sum'] = (('xofyear', 'lat', 'lon'), mom_sum)

    return data
コード例 #23
0
def vort_eq_hm(run, lev=150, lonin=[-1., 361.]):

    rcParams['figure.figsize'] = 15, 6.25
    rcParams['font.size'] = 18
    rcParams['text.usetex'] = True

    plot_dir = '/scratch/rg419/plots/paper_1_figs/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    #Load in vorticity budget term means
    data_vort = xr.open_dataset(
        '/scratch/rg419/Data_moist/climatologies/vort_eq_' + run + '.nc')
    data_vort = data_vort * 86400.**2.  #Convert to day^-2

    print 'vorticity budget data loaded'

    #Also load climatological data so that transient eddies can be calculated
    data = xr.open_dataset(
        '/scratch/rg419/Data_moist/climatologies/vort_eq_uv' + run + '.nc')
    print 'climatology loaded'

    if lonin[1] > lonin[0]:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
        ]
    else:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
        ]

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f

    abs_vort = vor.sel(lon=lons).mean('lon') * 86400.

    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)

    horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy)

    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
    stretching_mean = -86400.**2. * vor * div

    horiz_md_hm = horiz_md_mean.sel(lon=lons).mean('lon')
    stretching_hm = stretching_mean.sel(lon=lons).mean('lon')

    mn_dic = month_dic(1)
    tickspace = range(13, 72, 18)
    labels = [mn_dic[(k + 5) / 6] for k in tickspace]
    levels = np.arange(-1.5, 1.6, 0.25)

    print 'starting plotting'

    # Four subplots
    f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
                                                         3,
                                                         sharex='col',
                                                         sharey='row')
    plt.set_cmap('RdBu_r')
    #First plot
    f2 = horiz_md_hm.plot.contourf(x='xofyear',
                                   y='lat',
                                   levels=levels,
                                   ax=ax1,
                                   extend='both',
                                   add_labels=False)
    #ax1.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
    ax1.set_ylabel('Latitude')
    ax1.set_ylim(-60, 60)
    ax1.set_title('Horizontal advection', fontsize=17)
    ax1.set_yticks(np.arange(-60, 61, 30))
    ax1.grid(True, linestyle=':')
    ax1.text(-15, 60, 'a)')

    #Second plot
    dvordy.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                   y='lat',
                                                   levels=np.arange(
                                                       -1.e-10, 1.05e-10,
                                                       0.1e-10),
                                                   ax=ax2,
                                                   extend='both',
                                                   add_labels=False)
    #ax2.set_ylim(-60,60)
    ax2.set_yticks(np.arange(-60, 61, 30))
    ax2.set_title('$\partial (\overline{\zeta} + f) /\partial y$', fontsize=17)
    ax2.grid(True, linestyle=':')
    ax2.text(-7, 60, 'b)')

    #Third plot
    data.vcomp.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                       y='lat',
                                                       levels=np.arange(
                                                           -12., 13, 2.),
                                                       ax=ax3,
                                                       extend='both',
                                                       add_labels=False)
    ax3.set_ylim(-60, 60)
    ax3.set_yticks(np.arange(-60, 61, 30))
    ax3.set_title('$\overline{v}$', fontsize=17)
    ax3.grid(True, linestyle=':')
    ax3.text(-7, 60, 'c)')

    #Fourth plot
    stretching_hm.plot.contourf(x='xofyear',
                                y='lat',
                                levels=levels,
                                ax=ax4,
                                extend='both',
                                add_labels=False)
    #ax4.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
    ax4.set_ylabel('Latitude')
    ax4.set_ylim(-60, 60)
    ax4.set_yticks(np.arange(-60, 61, 30))
    ax4.set_xticks(tickspace)
    ax4.set_xticklabels(labels, rotation=25)
    ax4.set_title('Vortex stretching', fontsize=17)
    ax4.grid(True, linestyle=':')
    ax4.text(-15, 60, 'd)')

    #Fith plot
    (div.sel(lon=lons).mean('lon') * 86400.).plot.contourf(x='xofyear',
                                                           y='lat',
                                                           levels=np.arange(
                                                               -0.6, 0.7, 0.1),
                                                           ax=ax5,
                                                           extend='both',
                                                           add_labels=False)
    ax5.set_ylim(-60, 60)
    ax5.set_yticks(np.arange(-60, 61, 30))
    ax5.set_xticks(tickspace)
    ax5.set_xticklabels(labels, rotation=25)
    ax5.grid(True, linestyle=':')
    ax5.set_title('Divergence', fontsize=17)
    ax5.text(-7, 60, 'e)')

    #Sixth plot
    (vor.sel(lon=lons).mean('lon') * 86400.).plot.contourf(x='xofyear',
                                                           y='lat',
                                                           levels=np.arange(
                                                               -12., 13., 2.),
                                                           ax=ax6,
                                                           extend='both',
                                                           add_labels=False)
    ax6.set_ylim(-60, 60)
    ax6.set_yticks(np.arange(-60, 61, 30))
    ax6.set_xticks(tickspace)
    ax6.set_xticklabels(labels, rotation=25)
    ax6.grid(True, linestyle=':')
    ax6.set_title('Absolute vorticity', fontsize=17)
    ax6.text(-7, 60, 'f)')

    plt.subplots_adjust(right=0.97,
                        left=0.08,
                        top=0.93,
                        bottom=0.1,
                        hspace=0.25,
                        wspace=0.15)

    if lonin == [-1., 361.]:
        figname = 'vort_breakdown_' + run + '.pdf'
    else:
        figname = 'vort_breakdown_' + run + '_' + str(int(
            lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'

    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()
コード例 #24
0
def mom_budg_hm(run,
                lev=150,
                filename='plev_pentad',
                timeav='pentad',
                period_fac=1.,
                lonin=[-1., 361.]):

    a = 6376.0e3

    rcParams['figure.figsize'] = 12, 5.5
    rcParams['font.size'] = 18
    rcParams['text.usetex'] = True

    plot_dir = '/scratch/rg419/plots/itcz_theory_testing/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
                           '.nc')

    if lonin[1] > lonin[0]:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
        ]
    else:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
        ]

    #advective terms
    partition_advection(data, lons, lev=150)

    #Coriolis
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180.)
    fu = -1. * data.ucomp.sel(pfull=lev) * f * 86400.
    fu = fu.sel(lon=lons).mean('lon')

    metric_term = (-86400. * data.ucomp.sel(pfull=lev) *
                   data.ucomp.sel(pfull=lev) *
                   np.tan(data.lat * np.pi / 180.) / a).mean('lon')

    #totp = ((data.convection_rain + data.condensation_rain)*86400.).sel(lon=lons).mean('lon')
    #abs_vort = (data.vor + f).sel(lon=lons).mean('lon')*86400.

    #Geopotential gradient
    dphidy = gr.ddy(data.height.sel(pfull=lev), vector=False)
    dphidy = -86400. * 9.8 * dphidy.sel(lon=lons).mean('lon')

    mom_mean = data.v_dvdy_zav + data.w_dvdp_zav
    mom_trans = data.vv_trans_dy + data.vw_trans_dp
    mom_stat = data.v_dvdy_stat + data.w_dvdp_stat

    mom_sum = mom_mean + mom_trans + mom_stat + fu + dphidy + metric_term

    levels = np.arange(-300., 300.1, 30.)

    mn_dic = month_dic(1)
    tickspace = range(13, 72, 18)
    labels = [mn_dic[(k + 5) / 6] for k in tickspace]

    # Six subplots
    fig, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
                                                           3,
                                                           sharex='col',
                                                           sharey='row')
    plt.set_cmap('RdBu_r')
    #First plot
    f1 = fu.plot.contourf(ax=ax1,
                          x='xofyear',
                          y='lat',
                          extend='both',
                          levels=levels,
                          add_colorbar=False,
                          add_labels=False)
    #totp.plot.contour(ax=ax1, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
    ax1.set_ylabel('Latitude')
    ax1.set_title('Coriolis', fontsize=17)
    ax1.set_ylim(-60, 60)
    ax1.grid(True, linestyle=':')
    ax1.set_yticks(np.arange(-60., 61., 30.))
    ax1.text(-15, 60, 'a)')

    #Second plot
    mom_mean.plot.contourf(ax=ax2,
                           x='xofyear',
                           y='lat',
                           extend='both',
                           levels=levels,
                           add_colorbar=False,
                           add_labels=False)
    #ax2.contour(data.xofyear, data.lat, abs_vort.sel(pfull=lev).T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2)
    #abs_vort.sel(pfull=lev).plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', levels=np.arange(-2.,6.,2.), add_colorbar=False, add_labels=False, colors='k', linewidths=2)
    #totp.plot.contour(ax=ax2, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
    ax2.grid(True, linestyle=':')
    ax2.set_title('Mean state advection', fontsize=17)
    ax2.set_ylim(-60, 60)
    ax2.text(-5, 60, 'b)')

    #Third plot
    dphidy.plot.contourf(ax=ax3,
                         x='xofyear',
                         y='lat',
                         extend='both',
                         levels=levels,
                         add_colorbar=False,
                         add_labels=False)
    #totp.plot.contour(ax=ax3, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
    ax3.grid(True, linestyle=':')
    ax3.set_ylim(-60, 60)
    ax3.set_title('Geopotential gradient', fontsize=17)
    ax3.text(-5, 60, 'c)')

    #Fourth plot
    mom_trans.plot.contourf(ax=ax4,
                            x='xofyear',
                            y='lat',
                            extend='both',
                            levels=levels,
                            add_colorbar=False,
                            add_labels=False)
    #totp.plot.contour(ax=ax4, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
    ax4.grid(True, linestyle=':')
    ax4.set_ylabel('Latitude')
    ax4.set_xticks(tickspace)
    ax4.set_xticklabels(labels, rotation=25)
    ax4.set_ylim(-60, 60)
    ax4.set_title('Transient eddy flux conv.', fontsize=17)
    ax4.set_yticks(np.arange(-60., 61., 30.))
    ax4.text(-15, 60, 'd)')

    #Fifth plot
    metric_term.plot.contourf(ax=ax5,
                              x='xofyear',
                              y='lat',
                              extend='both',
                              levels=levels,
                              add_colorbar=False,
                              add_labels=False)
    #totp.plot.contour(ax=ax5, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
    ax5.grid(True, linestyle=':')
    ax5.set_xticks(tickspace)
    ax5.set_xticklabels(labels, rotation=25)
    ax5.set_title('Metric term', fontsize=17)
    ax5.set_ylim(-60, 60)
    ax5.text(-5, 60, 'e)')

    #Sixth plot
    mom_sum.plot.contourf(ax=ax6,
                          x='xofyear',
                          y='lat',
                          extend='both',
                          levels=levels,
                          add_colorbar=False,
                          add_labels=False)
    #totp.plot.contour(ax=ax6, x='xofyear', y='lat', extend='both', levels=np.arange(-92.,109.,100.), add_colorbar=False, add_labels=False, alpha=0.25, colors='k', linewidths=2)
    ax6.grid(True, linestyle=':')
    ax6.set_xticks(tickspace)
    ax6.set_xticklabels(labels, rotation=25)
    ax6.set_ylim(-60, 60)
    ax6.set_title('Residual', fontsize=17)
    ax6.text(-5, 60, 'f)')

    plt.subplots_adjust(right=0.97,
                        left=0.1,
                        top=0.95,
                        bottom=0.,
                        hspace=0.25,
                        wspace=0.12)
    #Colorbar
    cb1 = fig.colorbar(f1,
                       ax=[ax1, ax2, ax3, ax4, ax5, ax6],
                       use_gridspec=True,
                       orientation='horizontal',
                       fraction=0.15,
                       pad=0.15,
                       aspect=30,
                       shrink=0.5)
    #cb1=fig.colorbar(f1, ax=[ax1,ax2,ax3,ax4,ax5,ax6], use_gridspec=True,fraction=0.15, aspect=30)
    #cb1.set_label('$ms^{-1}day^{-1}$')

    if lonin == [-1., 361.]:
        figname = 'merid_mom_budg_' + run + '.pdf'
    else:
        figname = 'merid_mom_budg_' + run + '_' + str(int(
            lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'

    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()
コード例 #25
0
def vort_budg(run, ax_s, ax_h, pentad):

    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
                           '.nc')
    data['rain'] = (('xofyear', 'lat', 'lon'),
                    (data.convection_rain + data.condensation_rain) * 86400.)

    lons = [
        data.lon[i] for i in range(len(data.lon))
        if data.lon[i] >= 60. and data.lon[i] < 150.
    ]
    lats = [
        data.lat[i] for i in range(len(data.lat))
        if data.lat[i] >= -30. and data.lat[i] < 60.
    ]

    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp.sel(pfull=150))  # dvdx
    u_dy = gr.ddy(data.ucomp.sel(pfull=150))  # dudy
    vor = v_dx - u_dy + f
    div = gr.ddx(data.ucomp.sel(pfull=150)) + gr.ddy(data.vcomp.sel(pfull=150))

    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)

    horiz_md_mean = -86400.**2. * (data.ucomp.sel(pfull=150) * dvordx +
                                   data.vcomp.sel(pfull=150) * dvordy)
    stretching_mean = -86400.**2. * vor * div

    land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
    land = xr.open_dataset(land_file)

    f1 = stretching_mean[pentad:pentad +
                         4, :, :].mean('xofyear').plot.contourf(
                             x='lon',
                             y='lat',
                             levels=np.arange(-5., 5.1, 1.),
                             ax=ax_s,
                             add_labels=False,
                             extend='both',
                             add_colorbar=False)

    land.land_mask.plot.contour(x='lon',
                                y='lat',
                                levels=np.arange(0., 2., 1.),
                                ax=ax_s,
                                colors='k',
                                add_colorbar=False,
                                add_labels=False)

    cs = (vor * 86400.)[pentad:pentad + 4, :, :].sel(
        lon=lons,
        lat=lats).mean('xofyear').plot.contour(x='lon',
                                               y='lat',
                                               levels=np.arange(-14., 15., 2.),
                                               ax=ax_s,
                                               colors='0.6',
                                               add_labels=False,
                                               add_colorbar=False)
    plt.clabel(cs, fontsize=12, inline_spacing=-1, fmt='%1.0f')

    f1 = horiz_md_mean[pentad:pentad + 4, :, :].mean('xofyear').plot.contourf(
        x='lon',
        y='lat',
        levels=np.arange(-5., 5.1, 1.),
        ax=ax_h,
        add_labels=False,
        extend='both',
        add_colorbar=False)

    land.land_mask.plot.contour(x='lon',
                                y='lat',
                                levels=np.arange(0., 2., 1.),
                                ax=ax_h,
                                colors='k',
                                add_colorbar=False,
                                add_labels=False)

    cs = (vor * 86400.)[pentad:pentad + 4, :, :].sel(
        lon=lons,
        lat=lats).mean('xofyear').plot.contour(x='lon',
                                               y='lat',
                                               levels=np.arange(-14., 15., 2.),
                                               ax=ax_h,
                                               colors='0.6',
                                               add_labels=False,
                                               add_colorbar=False)
    plt.clabel(cs, fontsize=12, inline_spacing=-1, fmt='%1.0f')

    for ax in [ax_s, ax_h]:
        ax.set_xlim(60, 150)
        ax.set_ylim(-30, 60)
        ax.set_xticks(np.arange(60., 155., 30.))
        ax.set_yticks(np.arange(-30., 65., 30.))

    return f1
コード例 #26
0
def vort_eq_hm(run, lev=150, lonin=[-1., 361.]):

    rcParams['figure.figsize'] = 10, 7.5
    rcParams['font.size'] = 18
    rcParams['text.usetex'] = True

    plot_dir = '/scratch/rg419/plots/vort_schematic/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    #Load in vorticity budget term means
    data_vort = xr.open_dataset(
        '/scratch/rg419/Data_moist/climatologies/vort_eq_' + run + '.nc')
    data_vort = data_vort * 86400.**2.  #Convert to day^-2

    print 'vorticity budget data loaded'

    #Also load climatological data so that transient eddies can be calculated
    data = xr.open_dataset(
        '/scratch/rg419/Data_moist/climatologies/vort_eq_uv' + run + '.nc')
    print 'climatology loaded'

    if lonin[1] > lonin[0]:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
        ]
    else:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
        ]

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f

    abs_vort = vor.sel(lon=lons).mean('lon') * 86400.

    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)

    horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy)

    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
    stretching_mean = -86400.**2. * vor * div

    transient = data_vort.horiz_md.sel(
        pfull=lev).values + data_vort.stretching.sel(
            pfull=lev).values - horiz_md_mean.values - stretching_mean.values

    data_vort['transient'] = (('pentad', 'lat', 'lon'), transient)

    horiz_md_hm = horiz_md_mean.sel(lon=lons).mean('lon')
    stretching_hm = stretching_mean.sel(lon=lons).mean('lon')
    transient_hm = data_vort.transient.sel(lon=lons).mean('lon')
    total_hm = data_vort.total.sel(pfull=lev, lon=lons).mean('lon')

    mn_dic = month_dic(1)
    tickspace = range(13, 72, 18)
    labels = [mn_dic[(k + 5) / 6] for k in tickspace]
    levels = np.arange(-1.5, 1.6, 0.1)

    print 'starting plotting'

    # Four subplots
    f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,
                                               2,
                                               sharex='col',
                                               sharey='row')
    plt.set_cmap('RdBu_r')
    #First plot
    f2 = horiz_md_hm.plot.contourf(x='xofyear',
                                   y='lat',
                                   levels=levels,
                                   ax=ax1,
                                   extend='both',
                                   add_labels=False,
                                   add_colorbar=False)
    ax1.contour(data.xofyear,
                data.lat,
                abs_vort.T,
                levels=np.arange(-12., 13., 2.),
                colors='k',
                linewidths=2,
                alpha=0.25)
    ax1.set_ylabel('Latitude')
    ax1.set_ylim(-60, 60)
    ax1.set_yticks(np.arange(-60, 61, 30))
    ax1.grid(True, linestyle=':')
    ax1.text(-10, 60, 'a)')

    #Second plot
    stretching_hm.plot.contourf(x='xofyear',
                                y='lat',
                                levels=levels,
                                ax=ax2,
                                extend='both',
                                add_labels=False,
                                add_colorbar=False)
    ax2.contour(data.xofyear,
                data.lat,
                abs_vort.T,
                levels=np.arange(-12., 13., 2.),
                colors='k',
                linewidths=2,
                alpha=0.25)
    ax2.set_ylim(-60, 60)
    ax2.set_yticks(np.arange(-60, 61, 30))
    ax2.grid(True, linestyle=':')
    ax2.text(-5, 60, 'b)')

    #Third plot
    transient_hm.plot.contourf(x='pentad',
                               y='lat',
                               levels=levels,
                               ax=ax3,
                               extend='both',
                               add_labels=False,
                               add_colorbar=False)
    ax3.contour(data.xofyear,
                data.lat,
                abs_vort.T,
                levels=np.arange(-12., 13., 2.),
                colors='k',
                linewidths=2,
                alpha=0.25)
    ax3.set_ylabel('Latitude')
    ax3.set_ylim(-60, 60)
    ax3.set_yticks(np.arange(-60, 61, 30))
    ax3.set_xticks(tickspace)
    ax3.set_xticklabels(labels, rotation=25)
    ax3.grid(True, linestyle=':')
    ax3.text(-10, 60, 'c)')

    #Fourth plot
    total_hm.plot.contourf(x='pentad',
                           y='lat',
                           levels=levels,
                           ax=ax4,
                           extend='both',
                           add_labels=False,
                           add_colorbar=False)
    ax4.contour(data.xofyear,
                data.lat,
                abs_vort.T,
                levels=np.arange(-12., 13., 2.),
                colors='k',
                linewidths=2,
                alpha=0.25)
    ax4.set_ylim(-60, 60)
    ax4.set_yticks(np.arange(-60, 61, 30))
    ax4.set_xticks(tickspace)
    ax4.set_xticklabels(labels, rotation=25)
    ax4.grid(True, linestyle=':')
    ax4.text(-5, 60, 'd)')

    plt.subplots_adjust(right=0.97,
                        left=0.1,
                        top=0.95,
                        bottom=0.,
                        hspace=0.2,
                        wspace=0.1)
    #Colorbar
    cb1 = f.colorbar(f2,
                     ax=[ax1, ax2, ax3, ax4],
                     use_gridspec=True,
                     orientation='horizontal',
                     fraction=0.15,
                     pad=0.1,
                     aspect=30,
                     shrink=0.5)
    cb1.set_label('Vorticity tendency, day$^{-2}$')

    if lonin == [-1., 361.]:
        figname = 'vort_budg_hm_' + run + '.pdf'
    else:
        figname = 'vort_budg_hm_' + run + '_' + str(int(lonin[0])) + '_' + str(
            int(lonin[1])) + '.pdf'

    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()
コード例 #27
0
def t_dtdy_plot(run, lonin=[-1., 361.]):

    rcParams['figure.figsize'] = 6, 6
    rcParams['font.size'] = 16

    plot_dir = '/scratch/rg419/plots/surface_fluxes/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
                           '.nc')

    lons = pick_lons(data, lonin)

    try:
        precip_plot = (data.precipitation * 86400.).sel(lon=lons).mean('lon')
    except:
        precip_plot = ((data.convection_rain + data.condensation_rain) *
                       86400.).sel(lon=lons).mean('lon')

    data = data.sel(lon=lons).mean('lon')

    dTdy = gr.ddy(data.t_surf, vector=False)
    dTdy = dTdy * mc.a

    f, (ax1, ax2) = plt.subplots(2, sharex=True)

    f1 = data.t_surf.plot.contourf(x='xofyear',
                                   y='lat',
                                   levels=np.arange(270., 311., 2.5),
                                   ax=ax1,
                                   extend='both',
                                   add_labels=False)
    precip_plot.plot.contour(x='xofyear',
                             y='lat',
                             levels=np.arange(2., 15., 4.),
                             ax=ax1,
                             colors='k',
                             extend='both',
                             add_labels=False)
    ax1.set_title('SST')

    dTdy.plot.contourf(x='xofyear',
                       y='lat',
                       levels=np.arange(-45., 46., 5.),
                       ax=ax2,
                       extend='both',
                       add_labels=False)
    precip_plot.plot.contour(x='xofyear',
                             y='lat',
                             levels=np.arange(2., 15., 4.),
                             ax=ax2,
                             colors='k',
                             extend='both',
                             add_labels=False)
    ax2.set_title('dSSTdy')

    ax2.set_xlabel('Pentad')
    for ax in [ax1, ax2]:
        ax.set_ylabel('Latitude')
        ax.grid(True, linestyle=':')

    #plt.subplots_adjust(left=0.2, right=0.95, top=0.95, bottom=0., hspace=0.2)

    plt.savefig(plot_dir + 't_dtdy_plot_' + run + '.pdf', format='pdf')
    plt.close()
コード例 #28
0
def mom_budg_hm(run,
                lev=150,
                filename='plev_pentad',
                timeav='pentad',
                period_fac=1.,
                lonin=[-1., 361.]):

    a = 6376.0e3

    rcParams['figure.figsize'] = 12, 5.5
    rcParams['font.size'] = 18
    rcParams['text.usetex'] = True

    plot_dir = '/scratch/rg419/plots/itcz_theory_testing/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
                           '.nc')

    if lonin[1] > lonin[0]:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
        ]
    else:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
        ]

    #advective terms
    partition_advection(data, lons, lev=150)

    #Coriolis
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180.)
    fu = -1. * data.ucomp.sel(pfull=lev) * f * 86400.
    fu = fu.sel(lon=lons).mean('lon')

    metric_term = (-86400. * data.ucomp.sel(pfull=lev) *
                   data.ucomp.sel(pfull=lev) *
                   np.tan(data.lat * np.pi / 180.) / a).mean('lon')

    #totp = ((data.convection_rain + data.condensation_rain)*86400.).sel(lon=lons).mean('lon')
    #abs_vort = (data.vor + f).sel(lon=lons).mean('lon')*86400.

    #Geopotential gradient
    dphidy = gr.ddy(data.height.sel(pfull=lev), vector=False)
    dphidy = -86400. * 9.8 * dphidy.sel(lon=lons).mean('lon')

    mom_mean = data.v_dvdy_zav + data.w_dvdp_zav
    mom_trans = data.vv_trans_dy + data.vw_trans_dp
    mom_stat = data.v_dvdy_stat + data.w_dvdp_stat

    mom_sum = mom_mean + mom_trans + mom_stat + fu + dphidy + metric_term

    resid = -(fu + dphidy + metric_term)

    levels = np.arange(-10., 10.1, 1.)

    mn_dic = month_dic(1)
    tickspace = range(13, 72, 18)
    labels = [mn_dic[(k + 5) / 6] for k in tickspace]

    resid.plot.contourf(x='xofyear',
                        y='lat',
                        extend='both',
                        levels=levels,
                        add_labels=False)
    plt.ylabel('Latitude')
    plt.ylim(-60, 60)
    plt.xticks(tickspace, labels, rotation=25)
    plt.yticks(np.arange(-60., 61., 30.))
    plt.grid(True, linestyle=':')
    plt.tight_layout()
    plt.savefig(plot_dir + 'residual.pdf')
    plt.close()

    mom_mean.plot.contourf(x='xofyear',
                           y='lat',
                           extend='both',
                           levels=levels,
                           add_labels=False)
    plt.ylabel('Latitude')
    plt.ylim(-60, 60)
    plt.xticks(tickspace, labels, rotation=25)
    plt.yticks(np.arange(-60., 61., 30.))
    plt.grid(True, linestyle=':')
    plt.tight_layout()
    plt.savefig(plot_dir + 'mom_mean.pdf')
    plt.close()
コード例 #29
0
def vort_eq_hm(run, lev=150, lonin=[-1., 361.]):

    rcParams['figure.figsize'] = 15, 5.5
    rcParams['font.size'] = 18
    rcParams['text.usetex'] = True

    plot_dir = '/scratch/rg419/plots/paper_1_figs/revisions/'
    mkdir = sh.mkdir.bake('-p')
    mkdir(plot_dir)

    #Load in vorticity budget term means
    data_vort = xr.open_dataset(
        '/scratch/rg419/Data_moist/climatologies/vort_eq_' + run + '.nc')
    data_vort = data_vort * 86400.**2.  #Convert to day^-2

    print 'vorticity budget data loaded'

    #Also load climatological data so that transient eddies can be calculated
    data = xr.open_dataset(
        '/scratch/rg419/Data_moist/climatologies/vort_eq_uv' + run + '.nc')
    print 'climatology loaded'

    if lonin[1] > lonin[0]:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] and data.lon[i] < lonin[1]
        ]
    else:
        lons = [
            data.lon[i] for i in range(len(data.lon))
            if data.lon[i] >= lonin[0] or data.lon[i] < lonin[1]
        ]

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f

    abs_vort = vor.sel(lon=lons).mean('lon') * 86400.

    dvordx = gr.ddx(vor)
    dvordy = gr.ddy(vor, vector=False)

    horiz_md_mean = -86400.**2. * (data.ucomp * dvordx + data.vcomp * dvordy)

    div = gr.ddx(data.ucomp) + gr.ddy(data.vcomp)
    stretching_mean = -86400.**2. * vor * div

    # Zonal means
    horiz_zmean = -86400.**2. * (data.ucomp.mean('lon') * dvordx.mean('lon') +
                                 data.vcomp.mean('lon') * dvordy.mean('lon'))
    stretching_zmean = -86400.**2. * vor.mean('lon') * div.mean('lon')

    # Transient eddies
    transient_s_hm = (data_vort.stretching.sel(pfull=lev).values -
                      stretching_mean).sel(lon=lons).mean('lon')
    transient_h_hm = (data_vort.horiz_md.sel(pfull=lev).values -
                      horiz_md_mean).sel(lon=lons).mean('lon')

    transient_hm = transient_s_hm + transient_h_hm

    # Stationary eddies
    u_ed = data.ucomp - data.ucomp.mean('lon')
    v_ed = data.vcomp - data.vcomp.mean('lon')
    dvordx_ed = dvordx - dvordx.mean('lon')
    dvordy_ed = dvordy - dvordy.mean('lon')
    div_ed = div - div.mean('lon')
    vor_ed = vor - vor.mean('lon')

    stat_s = -86400.**2. * (div_ed * vor_ed)
    stat_h = -86400.**2. * (u_ed * dvordx_ed + v_ed * dvordy_ed)

    # Stationary eddy cross terms
    cross_s = -86400.**2. * (vor.mean('lon') * div_ed +
                             vor_ed * div.mean('lon'))
    cross_h = -86400.**2. * (u_ed * dvordx.mean('lon') +
                             v_ed * dvordy.mean('lon') +
                             data.ucomp.mean('lon') * dvordx_ed +
                             data.vcomp.mean('lon') * dvordy_ed)

    mn_dic = month_dic(1)
    tickspace = range(13, 72, 18)
    labels = [mn_dic[(k + 5) / 6] for k in tickspace]
    levels = np.arange(-1.5, 1.6, 0.25)

    print 'starting plotting'

    # Four subplots
    f, ((ax1, ax2, ax3), (ax4, ax5, ax6)) = plt.subplots(2,
                                                         3,
                                                         sharex='col',
                                                         sharey='row')
    plt.set_cmap('RdBu_r')
    #First plot
    f2 = horiz_zmean.plot.contourf(x='xofyear',
                                   y='lat',
                                   levels=levels,
                                   ax=ax1,
                                   extend='both',
                                   add_labels=False)
    #ax1.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
    ax1.set_ylabel('Latitude')
    ax1.set_ylim(-60, 60)
    ax1.set_yticks(np.arange(-60, 61, 30))
    ax1.grid(True, linestyle=':')
    ax1.text(-15, 60, 'a)')

    #Second plot
    stat_h.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                   y='lat',
                                                   levels=levels,
                                                   ax=ax2,
                                                   extend='both',
                                                   add_labels=False)
    ax2.set_ylim(-60, 60)
    ax2.set_yticks(np.arange(-60, 61, 30))
    ax2.grid(True, linestyle=':')
    ax2.text(-7, 60, 'b)')

    #Third plot
    cross_h.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                    y='lat',
                                                    levels=levels,
                                                    ax=ax3,
                                                    extend='both',
                                                    add_labels=False)
    ax3.set_ylim(-60, 60)
    ax3.set_yticks(np.arange(-60, 61, 30))
    ax3.grid(True, linestyle=':')
    ax3.text(-7, 60, 'c)')

    #Fourth plot
    stretching_zmean.plot.contourf(x='xofyear',
                                   y='lat',
                                   levels=levels,
                                   ax=ax4,
                                   extend='both',
                                   add_labels=False)
    #ax4.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
    ax4.set_ylabel('Latitude')
    ax4.set_ylim(-60, 60)
    ax4.set_yticks(np.arange(-60, 61, 30))
    ax4.set_xticks(tickspace)
    ax4.set_xticklabels(labels, rotation=25)
    ax4.grid(True, linestyle=':')
    ax4.text(-15, 60, 'd)')

    #Fifth plot
    stat_s.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                   y='lat',
                                                   levels=levels,
                                                   ax=ax5,
                                                   extend='both',
                                                   add_labels=False)
    ax5.set_ylim(-60, 60)
    ax5.set_yticks(np.arange(-60, 61, 30))
    ax5.set_xticks(tickspace)
    ax5.set_xticklabels(labels, rotation=25)
    ax5.grid(True, linestyle=':')
    ax5.text(-7, 60, 'e)')

    #Sixth plot
    cross_s.sel(lon=lons).mean('lon').plot.contourf(x='xofyear',
                                                    y='lat',
                                                    levels=levels,
                                                    ax=ax6,
                                                    extend='both',
                                                    add_labels=False)
    #ax4.contour(data.xofyear, data.lat, abs_vort.T, levels=np.arange(-12.,13.,2.), colors='k', linewidths=2, alpha=0.25)
    ax6.set_ylim(-60, 60)
    ax6.set_yticks(np.arange(-60, 61, 30))
    ax6.set_xticks(tickspace)
    ax6.set_xticklabels(labels, rotation=25)
    ax6.grid(True, linestyle=':')
    ax6.text(-7, 60, 'f)')

    plt.subplots_adjust(right=0.97,
                        left=0.08,
                        top=0.93,
                        bottom=0.1,
                        hspace=0.25,
                        wspace=0.15)

    if lonin == [-1., 361.]:
        figname = 'vort_eddies_' + run + '.pdf'
    else:
        figname = 'vort_eddies_' + run + '_' + str(int(lonin[0])) + '_' + str(
            int(lonin[1])) + '.pdf'

    plt.savefig(plot_dir + figname, format='pdf')
    plt.close()
コード例 #30
0
ファイル: vorticity_eq.py プロジェクト: subond/python_scripts
def vort_eq(run, months, filename='plev_daily', lev=150.):

    #Load in dataset
    name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc'
    names = [name_temp % m for m in range(months[0], months[1])]
    #read data into xarray
    data = xr.open_mfdataset(
        names,
        decode_times=False,  # no calendar so tell netcdf lib
        # choose how data will be broken down into manageable chunks.
        chunks={'time': 30})

    print 'files opened'

    # Calculate planetary vorticity
    omega = 7.2921150e-5
    f = 2 * omega * np.sin(data.lat * np.pi / 180)

    # Calculate vertical component of absolute vorticity = f + dv/dx - du/dy
    v_dx = gr.ddx(data.vcomp)  # dvdx
    u_dy = gr.ddy(data.ucomp)  # dudy
    vor = v_dx - u_dy + f
    #vor_rel = v_dx - u_dy
    print 'absolute vorticity calculated'

    dvordx = gr.ddx(vor.sel(pfull=lev))
    dvordy = gr.ddy(vor.sel(pfull=lev), vector=False)
    dvordp = gr.ddp(vor)

    # Calculate horizontal material derivative
    horiz_md = -1. * (data.ucomp * dvordx + data.vcomp * dvordy)

    # Calculate vertical material derivative
    vert_md = -1. * data.omega.sel(pfull=lev) * dvordp.sel(pfull=lev)
    print 'advective terms done'

    # Now do the terms involving gradients of windspeed
    div = gr.ddx(data.ucomp.sel(pfull=lev)) + gr.ddy(data.vcomp.sel(pfull=lev))
    stretching = -1. * vor * div
    tilting = (gr.ddp(data.ucomp)).sel(pfull=lev) * gr.ddy(
        data.omega.sel(pfull=lev), vector=False) - (gr.ddp(
            data.vcomp)).sel(pfull=lev) * gr.ddx(data.omega.sel(pfull=lev))
    print 'stretching and tilting done'

    total = horiz_md + vert_md + stretching + tilting

    ds = xr.Dataset(
        {
            'horiz_md': (['time', 'lat', 'lon'], horiz_md),
            'vert_md': (['time', 'lat', 'lon'], vert_md),
            'stretching': (['time', 'lat', 'lon'], stretching),
            'tilting': (['time', 'lat', 'lon'], tilting),
            'total': (['time', 'lat', 'lon'], total)
        },
        coords={
            'time': ('time', data.time),
            'pfull': ('pfull', data.pfull),
            'lat': ('lat', data.lat),
            'lon': ('lon', data.lon)
        })

    ds.coords['xofyear'] = np.mod(ds.time - 1., 360.) // 5 + 1.
    dsout = ds.groupby('xofyear').mean(('time'))

    GFDL_DATA = os.environ['GFDL_DATA']
    fileout = GFDL_DATA + 'climatologies/' + run + '_' + str(
        int(lev)) + '_vort_eq.nc'
    dsout.to_netcdf(path=fileout)

    print 'data written to ', fileout

    return dsout