Example #1
0
def testing(data, lat1, lat2, lon1, lon2, t, k):

    latname = get_coord(data, 'lat', 'name')
    lonname = get_coord(data, 'lon', 'name')

    data_sub = dat.subset(data, latname, lat1, lat2, lonname, lon1, lon2)

    plt.figure()
    ap.pcolor_latlon(data_sub[t, k],
                     axlims=(lat1, lat2, lon1, lon2),
                     cmap='jet')

    avg0 = data_sub.mean(axis=-1).mean(axis=-1)
    avg1 = mean_over_geobox(data, lat1, lat2, lon1, lon2, area_wtd=False)
    avg2 = mean_over_geobox(data, lat1, lat2, lon1, lon2, area_wtd=True)
    avg3 = mean_over_geobox(data,
                            lat1,
                            lat2,
                            lon1,
                            lon2,
                            area_wtd=True,
                            land_only=True)

    print(avg0[t, k].values)
    print(avg1[t, k].values)
    print(avg2[t, k].values)
    print(avg3[t, k].values)
def testing(data, lat1, lat2, lon1, lon2, t, k):

    latname = get_coord(data, 'lat', 'name')
    lonname = get_coord(data, 'lon', 'name')

    data_sub = dat.subset(data, latname, lat1, lat2, lonname, lon1, lon2)

    plt.figure()
    ap.pcolor_latlon(data_sub[t,k], axlims=(lat1,lat2,lon1,lon2), cmap='jet')

    avg0 = data_sub.mean(axis=-1).mean(axis=-1)
    avg1 = mean_over_geobox(data, lat1, lat2, lon1, lon2, area_wtd=False)
    avg2 = mean_over_geobox(data, lat1, lat2, lon1, lon2, area_wtd=True)
    avg3 = mean_over_geobox(data, lat1, lat2, lon1, lon2, area_wtd=True,
                            land_only=True)

    print(avg0[t, k].values)
    print(avg1[t, k].values)
    print(avg2[t, k].values)
    print(avg3[t, k].values)
Example #3
0
u = ds['u']
v = ds['v']
T = ds['T']
ps = ds['ps']
lat = get_coord(ps, 'lat')
lon = get_coord(ps, 'lon')
plev = get_coord(u, 'plev')

# ----------------------------------------------------------------------
# interp_latlon

lat_new = np.arange(-90, 90, 1.)
lon_new = np.arange(0, 360, 1.)

# DataArray input
ps_i = interp_latlon(ps, lat_new, lon_new)

# ndarray input
ps_i2 = interp_latlon(ps.values, lat_new, lon_new, lat, lon)

t = 0
cmap = 'jet'

plt.figure(figsize=(7, 10))
plt.subplot(311)
ap.pcolor_latlon(ps[t], cmap=cmap)
plt.subplot(312)
ap.pcolor_latlon(ps_i[t], cmap=cmap)
plt.subplot(313)
ap.pcolor_latlon(ps_i2[t], lat_new, lon_new, cmap=cmap)
Example #4
0
# ndarray
mfc2 = moisture_flux_conv(uq.values, vq.values, lat, lon, plev * 100)

# components
mfc3, mfc_x, mfc_y, uq_int, vq_int = moisture_flux_conv(uq,
                                                        vq,
                                                        return_comp=True)

# pmin, pmax
mfc_sub = moisture_flux_conv(uq, vq, pmin=600e2, pmax=900e2)

cmax = 12
plt.figure(figsize=(7, 10))
plt.subplot(311)
ap.pcolor_latlon(mfc)
plt.clim(-cmax, cmax)
plt.subplot(312)
ap.pcolor_latlon(mfc2, lat, lon)
plt.clim(-cmax, cmax)
plt.subplot(313)
ap.pcolor_latlon(mfc_sub)
plt.clim(-cmax, cmax)

plt.figure(figsize=(7, 8))
plt.subplot(211)
ap.pcolor_latlon(mfc_x)
plt.clim(-cmax, cmax)
plt.subplot(212)
ap.pcolor_latlon(mfc_y)
plt.clim(-cmax, cmax)
lon = get_coord(T, 'lon')
plev = get_coord(T, 'plev')
pname = get_coord(T, 'plev', 'name')
units_in = dat.pres_units(T[pname].units)
plev = dat.pres_convert(plev, units_in, 'Pa')
p0 = 1e5

