def streamfun_vid(run, ax_in, pentad, plot_land=True, tibet=True):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
uwnd = data.ucomp.sel(pfull=150.)
vwnd = data.vcomp.sel(pfull=150.)
# Create a VectorWind instance to handle the computation
w = VectorWind(uwnd, vwnd)
# Compute variables
streamfun, vel_pot = w.sfvp()
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.
]
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc'
land = xr.open_dataset(land_file)
streamfun_zanom = (streamfun - streamfun.mean('lon')) / 10.**6
f1 = streamfun_zanom[pentad, :, :].plot.contourf(x='lon',
y='lat',
ax=ax_in,
levels=np.arange(
-36., 37., 3.),
add_labels=False,
add_colorbar=False,
extend='both')
if plot_land:
land.land_mask.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
ax=ax_in,
colors='k',
add_colorbar=False,
add_labels=False)
if tibet:
land.zsurf.plot.contourf(x='lon',
y='lat',
ax=ax_in,
levels=np.arange(2500., 100001., 97500.),
add_labels=False,
extend='neither',
add_colorbar=False,
alpha=0.5,
cmap='Greys_r')
ax_in.set_xlim(60, 150)
ax_in.set_ylim(-30, 60)
ax_in.set_xticks(np.arange(60., 155., 30.))
ax_in.set_yticks(np.arange(-30., 65., 30.))
ax_in.grid(True, linestyle=':')
return f1
def walker_hm(regions=[[0,10], [10,20], [20,30], [30,40]]):
data = xr.open_dataset('/scratch/rg419/obs_and_reanalysis/era_v_clim_alllevs.nc' )
data_u = xr.open_dataset('/scratch/rg419/obs_and_reanalysis/era_u_clim_alllevs.nc' )
# Take pentad means
data.coords['pentad'] = data.day_of_yr //5 + 1.
data = data.groupby('pentad').mean(('day_of_yr'))
data_u.coords['pentad'] = data_u.day_of_yr //5 + 1.
data_u = data_u.groupby('pentad').mean(('day_of_yr'))
plot_dir = '/scratch/rg419/plots/overturning_monthly/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
# Create a VectorWind instance to handle the computation
w = VectorWind(data_u.u.sel(pfull=np.arange(50.,950.,50.)), data.v.sel(pfull=np.arange(50.,950.,50.)))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
# Set figure parameters
rcParams['figure.figsize'] = 10, 7
rcParams['font.size'] = 14
# Start figure with 4 subplots
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2, 2, sharex='col', sharey='row')
axes = [ax1,ax2,ax3,ax4]
i=0
for ax in axes:
psi_w = walker_cell(uchi, latin=regions[i], dp_in=-50.)
psi_w /= 1.e9
i=i+1
f1=psi_w.sel(pfull=500).plot.contourf(ax=ax, x='lon', y='pentad', add_labels=False, add_colorbar=False, levels=np.arange(-300.,301.,50.), extend='both')
ax1.set_title('0-10 N')
ax2.set_title('10-20 N')
ax3.set_title('20-30 N')
ax4.set_title('30-40 N')
for ax in [ax1,ax2,ax3,ax4]:
ax.grid(True,linestyle=':')
ax.set_yticks(np.arange(0.,73., 18.))
ax.set_xticks([0.,90.,180.,270.,360.])
ax3.set_xlabel('Longitude')
ax4.set_xlabel('Longitude')
ax1.set_ylabel('Pentad')
ax3.set_ylabel('Pentad')
plt.subplots_adjust(left=0.1, right=0.97, top=0.95, bottom=0.1, hspace=0.3, wspace=0.3)
cb1=fig.colorbar(f1, ax=axes, use_gridspec=True, orientation = 'horizontal',fraction=0.05, pad=0.1, aspect=30, shrink=0.5)
cb1.set_label('Zonal overturning streamfunction')
# Save as a pdf
plt.savefig(plot_dir + 'walker_cell_hm_era.pdf', format='pdf')
plt.close()
data.close()
def walker_cell_monthly(run, latin=None):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
plot_dir = '/scratch/rg419/plots/overturning_monthly/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data.coords['month'] = (data.xofyear - 1) // 6 + 1
data = data.groupby('month').mean(('xofyear'))
# Create a VectorWind instance to handle the computation, and compute variables
w = VectorWind(data.ucomp.sel(pfull=np.arange(50., 950., 50.)),
data.vcomp.sel(pfull=np.arange(50., 950., 50.)))
uchi, vchi, upsi, vpsi = w.helmholtz()
psi_w = walker_cell(uchi, latin=latin, dp_in=-50.)
psi_w /= 1.e9
# Set figure parameters
rcParams['figure.figsize'] = 12, 7
rcParams['font.size'] = 14
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3, 4)
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
psi_w.sel(month=i + 1).plot.contour(ax=ax,
x='lon',
y='pfull',
yincrease=False,
levels=np.arange(0., 301, 50.),
colors='k',
add_labels=False)
psi_w.sel(month=i + 1).plot.contour(ax=ax,
x='lon',
y='pfull',
yincrease=False,
levels=np.arange(-300., 0., 50.),
colors='k',
linestyles='dashed',
add_labels=False)
i = i + 1
ax.set_xticks(np.arange(0., 361., 120.))
ax.set_yticks(np.arange(0., 1001., 250.))
ax.grid(True, linestyle=':')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
plt.savefig(plot_dir + 'walker_' + run + '.pdf', format='pdf')
plt.close()
def test_truncation(da_urho,
da_vrho,
trunc_vals=[10, 20, 30, 40, 50, 100, 200]):
for idx, trunc in enumerate(trunc_vals):
print(f'Done {int(100*idx/da_urho.time.values.shape[0])} %')
da_urho_ = da_urho.isel(time=slice(0, 8))
da_vrho_ = da_vrho.isel(time=slice(0, 8))
w = VectorWind(da_urho_, da_vrho_)
da_urho_ = w.truncate(da_urho_, truncation=trunc)
da_vrho_ = w.truncate(da_vrho_, truncation=trunc)
lcs = LCS(timestep=-6 * 3600, timedim='time', SETTLS_order=4)
ftle = lcs(u=da_urho_,
v=da_vrho_,
isglobal=True,
s=4e6,
interp_to_common_grid=True,
traj_interp_order=3)
ftle = .5 * np.log(ftle)
toplot = ftle
fig, ax = plt.subplots(1,
1,
subplot_kw={'projection': ccrs.Robinson()})
toplot.plot(ax=ax,
transform=ccrs.PlateCarree(),
robust=True,
cbar_kwargs={'shrink': .6},
vmin=.5,
vmax=2)
ax.coastlines()
ax.set_title(f'N96 truncated at t{trunc}')
plt.savefig(f'upscale/figs/trunc_test/{trunc:03d}.png',
dpi=600,
boundary='trim',
bbox_inches='tight')
plt.close()
def walker_strength(data, lonin=[90.,180.], psi_min=True, sanity_check=False):
'''
data - a climatology from Isca with dimensions including xofyear
lonin - longitudes over which to look for min/max values
psi_min - if true look for minimum of walker streamfunction, otherwise maximum
sanity_check - produce a plot on screen to check the max/min has been identified correctly
'''
lons = np.arange(lonin[0], lonin[1], 0.1) # create hi res longitude array to interpolate to
# Create a VectorWind instance, calculate uchi, and use walker_cell to get the streamfunction from this.
w = VectorWind(data.ucomp.sel(pfull=np.arange(50.,950.,50.)), data.vcomp.sel(pfull=np.arange(50.,950.,50.)))
uchi, vchi, upsi, vpsi = w.helmholtz()
psi_w = walker_cell(uchi, latin=[-20,20], dp_in=-50.)
# Divide by 1e9 and select 500 hPa level
psi_w /= 1.e9
psi_w = psi_w.sel(pfull=500.)
# interpolate to hi res lons
f = spint.interp1d(psi_w.lon, psi_w, axis=1, fill_value='extrapolate')
psi_w = f(lons)
psi_w = xr.DataArray(psi_w, coords=[uchi.xofyear, lons], dims=['xofyear', 'lon'])
# mask to include only pos/neg values
psi_mask = np.ones(psi_w.values.shape)
if psi_min:
psi_mask[psi_w > 0.] = np.nan
else:
psi_mask[psi_w < 0.] = np.nan
psi_masked = psi_w * psi_mask
# At some times there will be no pos/neg value. Find times where there are lons with non-NaN values
times_defd = []
for i in range(0, len(psi_masked.xofyear)):
if np.any(np.isfinite(psi_masked[i,:])):
times_defd.append(np.float(psi_masked.xofyear[i]))
# reduce masked array times to include only non-Nan values
psi_red = psi_masked.sel(xofyear=times_defd)
# Find min/max and its location
if psi_min:
walker_mag = psi_red.min('lon')
walker_mag_loc = psi_red.lon.values[psi_red.argmin('lon').values]
else:
walker_mag = psi_red.max('lon')
walker_mag_loc = psi_red.lon.values[psi_red.argmax('lon').values]
walker_mag_loc = xr.DataArray(walker_mag_loc, coords=[psi_red.xofyear], dims=['xofyear'])
# plot the streamfunction and the location of the maximum as a sanity check
if sanity_check:
psi_w.plot.contourf(x='xofyear',y='lon')
walker_mag_loc.plot()
#plt.figure(2)
#plt.plot(walker_mag_loc,walker_mag)
plt.show()
return walker_mag
def calculate_streamfunction(uwnd, vwnd):
"""Calculate the Streamfunction
"""
uwnd_ds = xr.open_dataarray(uwnd)
vwnd_ds = xr.open_dataarray(vwnd)
wind = VectorWind(uwnd_ds, vwnd_ds)
psi = wind.streamfunction()
return psi
def cross_prod(run,
months,
filename='plev_pentad',
timeav='month',
period_fac=1.,
land_mask=False,
level=9):
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/cross_prod/' + str(
level) + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
mn_dic = month_dic(1)
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
uwnd = data.ucomp[:, level, :, :]
vwnd = data.vcomp[:, level, :, :]
w = VectorWind(uwnd, vwnd)
uchi, vchi, upsi, vpsi = w.helmholtz()
cross_prod = (upsi * vchi - vpsi * uchi)
for i in range(0, 12):
ax = cross_prod[i, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(
-200., 201., 10.),
extend='both',
add_colorbar=False,
add_label=False)
cb1 = plt.colorbar(ax)
cb1.set_label('Vpsi x Vchi')
if land_mask:
land = xr.open_dataset(
'/scratch/rg419/GFDL_model/GFDLmoistModel/input/land.nc')
land_plot = xr.DataArray(land.land_mask.values,
[('lat', data.lat), ('lon', data.lon)])
land_plot.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
colors='0.75',
add_colorbar=False,
add_labels=False)
plt.ylabel('Latitude')
plt.xlabel('Longitude')
plt.ylim(-10, 45)
plt.xlim(25, 150)
plt.tight_layout()
plt.savefig(plot_dir + str(i + 1) + '_' + str(mn_dic[i + 1]) + '.png')
plt.close()
def ke_partition(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.):
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
totp = (data.convection_rain + data.condensation_rain) * 86400.
cell_ar = cell_area(42, '/scratch/rg419/GFDL_model/GFDLmoistModel/')
cell_ar = xr.DataArray(cell_ar, [('lat', data.lat), ('lon', data.lon)])
uwnd = data.ucomp
vwnd = data.vcomp
w = VectorWind(uwnd, vwnd)
uchi, vchi, upsi, vpsi = w.helmholtz()
lats = [
i for i in range(len(data.lat))
if data.lat[i] >= 5. and data.lat[i] < 30.
]
lons = [
i for i in range(len(data.lon))
if data.lon[i] >= 60. and data.lon[i] < 150.
]
ke_chi = 0.5 * (uchi * uchi + vchi * vchi) * cell_ar
ke_chi_av = ke_chi[:, :, lats, lons].sum('lat').sum('lon') / cell_ar[
lats, lons].sum('lat').sum('lon')
ke_psi = 0.5 * (upsi * upsi + vpsi * vpsi) * cell_ar
ke_psi_av = ke_psi[:, :, lats, lons].sum('lat').sum('lon') / cell_ar[
lats, lons].sum('lat').sum('lon')
totp_av = (totp[:, lats, lons] * cell_ar[lats, lons]).sum('lat').sum(
'lon') / cell_ar[lats, lons].sum('lat').sum('lon')
ke_chi_av[:, 36].plot()
ke_psi_av[:, 36].plot()
plt.legend(['KE_chi', 'KE_psi'])
plt.xlabel('Pentad')
plt.ylabel('Kinetic Energy, m2/s2')
plt.savefig(plot_dir + 'KE_' + run + '.png')
plt.close()
totp_av.plot()
plt.xlabel('Pentad')
plt.ylabel('Precipitation, mm/day')
plt.savefig(plot_dir + 'precip_' + run + '.png')
plt.close()
def mass_streamfunction(data,
a=6376.0e3,
g=9.8,
lons=[-1000],
dp_in=0.,
intdown=True,
use_v_locally=False):
"""Calculate the mass streamfunction for the atmosphere.
Based on a vertical integral of the meridional wind.
Ref: Physics of Climate, Peixoto & Oort, 1992. p158.
`a` is the radius of the planet (default Isca value 6376km).
`g` is surface gravity (default Earth 9.8m/s^2).
lons allows a local area to be used by specifying boundaries as e.g. [70,100]
dp_in - if no phalf and if using regularly spaced pressure levels, use this increment for integral. Units hPa.
intdown - choose integratation direction (i.e. from surface to TOA, or TOA to surface).
Returns an xarray DataArray of mass streamfunction.
"""
if lons[0] == -1000: #Use large negative value to use all data if no lons provided
vbar = data.vcomp.mean('lon')
elif use_v_locally:
vbar = data.vcomp.sel(lon=lons).mean('lon').sortby('pfull',
ascending=False)
else:
from windspharm.xarray import VectorWind
# Create a VectorWind instance to handle the computation
w = VectorWind(data.ucomp.sel(pfull=np.arange(50., 950., 50.)),
data.vcomp.sel(pfull=np.arange(50., 950., 50.)))
# Compute variables
uchi, vbar, upsi, vpsi = w.helmholtz()
vbar = vbar.sel(lon=lons).mean('lon').sortby('pfull', ascending=False)
c = 2 * np.pi * a * np.cos(vbar.lat * np.pi / 180) / g
# take a diff of half levels, and assign to pfull coordinates
if dp_in == 0.:
dp = xr.DataArray(data.phalf.diff('phalf').values * 100.,
coords=[('pfull', data.pfull)])
return c * np.cumsum(vbar * dp, axis=vbar.dims.index('pfull'))
else:
dp = dp_in * 100.
if intdown:
psi = c * dp * np.flip(vbar, axis=vbar.dims.index('pfull')).cumsum(
'pfull') #[:,:,::-1]
else:
psi = -1. * c * dp * vbar.cumsum('pfull') #[:,:,::-1]
#psi.plot.contourf(x='lat', y='pfull', yincrease=False)
#plt.show()
return psi #c*dp*vbar[:,::-1,:].cumsum('pfull')#[:,:,::-1]
def calc_ftle(da_urho_, da_vrho_, truncation=20):
print('********************************** \n'
'Truncating wind and computing FTLE\n'
'**********************************')
w = VectorWind(da_urho_, da_vrho_)
da_urho_ = w.truncate(da_urho_, truncation=20)
da_vrho_ = w.truncate(da_vrho_, truncation=20)
lcs = LCS(timestep=-6 * 3600, timedim='time', SETTLS_order=4)
ftle = lcs(u=da_urho_,
v=da_vrho_,
isglobal=True,
s=4e6,
interp_to_common_grid=True,
traj_interp_order=3)
ftle = .5 * np.log(ftle)
return ftle
def calc_Koc(ds, detrend=False, plot=False):
if detrend:
ds['s'].data = signal.detrend(ds['s'].data,
axis=0) + ds['s'].mean(dim='time').data
s_mean = ds['s'].mean(dim='time')
s_anom = ds['s'] - s_mean
w = VectorWind(ds['u'], ds['v'])
grad_s_mean = w.gradient(s_mean)
grad_s_mean = VectorWind(grad_s_mean[0], grad_s_mean[1])
abs_grad_s_mean = grad_s_mean.magnitude()
sq_abs_grad_s_mean = abs_grad_s_mean * abs_grad_s_mean
grad_s_anom = w.gradient(s_anom)
grad_s_anom = VectorWind(grad_s_anom[0], grad_s_anom[1])
abs_grad_s_anom = grad_s_anom.magnitude()
sq_abs_grad_s_anom = abs_grad_s_anom * abs_grad_s_anom
sq_abs_grad_s_anom = sq_abs_grad_s_anom.mean(dim='time')
sq_abs_grad_s_mean = sq_abs_grad_s_mean.mean(dim='longitude')
sq_abs_grad_s_anom = sq_abs_grad_s_anom.mean(dim='longitude')
Koc = sq_abs_grad_s_anom / sq_abs_grad_s_mean
if plot:
plt.ion()
fig1, ax1 = plt.subplots()
ax1.plot(lats, sq_abs_grad_s_mean.data, label='mean')
ax1.plot(lats, sq_abs_grad_s_anom.data, label='anom')
ax11 = ax1.twinx()
ax11.plot(lats, ds['q'].mean(dim=['time', 'longitude']).data)
fig2, ax2 = plt.subplots()
ax2.plot(lats, Koc, label='Koc')
ax21 = ax2.twinx()
ax21.plot(lats, ds['q'].mean(dim=['time', 'longitude']).data)
plt.show()
return Koc
def potential_vorticity_baroclinic(uwnd, vwnd, theta, coord, **kwargs):
'''
Calculate potential vorticity on isobaric levels. Requires input uwnd,
vwnd and tmp arrays to have (lat, lon, ...) format for Windspharm.
Input
-----
uwnd : zonal winds, array-like
vwnd : meridional winds, array-like
tmp : temperature, array-like
theta : potential temperature, array-like
coord : dimension name for pressure axis (eg. 'pfull')
omega : planetary rotation rate, optional
g : planetary gravitational acceleration, optional
rsphere : planetary radius, in metres, optional
'''
omega = kwargs.pop('omega', 7.08822e-05)
g = kwargs.pop('g', 3.72076)
rsphere = kwargs.pop('rsphere', 3.3962e6)
w = VectorWind(uwnd.fillna(0), vwnd.fillna(0), rsphere=rsphere)
relvort = w.vorticity()
relvort = relvort.where(relvort != 0, other=np.nan)
planvort = w.planetaryvorticity(omega=omega)
absvort = relvort + planvort
dthtady, dthtadx = w.gradient(theta.fillna(0))
dthtadx = dthtadx.where(dthtadx != 0, other=np.nan)
dthtady = dthtady.where(dthtady != 0, other=np.nan)
dthtadp = wrapped_gradient(theta, coord)
dudp = wrapped_gradient(uwnd, coord)
dvdp = wrapped_gradient(vwnd, coord)
s = -dthtadp
f = dvdp * dthtadx - dudp * dthtady
ret = g * (absvort + f / s) * s
return ret
def overturning_hm(run,
regions=[[350, 10], [80, 100], [170, 190], [260, 280]]):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
plot_dir = '/scratch/rg419/plots/overturning_monthly/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
# Create a VectorWind instance to handle the computation
w = VectorWind(data.ucomp.sel(pfull=np.arange(50., 950., 50.)),
data.vcomp.sel(pfull=np.arange(50., 950., 50.)))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
ds_chi = xr.Dataset({'vcomp': (vchi)},
coords={
'xofyear': ('xofyear', vchi.xofyear),
'pfull': ('pfull', vchi.pfull),
'lat': ('lat', vchi.lat),
'lon': ('lon', vchi.lon)
})
def get_lons(lonin, data):
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]
]
return lons
# Set figure parameters
rcParams['figure.figsize'] = 10, 7
rcParams['font.size'] = 14
# Start figure with 4 subplots
fig, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,
2,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4]
i = 0
for ax in axes:
lons = get_lons(regions[i], data)
psi_chi = mass_streamfunction(ds_chi, lons=lons, dp_in=50.)
psi_chi /= 1.e9
i = i + 1
f1 = psi_chi.sel(pfull=500).plot.contourf(ax=ax,
x='xofyear',
y='lat',
add_labels=False,
add_colorbar=False,
levels=np.arange(
-500., 501., 100.),
extend='both')
ax1.set_title('West coast')
ax2.set_title('Land')
ax3.set_title('East coast')
ax4.set_title('Ocean')
for ax in [ax1, ax2, ax3, ax4]:
ax.grid(True, linestyle=':')
ax.set_ylim(-60, 60)
ax.set_yticks(np.arange(-60., 61., 30.))
ax.set_xticks([0, 18, 36, 54, 72])
ax3.set_xlabel('Pentad')
ax4.set_xlabel('Pentad')
ax1.set_ylabel('Latitude')
ax3.set_ylabel('Latitude')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label('Overturning streamfunction')
# Save as a pdf
plt.savefig(plot_dir + 'regional_overturning_hm_' + run + '.pdf',
format='pdf')
plt.close()
data.close()
def plot_vort_dev(run,
land_mask=None,
lev=200,
qscale=150.,
windtype='full',
ref_arrow=10.,
video=False):
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
sinphi = np.sin(data.lat * np.pi / 180.)
zeta = (2. * mc.omega * sinphi - 1. * gr.ddy(data.ucomp)) * 86400.
# Take zonal anomaly
data_zanom = data - data.mean('lon')
# Get rotational and divergent components of the flow
w = VectorWind(data.ucomp.sel(pfull=lev), data.vcomp.sel(pfull=lev))
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
uchi_zanom = (uchi - uchi.mean('lon')).sortby('lat')
vchi_zanom = (vchi - vchi.mean('lon')).sortby('lat')
upsi_zanom = (upsi - upsi.mean('lon')).sortby('lat')
vpsi_zanom = (vpsi - vpsi.mean('lon')).sortby('lat')
# Start figure with 1 subplots
rcParams['figure.figsize'] = 10, 5
rcParams['font.size'] = 14
for i in range(72):
fig, ax1 = plt.subplots()
title = 'Pentad ' + str(int(data.xofyear[i]))
f1 = zeta.sel(xofyear=i + 1,
pfull=lev).plot.contourf(x='lon',
y='lat',
ax=ax1,
add_labels=False,
add_colorbar=False,
extend='both',
zorder=1,
levels=np.arange(-10., 10., 2.))
if windtype == 'div':
b = ax1.quiver(data.lon[::6],
data.lat[::3],
uchi_zanom[i, ::3, ::6],
vchi_zanom[i, ::3, ::6],
scale=qscale,
angles='xy',
width=0.005,
headwidth=3.,
headlength=5.,
zorder=3)
ax1.quiverkey(b,
0.,
-0.5,
ref_arrow,
str(ref_arrow) + ' m/s',
fontproperties={
'weight': 'bold',
'size': 10
},
color='k',
labelcolor='k',
labelsep=0.03,
zorder=10)
elif windtype == 'rot':
b = ax1.quiver(data.lon[::6],
data.lat[::3],
upsi_zanom[i, ::3, ::6],
vpsi_zanom[i, ::3, ::6],
scale=qscale,
angles='xy',
width=0.005,
headwidth=3.,
headlength=5.,
zorder=3)
ax1.quiverkey(b,
0.,
-0.5,
ref_arrow,
str(ref_arrow) + ' m/s',
fontproperties={
'weight': 'bold',
'size': 10
},
color='k',
labelcolor='k',
labelsep=0.03,
zorder=10)
elif windtype == 'full':
b = ax1.quiver(data.lon[::6],
data.lat[::3],
data_zanom.ucomp.sel(pfull=lev)[i, ::3, ::6],
data_zanom.vcomp.sel(pfull=lev)[i, ::3, ::6],
scale=qscale,
angles='xy',
width=0.005,
headwidth=3.,
headlength=5.,
zorder=3)
ax1.quiverkey(b,
0.,
-0.5,
ref_arrow,
str(ref_arrow) + ' m/s',
fontproperties={
'weight': 'bold',
'size': 10
},
color='k',
labelcolor='k',
labelsep=0.03,
zorder=10)
else:
windtype = 'none'
ax1.grid(True, linestyle=':')
ax1.set_ylim(-60., 60.)
ax1.set_yticks(np.arange(-60., 61., 30.))
ax1.set_xticks(np.arange(0., 361., 90.))
ax1.set_title(title)
if not land_mask == None:
land = xr.open_dataset(land_mask)
land.land_mask.plot.contour(x='lon',
y='lat',
ax=ax1,
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k')
ax1.set_ylabel('Latitude')
ax1.set_xlabel('Longitude')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.93,
bottom=0.05,
hspace=0.25,
wspace=0.2)
cb1 = fig.colorbar(f1,
ax=ax1,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.15,
aspect=60,
shrink=0.5)
levstr = ''
windtypestr = ''
msestr = ''
vidstr = ''
if lev != 850:
levstr = '_' + str(lev)
if windtype != 'full':
windtypestr = '_' + windtype
if video:
vidstr = 'video/'
plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/' + run + '/' + vidstr + windtype + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
if video:
plt.savefig(plot_dir + 'wind_and_vort_zanom_' +
str(int(data.xofyear[i])) + levstr + windtypestr +
'.png',
format='png')
else:
plt.savefig(plot_dir + 'wind_and_vort_zanom_' +
str(int(data.xofyear[i])) + levstr + windtypestr +
'.pdf',
format='pdf')
plt.close()
# try:
# ws_ccmp.to_netcdf('datasets/CCMP_windspeed.nc')
# print('saved')
# except:
# pass
# %%
#ws_ccmp1=xr.open_dataset('datasets/CCMP_windspeed.nc')
#wu=xr.open_dataset('datasets/uwnd.10m.mon.mean.nc').sel(level=10).uwnd
#wv=xr.open_dataset('datasets/vwnd.10m.mon.mean.nc').sel(level=10).vwnd
ws_ccmp=xr.open_dataset('processed/CCMP_ws_1deg_global.nc')
wu=ws_ccmp.uwnd
wv=ws_ccmp.vwnd
# %% Test Horizontal Divergence
w = VectorWind(wu, wv)
#spd = w.magnitude()
divergence = w.divergence().sel(lat=slice(20,-20),lon=slice(120,290),time=slice('1997-07-01','2020-01-01'))
#div.mean(dim='time').plot()
wu=wu.sel(lat=slice(-20,20),lon=slice(120,290),time=slice('1997-07-01','2020-01-01'))
wv=wv.sel(lat=slice(-20,20),lon=slice(120,290),time=slice('1997-07-01','2020-01-01'))
# %% Prepare Figure
lanina=pd.read_csv('processed/indexes/la_nina_events.csv')
cp_nino=pd.read_csv('processed/indexes/cp_events.csv')
ep_nino=pd.read_csv('processed/indexes/ep_events.csv')
fp='processed/combined_dataset/month_data_exports.nc'
info=xr.open_mfdataset(fp).sel(Mooring=195).to_dataframe()
from windspharm.xarray import VectorWind
from windspharm.examples import example_data_path
mpl.rcParams['mathtext.default'] = 'regular'
# Read zonal and meridional wind components from file using the xarray module.
# The components are in separate files.
ds = xr.open_mfdataset([example_data_path(f)
for f in ('uwnd_mean.nc', 'vwnd_mean.nc')])
uwnd = ds['uwnd']
vwnd = ds['vwnd']
# Create a VectorWind instance to handle the computation of streamfunction and
# velocity potential.
w = VectorWind(uwnd, vwnd)
# Compute the streamfunction and velocity potential.
sf, vp = w.sfvp()
# Pick out the field for December.
sf_dec = sf[sf['time.month'] == 12]
vp_dec = vp[vp['time.month'] == 12]
# Plot streamfunction.
clevs = [-120, -100, -80, -60, -40, -20, 0, 20, 40, 60, 80, 100, 120]
ax = plt.subplot(111, projection=ccrs.PlateCarree(central_longitude=180))
sf_dec *= 1e-6
fill_sf = sf_dec[0].plot.contourf(ax=ax, levels=clevs, cmap=plt.cm.RdBu_r,
transform=ccrs.PlateCarree(), extend='both',
add_colorbar=False)
def abs_vort_hm(run,
lev=150,
filename='plev_daily',
timeav='pentad',
period_fac=1.,
latin=25.):
rcParams['figure.figsize'] = 9, 9
rcParams['font.size'] = 18
rcParams['text.usetex'] = True
plot_dir = '/scratch/rg419/plots/clean_diags/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
name_temp = '/scratch/rg419/Data_moist/' + run + '/run%03d/' + filename + '.nc'
names = [name_temp % m for m in range(397, 409)]
#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})
uwnd = data.ucomp.sel(pfull=lev)
vwnd = data.vcomp.sel(pfull=lev)
# Create a VectorWind instance to handle the computation of streamfunction and
# velocity potential.
w = VectorWind(uwnd, vwnd)
# Compute the streamfunction and velocity potential.
data['vor'], data['div'] = w.vrtdiv()
#data = xr.open_dataset('/scratch/rg419/Data_moist/'+run+'climatologies/'+run+'.nc')
#Coriolis
omega = 7.2921150e-5
f = 2 * omega * np.sin(data.lat * np.pi / 180)
lat_hm = data.lat[np.argmin(np.abs(data.lat - latin))]
abs_vort = (f + data.vor).sel(lat=lat_hm) * 86400.
levels = np.arange(0., 14.1, 2.)
mn_dic = month_dic(1)
tickspace = range(13, 72, 18)
labels = [mn_dic[(k + 5) / 6] for k in tickspace]
# Plot
f1 = abs_vort.plot.contourf(x='lon',
y='time',
extend='both',
levels=levels,
add_colorbar=False,
add_labels=False)
plt.set_cmap('inferno_r')
plt.ylabel('Time')
#plt.yticks(tickspace, labels, rotation=25)
plt.xlabel('Longitude')
plt.xlim(60, 150)
plt.ylim(240 + 33 * 360, 90 + 33 * 360)
plt.grid(True, linestyle=':')
plt.tight_layout()
#Colorbar
cb1 = plt.colorbar(f1,
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.1,
aspect=30)
cb1.set_label('$day^{-1}$')
figname = 'abs_vort_lon_hm.pdf'
plt.savefig(plot_dir + figname, format='pdf')
plt.close()
def plot_gill_dev(run, land_mask=None, lev=850, qscale=100., windtype='full', ref_arrow=5, mse=False, video=False):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' + run + '.nc')
data['mse'] = (mc.cp_air * data.temp + mc.L * data.sphum + mc.grav * data.height)/1000.
data['precipitation'] = (data.precipitation*86400.)
# Take zonal anomaly
data_zanom = data - data.mean('lon')
# Get rotational and divergent components of the flow
if windtype == 'div' or windtype=='rot':
w = VectorWind(data.ucomp.sel(pfull=lev), data.vcomp.sel(pfull=lev))
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
uchi_zanom = (uchi - uchi.mean('lon')).sortby('lat')
vchi_zanom = (vchi - vchi.mean('lon')).sortby('lat')
upsi_zanom = (upsi - upsi.mean('lon')).sortby('lat')
vpsi_zanom = (vpsi - vpsi.mean('lon')).sortby('lat')
# Start figure with 1 subplots
rcParams['figure.figsize'] = 10, 5
rcParams['font.size'] = 14
for i in range(72):
fig, ax1 = plt.subplots()
title = 'Pentad ' + str(int(data.xofyear[i]))
if mse:
f1 = data.mse.sel(xofyear=i+1, pfull=850.).plot.contourf(x='lon', y='lat', ax=ax1, add_labels=False, add_colorbar=False, extend='both', cmap='Blues', zorder=1, levels = np.arange(290.,341.,5.))
else:
f1 = data.precipitation[i,:,:].plot.contourf(x='lon', y='lat', ax=ax1, levels = np.arange(2.,21.,2.), add_labels=False, add_colorbar=False, extend='both', cmap='Blues', zorder=1)
#data_zanom.slp[i+4,:,:].plot.contour(x='lon', y='lat', ax=axes[i], levels = np.arange(-15.,16.,3.), add_labels=False, colors='0.5', alpha=0.5)
ax1.contour(data_zanom.lon, data_zanom.lat, data_zanom.slp[i,:,:], levels = np.arange(0.,16.,3.), colors='0.4', alpha=0.5, zorder=2)
ax1.contour(data_zanom.lon, data_zanom.lat, data_zanom.slp[i,:,:], levels = np.arange(-15.,0.,3.), colors='0.4', alpha=0.5, linestyle='--', zorder=2)
if windtype=='div':
b = ax1.quiver(data.lon[::6], data.lat[::3], uchi_zanom[i,::3,::6], vchi_zanom[i,::3,::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3)
ax1.quiverkey(b, 5.,65., ref_arrow, str(ref_arrow) + ' m/s', fontproperties={'weight': 'bold', 'size': 10}, color='k', labelcolor='k', labelsep=0.03, zorder=10)
elif windtype=='rot':
b = ax1.quiver(data.lon[::6], data.lat[::3], upsi_zanom[i,::3,::6], vpsi_zanom[i,::3,::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3)
ax1.quiverkey(b, 5.,65., ref_arrow, str(ref_arrow) + ' m/s', fontproperties={'weight': 'bold', 'size': 10}, color='k', labelcolor='k', labelsep=0.03, zorder=10)
elif windtype=='full':
b = ax1.quiver(data.lon[::6], data.lat[::3], data_zanom.ucomp.sel(pfull=lev)[i,::3,::6], data_zanom.vcomp.sel(pfull=lev)[i,::3,::6], scale=qscale, angles='xy', width=0.005, headwidth=3., headlength=5., zorder=3)
ax1.quiverkey(b, 5.,65., ref_arrow, str(ref_arrow) + ' m/s', coordinates='data', fontproperties={'weight': 'bold', 'size': 10}, color='k', labelcolor='k', labelsep=0.03, zorder=10)
else:
windtype='none'
ax1.grid(True,linestyle=':')
ax1.set_ylim(-60.,60.)
ax1.set_yticks(np.arange(-60.,61.,30.))
ax1.set_xticks(np.arange(0.,361.,90.))
ax1.set_title(title)
if not land_mask==None:
land = xr.open_dataset(land_mask)
land.land_mask.plot.contour(x='lon', y='lat', ax=ax1, levels=np.arange(-1.,2.,1.), add_labels=False, colors='k')
land.zsurf.plot.contour(ax=ax1, x='lon', y='lat', levels=np.arange(0.,2001.,1000.), add_labels=False, colors='k')
ax1.set_ylabel('Latitude')
ax1.set_xlabel('Longitude')
plt.subplots_adjust(left=0.1, right=0.97, top=0.93, bottom=0.05, hspace=0.25, wspace=0.2)
cb1=fig.colorbar(f1, ax=ax1, use_gridspec=True, orientation = 'horizontal',fraction=0.05, pad=0.15, aspect=60, shrink=0.5)
levstr=''; windtypestr=''; msestr=''; vidstr=''
if lev != 850:
levstr = '_' + str(lev)
if windtype != 'full':
windtypestr = '_' + windtype
if mse:
msestr = '_mse'
if video:
vidstr='video/'
plot_dir = '/scratch/rg419/plots/zonal_asym_runs/gill_development/' + run +'/' + vidstr + windtype + msestr + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
if video:
plt.savefig(plot_dir + 'wind_and_slp_zanom_' + str(int(data.xofyear[i])) + levstr + windtypestr + msestr + '.png', format='png')
else:
plt.savefig(plot_dir + 'wind_and_slp_zanom_' + str(int(data.xofyear[i])) + levstr + windtypestr + msestr + '.pdf', format='pdf')
plt.close()
from windspharm.xarray import VectorWind
from windspharm.examples import example_data_path
mpl.rcParams['mathtext.default'] = 'regular'
# Read zonal and meridional wind components from file using the xarray module.
# The components are in separate files.
ds = xr.open_mfdataset([example_data_path(f)
for f in ('uwnd_mean.nc', 'vwnd_mean.nc')])
uwnd = ds['uwnd']
vwnd = ds['vwnd']
# Create a VectorWind instance to handle the computations.
w = VectorWind(uwnd, vwnd)
# Compute components of rossby wave source: absolute vorticity, divergence,
# irrotational (divergent) wind components, gradients of absolute vorticity.
eta = w.absolutevorticity()
div = w.divergence()
uchi, vchi = w.irrotationalcomponent()
etax, etay = w.gradient(eta)
etax.attrs['units'] = 'm**-1 s**-1'
etay.attrs['units'] = 'm**-1 s**-1'
# Combine the components to form the Rossby wave source term.
S = eta * -1. * div - (uchi * etax + vchi * etay)
# Pick out the field for December at 200 hPa.
S_dec = S[S['time.month'] == 12]
else:
dp = dp_in * 100.
if intdown:
psi = c * dp * np.flip(ubar, axis=ubar.dims.index('pfull')).cumsum(
'pfull') #[:,:,::-1]
else:
psi = -1. * c * dp * ubar.cumsum('pfull') #[:,:,::-1]
#psi.plot.contourf(x='lat', y='pfull', yincrease=False)
#plt.show()
return psi #c*dp*vbar[:,::-1,:].cumsum('pfull')#[:,:,::-1]
if __name__ == '__main__':
# example calculating Walker cell for a GFDL dataset
from windspharm.xarray import VectorWind
d = xarray.open_dataset(
'/disca/share/rg419/Data_moist/climatologies/3q_shallow.nc')
w = VectorWind(d.ucomp.sel(pfull=np.arange(50., 950., 50.)),
d.vcomp.sel(pfull=np.arange(50., 950., 50.)))
# Compute variables
uchi, vchi, upsi, vpsi = w.helmholtz()
walker = walker_cell(uchi, dp_in=-50., intdown=True)
walker /= 1.e9
walker.sel(xofyear=1).plot.contourf(x='lon',
y='pfull',
yincrease=False,
levels=np.arange(-300., 301., 50.))
plt.show()
def h_w_mass_flux_monthly(run, lev=500., dp=5000.):
data = xr.open_dataset(
'/scratch/rg419/obs_and_reanalysis/era_v_clim_alllevs.nc')
data_u = xr.open_dataset(
'/scratch/rg419/obs_and_reanalysis/era_u_clim_alllevs.nc')
plot_dir = '/scratch/rg419/plots/overturning_monthly/era/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data.coords['month'] = (data.xofyear - 1) // 6 + 1
data = data.groupby('month').mean(('xofyear'))
# Create a VectorWind instance to handle the computation
w = VectorWind(data.ucomp.sel(pfull=np.arange(50., 950., 50.)),
data.vcomp.sel(pfull=np.arange(50., 950., 50.)))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
coslat = np.cos(data.lat * np.pi / 180)
# Evaluate mass fluxes for the zonal and meridional components (Walker and Hadley) following Schwendike et al. 2014
mass_flux_zon = (gr.ddx(uchi)).cumsum('pfull') * dp * coslat / mc.grav
mass_flux_merid = (gr.ddy(vchi)).cumsum('pfull') * dp * coslat / mc.grav
# Set figure parameters
rcParams['figure.figsize'] = 15, 11
rcParams['font.size'] = 14
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3,
4,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
f1 = mass_flux_merid.sel(pfull=lev)[i, :, :].plot.contourf(
ax=ax,
x='lon',
y='lat',
add_labels=False,
add_colorbar=False,
levels=np.arange(-0.0065, 0.0066, 0.001),
extend='both')
mass_flux_zon.sel(pfull=lev)[i, :, :].plot.contour(ax=ax,
x='lon',
y='lat',
add_labels=False,
colors='k',
levels=np.arange(
0.0005, 0.0066,
0.001))
mass_flux_zon.sel(pfull=lev)[i, :, :].plot.contour(
ax=ax,
x='lon',
y='lat',
add_labels=False,
colors='0.5',
levels=np.arange(-0.0065, -0.00049, 0.001))
i = i + 1
ax.set_ylim(-60, 60)
ax.set_xticks(np.arange(0, 361, 90))
ax.set_yticks(np.arange(-60, 61, 30))
ax.grid(True, linestyle=':')
plt.subplots_adjust(left=0.05,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.2,
wspace=0.2)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label(
'Vertical mass flux associated with meridional circulation, kgm$^{-2}$s$^{-1}$'
)
figname = plot_dir + 'h_w_' + run + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
import xarray as xr
from windspharm.xarray import VectorWind
f = xr.open_dataset(
"/Users/brianpm/Documents/www.ncl.ucar.edu/Applications/Data/cdf/uv300.nc")
u = f["U"]
v = f["V"]
w = VectorWind(u, v)
## VERY IMPORTANT: VectorWind apparently reverses latitude to be decreasing (90 to -90)
vort, div = w.vrtdiv() # Relative vorticity and horizontal divergence.
sf, vp = w.sfvp() # The streamfunction and velocity potential respectively.
uchi, vchi, upsi, vpsi = w.helmholtz()
# plot the results
import matplotlib as mpl
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import numpy as np
fig, ax = plt.subplots(figsize=(12, 12),
nrows=2,
subplot_kw={"projection": ccrs.PlateCarree()},
constrained_layout=True)
N = mpl.colors.Normalize(vmin=-8e6, vmax=8e6)
x, y = np.meshgrid(f['lon'], f['lat'])
im0 = ax[0].contourf(x,
y,
vp[0, ::-1, :],
norm=N,
transform=ccrs.PlateCarree(),
import xarray as xr
from windspharm.xarray import VectorWind
from windspharm.examples import example_data_path
mpl.rcParams['mathtext.default'] = 'regular'
# Read zonal and meridional wind components from file using the xarray module.
# The components are in separate files.
ds = xr.open_mfdataset(
[example_data_path(f) for f in ('uwnd_mean.nc', 'vwnd_mean.nc')])
uwnd = ds['uwnd']
vwnd = ds['vwnd']
# Create a VectorWind instance to handle the computations.
w = VectorWind(uwnd, vwnd)
# Compute components of rossby wave source: absolute vorticity, divergence,
# irrotational (divergent) wind components, gradients of absolute vorticity.
eta = w.absolutevorticity()
div = w.divergence()
uchi, vchi = w.irrotationalcomponent()
etax, etay = w.gradient(eta)
etax.attrs['units'] = 'm**-1 s**-1'
etay.attrs['units'] = 'm**-1 s**-1'
# Combine the components to form the Rossby wave source term.
S = eta * -1. * div - (uchi * etax + vchi * etay)
# Pick out the field for December at 200 hPa.
S_dec = S[S['time.month'] == 12]
#S.to_netcdf(RUTA + 'monthly_RWS.nc4')
#uchi_monthly = xr.concat(uchi_monthly, dim='month')
#uchi_monthly['month'] = np.array([8, 9, 10, 11, 12, 1, 2])
#uchi_monthly.to_netcdf(RUTA + 'monthly_uchi.nc4')
#vchi_monthly = xr.concat(vchi_monthly, dim='month')
#vchi_monthly['month'] = np.array([8, 9, 10, 11, 12, 1, 2])
#vchi_monthly.to_netcdf(RUTA + 'monthly_vchi.nc4')
S_seasonal = []
uchi_seasonal = []
vchi_seasonal = []
for i in np.arange(5):
ds1 = ds.isel(month=range(i, i + 3)).mean(dim='month')
#create vectorwind instance
w = VectorWind(ds1['u'][:, :, :], ds1['v'][:, :, :])
eta = w.absolutevorticity()
div = w.divergence()
uchi, vchi = w.irrotationalcomponent()
etax, etay = w.gradient(eta)
etax.attrs['units'] = 'm**-1 s**-1'
etay.attrs['units'] = 'm**-1 s**-1'
# Combine the components to form the Rossby wave source term.
S = eta * -1. * div - (uchi * etax + vchi * etay)
S *= 1e11
S_seasonal.append(S)
uchi_seasonal.append(uchi)
vchi_seasonal.append(vchi)
S_seasonal = xr.concat(S_seasonal, dim='season')
S_seasonal['season'] = np.array(seas)
data_vo = (dvdx - dudy) * 86400.
data_t = xr.open_dataset(
'/disca/share/reanalysis_links/jra_55/1958_2016/temp_monthly/atmos_monthly_together.nc'
)
data_q = xr.open_dataset(
'/disca/share/rg419/JRA_55/sphum_monthly/atmos_monthly_together.nc')
data_z = xr.open_dataset(
'/disca/share/reanalysis_links/jra_55/1958_2016/height_monthly/atmos_monthly_together.nc'
)
data_mse = (mc.cp_air * data_t.var11 + mc.L * data_q.var51 +
9.81 * data_z.var7) / 1000.
# Create a VectorWind instance to handle the computation
w = VectorWind(data_u.sel(lev=np.arange(5000., 100001., 5000.)),
data_v.sel(lev=np.arange(5000., 100001., 5000.)))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
coslat = np.cos(data_u.lat * np.pi / 180)
dp = 5000.
# Evaluate mass fluxes for the zonal and meridional components (Walker and Hadley) following Schwendike et al. 2014
mass_flux_zon = (gr.ddx(uchi)).cumsum('lev') * dp * coslat / 9.81
mass_flux_merid = (gr.ddy(vchi)).cumsum('lev') * dp * coslat / 9.81
land_mask = '/scratch/rg419/python_scripts/land_era/ERA-I_Invariant_0125.nc'
land = xr.open_dataset(land_mask)
def plot_winter_climate(data,
from windspharm.xarray import VectorWind
GFDL_DATA = os.environ['GFDL_DATA']
filename = GFDL_DATA + 'full_qflux/run121/atmos_pentad.nc'
fileout = GFDL_DATA + 'full_qflux/run121/atmos_test.nc'
dsin= Dataset(filename, 'r', format='NETCDF3_CLASSIC')
data = xr.open_dataset(filename,decode_times=False)
uwnd = data.ucomp
vwnd = data.vcomp
# Create a VectorWind instance to handle the computation of streamfunction and
# velocity potential.
w = VectorWind(uwnd, vwnd)
# Compute the streamfunction and velocity potential.
streamfun, vel_pot = w.sfvp()
dsout= Dataset(fileout, 'w', format='NETCDF3_CLASSIC')
dsout = copy_netcdf_attrs(dsin, dsout)
sf_out = dsout.createVariable('streamfun', 'f4', ('time', 'pfull', 'lat', 'lon',))
sf_out.setncatts({k: dsout.variables['ps'].getncattr(k) for k in dsout.variables['ps'].ncattrs()})
sf_out.setncattr('long_name', 'streamfunction')
sf_out.setncattr('units', 'm**2 s**-1')
sf_out[:] = streamfun.load().data
def plot_sf_vp(land_mask=None):
rcParams['figure.figsize'] = 10, 5
rcParams['font.size'] = 16
plot_dir = '/scratch/rg419/plots/era_wn2/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
land = xr.open_dataset(land_mask)
data = xr.open_dataset('/scratch/rg419/obs_and_reanalysis/sep_levs_u/era_u_200_mm.nc')
print data.time.units
uwnd = data.u.load().squeeze('level')
uwnd = uwnd[0:456,:,:]
data_sn = data.resample(time='Q-NOV').mean()
uwnd_sn = data_sn.u.load().squeeze('level')
data = xr.open_dataset('/scratch/rg419/obs_and_reanalysis/sep_levs_v/era_v_200_mm.nc')
vwnd = data.v.load()
data_sn = data.resample(time='Q-NOV').mean()
vwnd_sn = data_sn.v.load()
# Create a VectorWind instance to handle the computation
w = VectorWind(uwnd, vwnd)
# Compute variables
streamfun, vel_pot = w.sfvp()
streamfun = streamfun/10.**6
# Create a VectorWind instance to handle the computation
w = VectorWind(uwnd_sn, vwnd_sn)
# Compute variables
streamfun_sn, vel_pot_sn = w.sfvp()
streamfun_sn = streamfun_sn/10.**6
lats = [data.latitude[i] for i in range(len(data.latitude)) if data.latitude[i] >= 5.]
for year in range(1979,2017):
sf_max_lat = np.zeros((12,))
sf_max_lon = np.zeros((12,))
for month in range(1,13):
streamfun_asia_i = streamfun.sel(time=str(year)+'-%02d' % month , latitude=lats)
sf_max_i = streamfun_asia_i.where(streamfun_asia_i==streamfun_asia_i.max(), drop=True)
sf_max_lon[month-1] = sf_max_i.longitude.values
sf_max_lat[month-1] = sf_max_i.latitude.values
f1 = streamfun_sn.sel(time=str(year) + '-08').squeeze('time').plot.contourf(x='longitude', y='latitude', levels = np.arange(-140.,141.,10.), add_labels=False, add_colorbar=False, extend='both')
for i in range(12):
plt.text(sf_max_lon[i]+3, sf_max_lat[i], str(i+1), fontsize=10)
plt.plot(sf_max_lon, sf_max_lat, 'kx-', mew=1.5)
plt.grid(True,linestyle=':')
plt.colorbar(f1)
land.lsm[0,:,:].plot.contour(x='longitude', y='latitude', levels=np.arange(-1.,2.,1.), add_labels=False, colors='k')
plt.savefig(plot_dir + 'streamfun_era_' + str(year) + '.pdf', format='pdf')
plt.close()
def overturning_monthly(run, lonin=[-1., 361.]):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
plot_dir = '/scratch/rg419/plots/overturning_monthly/' + run + '/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data.coords['month'] = (data.xofyear - 1) // 6 + 1
data = data.groupby('month').mean(('xofyear'))
#data['vcomp'] = data.vcomp.fillna(0.)
#data['ucomp'] = data.ucomp.fillna(0.)
# Create a VectorWind instance to handle the computation
w = VectorWind(data.ucomp.sel(pfull=np.arange(50., 950., 50.)),
data.vcomp.sel(pfull=np.arange(50., 950., 50.)))
#w = VectorWind(data.ucomp, data.vcomp)
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi, vchi, upsi, vpsi = w.helmholtz()
#print(vchi.pfull)
#print(data.pfull)
#data.vcomp.mean('lon')[0,:,:].plot.contourf(x='lat', y='pfull', yincrease=False, add_labels=False)
#plt.figure(2)
#vchi.mean('lon')[0,:,:].plot.contourf(x='lat', y='pfull', yincrease=False, add_labels=False)
#plt.show()
ds_chi = xr.Dataset({'vcomp': (vchi)},
coords={
'month': ('month', vchi.month),
'pfull': ('pfull', vchi.pfull),
'lat': ('lat', vchi.lat),
'lon': ('lon', vchi.lon)
})
ds_psi = xr.Dataset({'vcomp': (vpsi)},
coords={
'month': ('month', vchi.month),
'pfull': ('pfull', vchi.pfull),
'lat': ('lat', vchi.lat),
'lon': ('lon', vchi.lon)
})
def get_lons(lonin, data):
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]
]
return lons
lons = get_lons(lonin, data)
psi = mass_streamfunction(data, lons=lons, dp_in=50., use_v_locally=True)
psi /= 1.e9
psi_chi = mass_streamfunction(ds_chi, lons=lons, dp_in=50.)
psi_chi /= 1.e9
# Set figure parameters
rcParams['figure.figsize'] = 10, 7
rcParams['font.size'] = 14
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3, 4)
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
psi[:, i, :].plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(0., 601, 100.),
colors='k',
add_labels=False)
psi[:, i, :].plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-600., 0., 100.),
colors='k',
linestyles='dashed',
add_labels=False)
f1 = data.ucomp.sel(month=i +
1).sel(lon=lons).mean('lon').plot.contourf(
ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-50., 50.1, 5.),
extend='both',
add_labels=False,
add_colorbar=False)
m = mc.omega * mc.a**2. * np.cos(
psi.lat * np.pi / 180.)**2. + data.ucomp.sel(
lon=lons).mean('lon') * mc.a * np.cos(psi.lat * np.pi / 180.)
m_levs = mc.omega * mc.a**2. * np.cos(
np.arange(-60., 1., 5.) * np.pi / 180.)**2.
m.sel(month=i + 1).plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=m_levs,
colors='0.7',
add_labels=False)
i = i + 1
ax.set_xlim(-35, 35)
ax.set_xticks(np.arange(-30, 31, 15))
ax.grid(True, linestyle=':')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label('Zonal wind speed, m/s')
if lonin == [-1., 361.]:
plt.savefig(plot_dir + 'psi_u_' + run + '.pdf', format='pdf')
else:
figname = plot_dir + 'psi_u_' + run + '_' + str(int(
lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3, 4)
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
psi_chi[:, i, :].plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(0., 601, 100.),
colors='k',
add_labels=False)
psi_chi[:, i, :].plot.contour(ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-600., 0., 100.),
colors='k',
linestyles='dashed',
add_labels=False)
f1 = uchi.sel(month=i + 1).sel(lon=lons).mean('lon').plot.contourf(
ax=ax,
x='lat',
y='pfull',
yincrease=False,
levels=np.arange(-3., 3.1, 0.5),
extend='both',
add_labels=False,
add_colorbar=False)
i = i + 1
ax.set_xlim(-35, 35)
ax.set_xticks(np.arange(-30, 31, 15))
ax.grid(True, linestyle=':')
plt.subplots_adjust(left=0.1,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label('Zonal wind speed, m/s')
if lonin == [-1., 361.]:
plt.savefig(plot_dir + 'psi_chi_u_' + run + '.pdf', format='pdf')
else:
figname = plot_dir + 'psi_chi_u_' + run + '_' + str(int(
lonin[0])) + '_' + str(int(lonin[1])) + '.pdf'
plt.savefig(figname, format='pdf')
plt.close()
def horiz_streamfun_monthly(run, land_mask=None):
data = xr.open_dataset('/disca/share/rg419/Data_moist/climatologies/' +
run + '.nc')
data.coords['month'] = (data.xofyear - 1) // 6 + 1
data = data.groupby('month').mean(('xofyear'))
# Create a VectorWind instance to handle the computation
w = VectorWind(data.ucomp.sel(pfull=150.), data.vcomp.sel(pfull=150.))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi_150, vchi_150, upsi_150, vpsi_150 = w.helmholtz()
vel_pot_150 = (vel_pot - vel_pot.mean('lon')) / 10.**6
streamfun_150 = (streamfun - streamfun.mean('lon')) / 10.**6
w = VectorWind(data.ucomp.sel(pfull=850.), data.vcomp.sel(pfull=850.))
# Compute variables
streamfun, vel_pot = w.sfvp()
uchi_850, vchi_850, upsi_850, vpsi_850 = w.helmholtz()
vel_pot_850 = (vel_pot - vel_pot.mean('lon')) / 10.**6
streamfun_850 = (streamfun - streamfun.mean('lon')) / 10.**6
# Set figure parameters
rcParams['figure.figsize'] = 15, 8
rcParams['font.size'] = 14
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3,
4,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
f1 = streamfun_150[i, :, :].plot.contourf(ax=ax,
x='lon',
y='lat',
add_labels=False,
add_colorbar=False,
levels=np.arange(
-30., 30.1, 5.),
extend='both')
streamfun_850[i, :, :].plot.contour(ax=ax,
x='lon',
y='lat',
add_labels=False,
add_colorbar=False,
colors='k',
levels=np.arange(0, 12.1, 2.))
ax.contour(streamfun_850.lon,
streamfun_850.lat,
streamfun_850[i, :, :],
colors='k',
ls='--',
levels=np.arange(-12., 0., 2.))
if not land_mask == None:
land = xr.open_dataset(land_mask)
land.land_mask.plot.contour(x='lon',
y='lat',
ax=axes[i],
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k',
alpha=0.5)
#streamfun_850[i,:,:].plot.contour(ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, cmap='PRGn', levels=np.arange(-12.,12.1,2.))
i = i + 1
ax.set_ylim(-35, 35)
ax.set_yticks(np.arange(-30, 31, 15))
ax.set_xlim(0, 360)
ax.set_xticks(np.arange(0, 361, 60))
ax.grid(True, linestyle=':')
for ax in [ax1, ax5, ax9]:
ax.set_ylabel('Latitude')
for ax in [ax9, ax10, ax11, ax12]:
ax.set_xlabel('Longitude')
plt.subplots_adjust(left=0.08,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label('Horizontal streamfunction')
plt.savefig(plot_dir + 'streamfun_' + run + '.pdf', format='pdf')
plt.close()
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3,
4,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
f1 = vel_pot_150[i, :, :].plot.contourf(ax=ax,
x='lon',
y='lat',
add_labels=False,
add_colorbar=False,
levels=np.arange(
-30., 30.1, 5.),
extend='both')
vel_pot_850[i, :, :].plot.contour(ax=ax,
x='lon',
y='lat',
add_labels=False,
add_colorbar=False,
colors='k',
levels=np.arange(0, 12.1, 2.))
ax.contour(vel_pot_850.lon,
vel_pot_850.lat,
vel_pot_850[i, :, :],
colors='k',
ls='--',
levels=np.arange(-12., 0., 2.))
if not land_mask == None:
land = xr.open_dataset(land_mask)
land.land_mask.plot.contour(x='lon',
y='lat',
ax=axes[i],
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k',
alpha=0.5)
#vel_pot_850[i,:,:].plot.contour(ax=ax, x='lon', y='lat', add_labels=False, add_colorbar=False, cmap='PRGn', levels=np.arange(-12.,12.1,2.))
i = i + 1
ax.set_ylim(-35, 35)
ax.set_yticks(np.arange(-30, 31, 15))
ax.set_xlim(0, 360)
ax.set_xticks(np.arange(0, 361, 60))
ax.grid(True, linestyle=':')
for ax in [ax1, ax5, ax9]:
ax.set_ylabel('Latitude')
for ax in [ax9, ax10, ax11, ax12]:
ax.set_xlabel('Longitude')
plt.subplots_adjust(left=0.08,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label('Velocity potential')
plt.savefig(plot_dir + 'vel_pot_' + run + '.pdf', format='pdf')
plt.close()
# Start figure with 12 subplots
fig, ((ax1, ax2, ax3, ax4), (ax5, ax6, ax7, ax8),
(ax9, ax10, ax11, ax12)) = plt.subplots(3,
4,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4, ax5, ax6, ax7, ax8, ax9, ax10, ax11, ax12]
i = 0
for ax in axes:
f1 = vel_pot_150[i, :, :].plot.contourf(ax=ax,
x='lon',
y='lat',
add_labels=False,
add_colorbar=False,
levels=np.arange(
-30., 30.1, 5.),
extend='both')
b = axes[i].quiver(data.lon[::5],
data.lat[::2],
uchi_150[i, ::2, ::5],
vchi_150[i, ::2, ::5],
scale=100.,
angles='xy')
if not land_mask == None:
land = xr.open_dataset(land_mask)
land.land_mask.plot.contour(x='lon',
y='lat',
ax=axes[i],
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k',
alpha=0.5)
i = i + 1
ax.set_ylim(-35, 35)
ax.set_yticks(np.arange(-30, 31, 15))
ax.set_xlim(0, 360)
ax.set_xticks(np.arange(0, 361, 60))
ax.grid(True, linestyle=':')
for ax in [ax1, ax5, ax9]:
ax.set_ylabel('Latitude')
for ax in [ax9, ax10, ax11, ax12]:
ax.set_xlabel('Longitude')
plt.subplots_adjust(left=0.08,
right=0.97,
top=0.95,
bottom=0.1,
hspace=0.3,
wspace=0.3)
cb1 = fig.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.05,
pad=0.1,
aspect=30,
shrink=0.5)
cb1.set_label('Velocity potential')
plt.savefig(plot_dir + 'vel_pot_vchi_' + run + '.pdf', format='pdf')
plt.close()
def plot_sf_clim(run, land_mask=None):
rcParams['figure.figsize'] = 10, 8
rcParams['font.size'] = 16
plot_dir = '/scratch/rg419/plots/egu_2018_talk_plots/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
data = xr.open_dataset('/scratch/rg419/Data_moist/climatologies/' + run +
'.nc')
uwnd = data.ucomp.sel(pfull=150.)
vwnd = data.vcomp.sel(pfull=150.)
# Create a VectorWind instance to handle the computation
w = VectorWind(uwnd, vwnd)
# Compute variables
streamfun, vel_pot = w.sfvp()
def sn_av(da):
#Take seasonal and monthly averages
da = da / 10.**6
da.coords['season'] = np.mod(da.xofyear + 5., 72.) // 18.
da_sn = da.groupby('season').mean(('xofyear'))
return da_sn
streamfun_sn = sn_av(streamfun)
streamfun_sn = streamfun_sn - streamfun_sn.mean('lon')
f, ((ax1, ax2), (ax3, ax4)) = plt.subplots(2,
2,
sharex='col',
sharey='row')
axes = [ax1, ax2, ax3, ax4]
title = ['DJF', 'MAM', 'JJA', 'SON']
for i in range(4):
f1 = streamfun_sn[i, :, :].plot.contourf(x='lon',
y='lat',
ax=axes[i],
levels=np.arange(
-36., 37., 3.),
add_labels=False,
add_colorbar=False,
extend='both')
axes[i].grid(True, linestyle=':')
if not land_mask == None:
land = xr.open_dataset(land_mask)
land.land_mask.plot.contour(x='lon',
y='lat',
ax=axes[i],
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k')
for ax in [ax1, ax3]:
ax.set_ylabel('Latitude')
for ax in [ax3, ax4]:
ax.set_xlabel('Longitude')
cb1 = f.colorbar(f1,
ax=axes,
use_gridspec=True,
orientation='horizontal',
fraction=0.15,
pad=0.15,
aspect=30,
shrink=0.5)
plt.savefig(plot_dir + 'streamfun_' + run + '.pdf', format='pdf')
plt.close()
import numpy as np
import xarray as xr
from windspharm.xarray import VectorWind
from windspharm.examples import example_data_path
# Read zonal and meridional wind components from file using the xarray module.
# The components are in separate files.
ds = xr.open_mfdataset(
[example_data_path(f) for f in ('uwnd_mean.nc', 'vwnd_mean.nc')])
uwnd = ds['uwnd']
vwnd = ds['vwnd']
# Create a VectorWind instance to handle the computation of streamfunction and
# velocity potential.
w = VectorWind(uwnd, vwnd)
# Compute the streamfunction and velocity potential.
sf, vp = w.sfvp()
# Pick out the field for December.
sf_dec = sf[sf['time.month'] == 12]
vp_dec = vp[vp['time.month'] == 12]
# Plot streamfunction.
clevs = [-120, -100, -80, -60, -40, -20, 0, 20, 40, 60, 80, 100, 120]
ax = plt.subplot(111, projection=ccrs.PlateCarree(central_longitude=180))
sf_dec *= 1e-6
fill_sf = sf_dec[0].plot.contourf(ax=ax,
levels=clevs,
cmap=plt.cm.RdBu_r,
def plot_sf_vp(land_mask=None):
rcParams['figure.figsize'] = 10, 5
rcParams['font.size'] = 16
plot_dir = '/scratch/rg419/plots/era_wn2/'
mkdir = sh.mkdir.bake('-p')
mkdir(plot_dir)
land = xr.open_dataset(land_mask)
data = xr.open_dataset(
'/scratch/rg419/obs_and_reanalysis/sep_levs_u/era_u_200_mm.nc')
uwnd = data.u.load().squeeze('level')
uwnd = uwnd[0:456, :, :]
data_sn = data.resample(time='Q-NOV').mean()
uwnd_sn = data_sn.u.load().squeeze('level')
data = xr.open_dataset(
'/scratch/rg419/obs_and_reanalysis/sep_levs_v/era_v_200_mm.nc')
vwnd = data.v.load()
data_sn = data.resample(time='Q-NOV').mean()
vwnd_sn = data_sn.v.load()
# Create a VectorWind instance to handle the computation
w = VectorWind(uwnd, vwnd)
# Compute variables
streamfun, vel_pot = w.sfvp()
streamfun = streamfun / 10.**6
# Create a VectorWind instance to handle the computation
w = VectorWind(uwnd_sn, vwnd_sn)
# Compute variables
streamfun_sn, vel_pot_sn = w.sfvp()
streamfun_sn = streamfun_sn / 10.**6
streamfun_sn_mean = streamfun_sn.groupby('time.month').mean('time')
#time_list = [str(year) + '-08' for year in range(1979,2017)]
#print time_list[0]
#streamfun_JJA_mean = streamfun_sn.sel(time=time_list[2]).mean('time')
#print streamfun_JJA_mean
#streamfun_JJA_anom = streamfun_sn.sel(time=time_list[0]) - streamfun_JJA_mean
#print streamfun_JJA_anom
for year in range(1979, 2017):
f1 = (streamfun_sn.sel(time=str(year) + '-08') -
streamfun_sn_mean.sel(month=8)).squeeze('time').plot.contourf(
x='longitude',
y='latitude',
levels=np.arange(-20., 21., 2.),
add_labels=False,
add_colorbar=False,
extend='both')
plt.grid(True, linestyle=':')
plt.colorbar(f1)
land.lsm[0, :, :].plot.contour(x='longitude',
y='latitude',
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k')
plt.savefig(plot_dir + 'streamfun_era_anom_' + str(year) + '.pdf',
format='pdf')
plt.close()
f1 = (streamfun_sn.sel(time=str(year) + '-08') -
streamfun_sn.sel(time=str(year) + '-08').mean('longitude')
).squeeze('time').plot.contourf(x='longitude',
y='latitude',
levels=np.arange(-50., 51., 5.),
add_labels=False,
add_colorbar=False,
extend='both')
plt.grid(True, linestyle=':')
plt.colorbar(f1)
land.lsm[0, :, :].plot.contour(x='longitude',
y='latitude',
levels=np.arange(-1., 2., 1.),
add_labels=False,
colors='k')
plt.savefig(plot_dir + 'streamfun_era_zanom_' + str(year) + '.pdf',
format='pdf')
plt.close()