print('Calculating potential temperature')
theta = av.potential_temp(T, plev, p0)
print('Calculating equivalent potential temperature')
theta_e = av.equiv_potential_temp(T, plev, q, p0)

# t, k = 0, 6
# t, k = 0, 22
t, k = 0, 14
pstr = '%d hPa' % (plev[k]/100)
plt.figure(figsize=(7,10))
plt.subplot(311)
ap.pcolor_latlon(T[t,k])
plt.title('Temperature ' + pstr)
plt.subplot(312)
ap.pcolor_latlon(theta[t,k])
plt.title('Potential Temperature ' + pstr)
plt.subplot(313)
ap.pcolor_latlon(theta_e[t,k])
plt.title('Equiv Potential Temperature ' + pstr)

ps.dims
plt.figure()
ap.pcolor_latlon(ps,cmap='hot')
x1, y1 = np.meshgrid(lon_m, lat_m)
x2, y2 = np.meshgrid(ps_m2.XDim, ps_m2.YDim)
cmap='jet'
ncontours = 30

plt.figure(figsize=(7,8))
plt.subplot(211)
plt.contourf(x1, y1, ps_m[t].values, ncontours, cmap=cmap)
plt.colorbar()
plt.subplot(212)
plt.contourf(x2, y2, ps_m2[t].values, ncontours, cmap=cmap)
plt.colorbar()

plt.figure(figsize=(7,8))
plt.subplot(211)
ap.pcolor_latlon(ps_m[t], cmap=cmap)
plt.subplot(212)
ap.pcolor_latlon(ps_m2[t], cmap=cmap)

# # ----------------------------------------------------------------------
# # Correct for topography
#
# u_orig = u
# u = dat.correct_for_topography(u_orig, topo)
#
# # ----------------------------------------------------------------------
# # Interpolate onto new pressure grid
#
# plev_new = np.arange(1000,0,-50.)
# u_i = dat.interp_plevels(u, plev_new, pdim=-3)
#
ps = ds['ps']

topo = get_ps_clim(lat, lon)
topo.values /= 100
topo.attrs['units'] = 'hPa'

# ----------------------------------------------------------------------
# Correct for topography

u_orig = u
u = correct_for_topography(u_orig, topo)

m, k = 3, 1
plt.figure(figsize=(7,8))
plt.subplot(211)
ap.pcolor_latlon(u_orig[m,k], cmap='jet')
plt.subplot(212)
ap.pcolor_latlon(u[m,k], cmap='jet')


# ----------------------------------------------------------------------
# Zonal mean zonal wind
season='jjas'
lon1, lon2 = 60, 100
cint = 5
months = utils.season_months(season)

uplot = dat.subset(u, 'lon', lon1, lon2, 'mon', months)
uplot = uplot.mean(['lon', 'mon'])

ps_plot = dat.subset(topo, 'lon', lon1, lon2)
ps_m2 = set_lon(ps_m, lonmax=360)

# ndarray as input
ps_m3, lon_out = set_lon(ps_m.values, lonmax=360, lon=lon_m)

t = 0
x1, y1 = np.meshgrid(lon_m, lat_m)
x2, y2 = np.meshgrid(ps_m2.XDim, ps_m2.YDim)
x3, y3 = np.meshgrid(lon_out, lat_m)
cmap='jet'
ncontours = 30

plt.figure(figsize=(7,10))
plt.subplot(311)
plt.contourf(x1, y1, ps_m[t].values, ncontours, cmap=cmap)
plt.colorbar()
plt.subplot(312)
plt.contourf(x2, y2, ps_m2[t].values, ncontours, cmap=cmap)
plt.colorbar()
plt.subplot(313)
plt.contourf(x3, y3, ps_m3[t], ncontours, cmap=cmap)
plt.colorbar()

plt.figure(figsize=(7,10))
plt.subplot(311)
ap.pcolor_latlon(ps_m[t], cmap=cmap)
plt.subplot(312)
ap.pcolor_latlon(ps_m2[t], cmap=cmap)
plt.subplot(313)
ap.pcolor_latlon(ps_m3[t], lat_m, lon_out, cmap=cmap)
topo.units = 'hPa'

# ----------------------------------------------------------------------
# Relative and absolute vorticity

# DataArray
rel_vort, abs_vort, f = vorticity(u, v)

# ndarray
rel_vort2, abs_vort2, f2 = vorticity(u.values, v.values, lat, lon)


t, k = 0, 22
plt.figure(figsize=(12,8))
plt.subplot(221)
ap.pcolor_latlon(rel_vort[t,k])
plt.subplot(222)
ap.pcolor_latlon(abs_vort[t,k])
plt.subplot(223)
ap.pcolor_latlon(rel_vort2[t,k], lat, lon)
plt.subplot(224)
ap.pcolor_latlon(abs_vort2[t,k], lat, lon)

# ----------------------------------------------------------------------
# Rossby number

Ro = rossby_num(u, v)
Ro2 = rossby_num(u.values, v.values, lat, lon)

t, k = 0, 22
lon1, lon2 = 60, 100
Example #10
0
ps = ds['ps']

topo = dat.get_ps_clim(lat, lon) / 100
topo.units = 'hPa'

# ----------------------------------------------------------------------
# Correct for topography

u_orig = u
u = dat.correct_for_topography(u_orig, topo)

# ----------------------------------------------------------------------
# Near-surface data

# DataArray
u_s, ind_s = near_surface(u, pdim=-3, return_inds=True)

# ndarray
u_s2 = near_surface(u.values, pdim=-3, return_inds=False)

m = 0
plt.figure(figsize=(12, 10))
plt.subplot(221)
ap.pcolor_latlon(u[m, 0], cmap='jet')
plt.subplot(222)
ap.pcolor_latlon(u_s[m], cmap='jet')
plt.subplot(223)
ap.pcolor_latlon(u_s2[m], lat, lon, cmap='jet')
plt.subplot(224)
ap.pcolor_latlon(ind_s[m], lat, lon, cmap='jet')
lon = get_coord(T, 'lon')
plev = get_coord(T, 'plev')
pname = get_coord(T, 'plev', 'name')
units_in = dat.pres_units(T[pname].units)
plev = dat.pres_convert(plev, units_in, 'Pa')
p0 = 1e5

print('Calculating potential temperature')
theta = av.potential_temp(T, plev, p0)
print('Calculating equivalent potential temperature')
theta_e = av.equiv_potential_temp(T, plev, q, p0)

# t, k = 0, 6
# t, k = 0, 22
t, k = 0, 14
pstr = '%d hPa' % (plev[k] / 100)
plt.figure(figsize=(7, 10))
plt.subplot(311)
ap.pcolor_latlon(T[t, k])
plt.title('Temperature ' + pstr)
plt.subplot(312)
ap.pcolor_latlon(theta[t, k])
plt.title('Potential Temperature ' + pstr)
plt.subplot(313)
ap.pcolor_latlon(theta_e[t, k])
plt.title('Equiv Potential Temperature ' + pstr)

ps.dims
plt.figure()
ap.pcolor_latlon(ps, cmap='hot')
# ----------------------------------------------------------------------
# Initial plots

ubar = u.mean(axis=-1)
vbar = v.mean(axis=-1)

plt.figure(figsize=(7,8))
plt.subplot(211)
ap.contour_latpres(ubar, clev=5)
plt.subplot(212)
ap.contour_latpres(vbar, clev=0.5)

t, k = 0, 4
plt.figure(figsize=(7,8))
plt.subplot(211)
ap.pcolor_latlon(u[t,k])
plt.subplot(212)
ap.pcolor_latlon(v[t,k])

# ----------------------------------------------------------------------
# Moisture fluxes

uq = u * q
vq = v * q

uq_int = dat.int_pres(uq, pdim=-3)
vq_int = dat.int_pres(vq, pdim=-3)

t, k = 0, 4
plt.figure(figsize=(7,8))
plt.subplot(211)
# ----------------------------------------------------------------------
# mask_oceans

# DataArray input
ps_land = mask_oceans(ps)

# ndarray input
ps_land2 = mask_oceans(ps, lat, lon)

t = 0
cmap='jet'

# pcolor
plt.figure(figsize=(7,10))
plt.subplot(311)
ap.pcolor_latlon(ps[t], cmap=cmap)
plt.subplot(312)
ap.pcolor_latlon(ps_land[t], cmap=cmap)
plt.subplot(313)
ap.pcolor_latlon(ps_land2[t], lat, lon, cmap=cmap)

# contourf, contour
ps_plot = ps[0]/100
cint = 100
plt.figure(figsize=(7,8))
plt.subplot(211)
ap.contourf_latlon(ps_plot, clev=cint, cmap=cmap, symmetric=False)
plt.subplot(212)
ap.contour_latlon(ps_plot, clev=cint)
# ----------------------------------------------------------------------
# mask_oceans

# DataArray input
ps_land = mask_oceans(ps)

# ndarray input
ps_land2 = mask_oceans(ps, lat, lon)

t = 0
cmap = 'jet'

# pcolor
plt.figure(figsize=(7, 10))
plt.subplot(311)
ap.pcolor_latlon(ps[t], cmap=cmap)
plt.subplot(312)
ap.pcolor_latlon(ps_land[t], cmap=cmap)
plt.subplot(313)
ap.pcolor_latlon(ps_land2[t], lat, lon, cmap=cmap)

# contourf, contour
ps_plot = ps[0] / 100
cint = 100
plt.figure(figsize=(7, 8))
plt.subplot(211)
ap.contourf_latlon(ps_plot, clev=cint, cmap=cmap, symmetric=False)
plt.subplot(212)
ap.contour_latlon(ps_plot, clev=cint)
Example #15
0
x1, y1 = np.meshgrid(lon_m, lat_m)
x2, y2 = np.meshgrid(ps_m2.XDim, ps_m2.YDim)
cmap = 'jet'
ncontours = 30

plt.figure(figsize=(7, 8))
plt.subplot(211)
plt.contourf(x1, y1, ps_m[t].values, ncontours, cmap=cmap)
plt.colorbar()
plt.subplot(212)
plt.contourf(x2, y2, ps_m2[t].values, ncontours, cmap=cmap)
plt.colorbar()

plt.figure(figsize=(7, 8))
plt.subplot(211)
ap.pcolor_latlon(ps_m[t], cmap=cmap)
plt.subplot(212)
ap.pcolor_latlon(ps_m2[t], cmap=cmap)

# # ----------------------------------------------------------------------
# # Correct for topography
#
# u_orig = u
# u = dat.correct_for_topography(u_orig, topo)
#
# # ----------------------------------------------------------------------
# # Interpolate onto new pressure grid
#
# plev_new = np.arange(1000,0,-50.)
# u_i = dat.interp_plevels(u, plev_new, pdim=-3)
#
Example #16
0
ps_m2 = set_lon(ps_m, lonmax=360)

# ndarray as input
ps_m3, lon_out = set_lon(ps_m.values, lonmax=360, lon=lon_m)

t = 0
x1, y1 = np.meshgrid(lon_m, lat_m)
x2, y2 = np.meshgrid(ps_m2.XDim, ps_m2.YDim)
x3, y3 = np.meshgrid(lon_out, lat_m)
cmap = 'jet'
ncontours = 30

plt.figure(figsize=(7, 10))
plt.subplot(311)
plt.contourf(x1, y1, ps_m[t].values, ncontours, cmap=cmap)
plt.colorbar()
plt.subplot(312)
plt.contourf(x2, y2, ps_m2[t].values, ncontours, cmap=cmap)
plt.colorbar()
plt.subplot(313)
plt.contourf(x3, y3, ps_m3[t], ncontours, cmap=cmap)
plt.colorbar()

plt.figure(figsize=(7, 10))
plt.subplot(311)
ap.pcolor_latlon(ps_m[t], cmap=cmap)
plt.subplot(312)
ap.pcolor_latlon(ps_m2[t], cmap=cmap)
plt.subplot(313)
ap.pcolor_latlon(ps_m3[t], lat_m, lon_out, cmap=cmap)
ps = ds['ps']

topo = dat.get_ps_clim(lat, lon) / 100
topo.units = 'hPa'

# ----------------------------------------------------------------------
# Correct for topography

u_orig = u
u = dat.correct_for_topography(u_orig, topo)

# ----------------------------------------------------------------------
# Near-surface data

# DataArray
u_s, ind_s = near_surface(u, pdim=-3, return_inds=True)

# ndarray
u_s2 = near_surface(u.values, pdim=-3, return_inds=False)

m = 0
plt.figure(figsize=(12,10))
plt.subplot(221)
ap.pcolor_latlon(u[m,0], cmap='jet')
plt.subplot(222)
ap.pcolor_latlon(u_s[m], cmap='jet')
plt.subplot(223)
ap.pcolor_latlon(u_s2[m], lat, lon, cmap='jet')
plt.subplot(224)
ap.pcolor_latlon(ind_s[m], lat, lon, cmap='jet')
u = ds['u']
v = ds['v']
T = ds['T']
ps = ds['ps']
lat = get_coord(ps, 'lat')
lon = get_coord(ps, 'lon')
plev = get_coord(u, 'plev')

# ----------------------------------------------------------------------
# interp_latlon

lat_new = np.arange(-90, 90, 1.)
lon_new = np.arange(0, 360, 1.)

# DataArray input
ps_i = interp_latlon(ps, lat_new, lon_new)

# ndarray input
ps_i2 = interp_latlon(ps.values, lat_new, lon_new, lat, lon)

t = 0
cmap='jet'

plt.figure(figsize=(7,10))
plt.subplot(311)
ap.pcolor_latlon(ps[t], cmap=cmap)
plt.subplot(312)
ap.pcolor_latlon(ps_i[t], cmap=cmap)
plt.subplot(313)
ap.pcolor_latlon(ps_i2[t], lat_new, lon_new, cmap=cmap)
Example #19
0
topo = dat.get_ps_clim(lat, lon) / 100
topo.units = 'hPa'

# ----------------------------------------------------------------------
# Relative and absolute vorticity

# DataArray
rel_vort, abs_vort, f = vorticity(u, v)

# ndarray
rel_vort2, abs_vort2, f2 = vorticity(u.values, v.values, lat, lon)

t, k = 0, 22
plt.figure(figsize=(12, 8))
plt.subplot(221)
ap.pcolor_latlon(rel_vort[t, k])
plt.subplot(222)
ap.pcolor_latlon(abs_vort[t, k])
plt.subplot(223)
ap.pcolor_latlon(rel_vort2[t, k], lat, lon)
plt.subplot(224)
ap.pcolor_latlon(abs_vort2[t, k], lat, lon)

# ----------------------------------------------------------------------
# Rossby number

Ro = rossby_num(u, v)
Ro2 = rossby_num(u.values, v.values, lat, lon)

t, k = 0, 22
lon1, lon2 = 60, 100
Example #20
0
ps = ds['ps']

topo = get_ps_clim(lat, lon)
topo.values /= 100
topo.attrs['units'] = 'hPa'

# ----------------------------------------------------------------------
# Correct for topography

u_orig = u
u = correct_for_topography(u_orig, topo)

m, k = 3, 1
plt.figure(figsize=(7, 8))
plt.subplot(211)
ap.pcolor_latlon(u_orig[m, k], cmap='jet')
plt.subplot(212)
ap.pcolor_latlon(u[m, k], cmap='jet')

# ----------------------------------------------------------------------
# Zonal mean zonal wind
season = 'jjas'
lon1, lon2 = 60, 100
cint = 5
months = utils.season_months(season)

uplot = dat.subset(u, 'lon', lon1, lon2, 'mon', months)
uplot = uplot.mean(['lon', 'mon'])

ps_plot = dat.subset(topo, 'lon', lon1, lon2)
ps_plot = ps_plot.mean('lon')
mfc = moisture_flux_conv(uq,vq)

# ndarray
mfc2 = moisture_flux_conv(uq.values, vq.values, lat, lon, plev*100)

# components
mfc3, mfc_x, mfc_y, uq_int, vq_int = moisture_flux_conv(
    uq, vq, return_comp=True)

# pmin, pmax
mfc_sub = moisture_flux_conv(uq, vq, pmin=600e2, pmax=900e2)

cmax = 12
plt.figure(figsize=(7,10))
plt.subplot(311)
ap.pcolor_latlon(mfc)
plt.clim(-cmax, cmax)
plt.subplot(312)
ap.pcolor_latlon(mfc2, lat, lon)
plt.clim(-cmax, cmax)
plt.subplot(313)
ap.pcolor_latlon(mfc_sub)
plt.clim(-cmax, cmax)

plt.figure(figsize=(7,8))
plt.subplot(211)
ap.pcolor_latlon(mfc_x)
plt.clim(-cmax, cmax)
plt.subplot(212)
ap.pcolor_latlon(mfc_y)
plt.clim(-cmax, cmax)