def precip_hm_fn(inp_fol, nproc):
#set run name
run_fol = inp_fol + '/np' + nproc
year = 1
rundata = load_year_xr(run_fol, year)
rundata.coords['pentad'] = (rundata.time // 5) + 1
conv_p = rundata.convection_rain[:, 37, :].groupby('pentad').mean(('time'))
tot_p = xr.DataArray(np.zeros((72, 128, 5)), [('pentad', conv_p.pentad),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
for year in range(2, 7):
rundata = load_year_xr(run_fol, year)
rundata.coords['pentad'] = (rundata.time // 5) + 1
#take time and zonal mean
conv_p = rundata.convection_rain[:, 37, :].groupby('pentad').mean(
('time'))
cond_p = rundata.condensation_rain[:, 37, :].groupby('pentad').mean(
('time'))
tot_p[:, :, year - 2] = (conv_p + cond_p) * 86400
tot_p_clim = tot_p.mean(('year'))
#plot
tot_p_clim.plot.contourf(x='lon',
y='pentad',
levels=np.arange(3., 31., 3.),
add_label=False)
plt.plot([float(rundata.lonb[21]),
float(rundata.lonb[21])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[26]),
float(rundata.lonb[26])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[30]),
float(rundata.lonb[30])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[35]),
float(rundata.lonb[35])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[39]),
float(rundata.lonb[39])], [20., 60.], 'k')
plt.ylim(20, 60)
plt.xlim(60, 180)
plt.xlabel('Longitude')
plt.ylabel('Pentad')
plt.title('Precipitation at 15N, mm/day')
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol + '/precip_hm.png')
plt.clf()
def ro_hm_fn(inp_fol, nproc):
#set run name
run_fol = inp_fol + '/np' + nproc
year = 1
rundata = load_year_xr(run_fol, year)
rundata.coords['pentad'] = (rundata.time // 5) + 1
conv_p = rundata.convection_rain[:, 37, :].groupby('pentad').mean(('time'))
vor = xr.DataArray(np.zeros((72, 64, 5)), [('pentad', conv_p.pentad),
('lat', rundata.lat),
('year', [2, 3, 4, 5, 6])])
for year in range(2, 7):
rundata = load_year_xr(run_fol, year)
rundata.coords['pentad'] = (rundata.time // 5) + 1
#take time and zonal mean
vor[:, :, year - 2] = rundata.vor[:, 28, :, :].groupby('pentad').mean(
('time', 'lon'))
vor_clim = vor.mean(('year'))
omega = 7.2921150e-5
f = 2 * omega * np.sin(rundata.lat * np.pi / 180)
ro_clim = -1. * vor_clim / f
#plot
ro_clim.plot.contourf(x='lat',
y='pentad',
levels=np.arange(0., 1., 0.1),
add_label=False)
plt.ylim(20, 60)
#plt.xlim(60,180)
plt.xlabel('Latitude')
plt.ylabel('Pentad')
plt.title('206 hPa -vorticity/f')
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol + '/ro206_hm.png')
plt.clf()
def mombudg_fn(inp_fol, nproc, dim):
#inp_fol = experiment name
#nproc = no of processors
#dim = dimension to take eddies around (time or lon)
#load in run data
run_fol = inp_fol + '/np' + nproc
year = 5
rundata = load_year_xr(run_fol, year, pinterp=True)
#Define constants
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi / 180)
f = 2 * omega * np.sin(rundata.lat * np.pi / 180)
#Take first averages: u, v, uv, w, phi, damping terms
u_av = rundata.ucomp.mean((dim))
v_av = rundata.vcomp.mean((dim))
w_av = rundata.omega.mean((dim))
uu_av = rundata.ucomp_sq.mean((dim))
uv_av = rundata.ucomp_vcomp.mean((dim))
uw_av = rundata.ucomp_omega.mean((dim))
phi_av = rundata.height.mean((dim)) * 9.8
#terms that don't need differentiating etc can be immediately averaged
ddamp_tzav = rundata.dt_ug_diffusion.mean(('time', 'lon'))
rdamp_tzav = rundata.udt_rdamp.mean(('time', 'lon'))
fv_tzav = f * rundata.vcomp.mean(('time', 'lon'))
#calculate eddy products
uued_av = uu_av - u_av * u_av
uved_av = (uv_av - u_av * v_av) * coslat * coslat
uwed_av = uw_av - u_av * w_av
#Evaluate gradients needed
if dim == 'time':
dudx_av = xr.DataArray(cfd(u_av.values, u_av.lon * np.pi / 180, 2),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
ududx_av = u_av * dudx_av / coslat / a
dudy_av = xr.DataArray(
cfd((u_av * coslat).values, u_av.lat * np.pi / 180, 1),
[('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
vdudy_av = v_av * dudy_av / coslat / a
dudp_av = xr.DataArray(cfd(u_av.values, u_av.pfull * 100, 0),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
wdudp_av = w_av * dudp_av
duueddx_av = xr.DataArray(
cfd(uued_av.values, uued_av.lon * np.pi / 180, 2),
[('pfull', uued_av.pfull), ('lat', uued_av.lat),
('lon', uued_av.lon)])
duueddx_av = duueddx_av / coslat / a
duveddy_av = xr.DataArray(
cfd(uved_av.values, uved_av.lat * np.pi / 180, 1),
[('pfull', uued_av.pfull), ('lat', uued_av.lat),
('lon', uued_av.lon)])
duveddy_av = duveddy_av / coslat / coslat / a
duweddp_av = xr.DataArray(cfd(uwed_av.values, uwed_av.pfull * 100, 0),
[('pfull', uued_av.pfull),
('lat', uued_av.lat), ('lon', uued_av.lon)])
dphidx_av = xr.DataArray(
cfd(phi_av.values, phi_av.lon * np.pi / 180,
2), [('pfull', phi_av.pfull), ('lat', phi_av.lat),
('lon', phi_av.lon)])
dphidx_av = dphidx_av / coslat / a
dphidx_tzav = dphidx_av.mean(('lon'))
ududx_tzav = ududx_av.mean(('lon'))
vdudy_tzav = vdudy_av.mean(('lon'))
wdudp_tzav = wdudp_av.mean(('lon'))
duueddx_tzav = duueddx_av.mean(('lon'))
duveddy_tzav = duveddy_av.mean(('lon'))
duweddp_tzav = duweddp_av.mean(('lon'))
elif dim == 'lon':
dudy_av = xr.DataArray(
cfd((u_av * coslat).values, u_av.lat * np.pi / 180, 2),
[('time', u_av.time), ('pfull', u_av.pfull), ('lat', u_av.lat)])
vdudy_av = v_av * dudy_av / coslat / a
dudp_av = xr.DataArray(cfd(u_av.values, u_av.pfull * 100, 1),
[('time', u_av.time), ('pfull', u_av.pfull),
('lat', u_av.lat)])
wdudp_av = w_av * dudp_av
duveddy_av = xr.DataArray(
cfd(uved_av.values, uved_av.lat * np.pi / 180,
2), [('time', u_av.time), ('pfull', uued_av.pfull),
('lat', uued_av.lat)])
duveddy_av = duveddy_av / coslat / coslat / a
duweddp_av = xr.DataArray(cfd(uwed_av.values, uwed_av.pfull * 100,
1), [('time', u_av.time),
('pfull', uued_av.pfull),
('lat', uued_av.lat)])
vdudy_tzav = vdudy_av.mean(('time'))
wdudp_tzav = wdudp_av.mean(('time'))
duveddy_tzav = duveddy_av.mean(('time'))
duweddp_tzav = duweddp_av.mean(('time'))
dphidx_tzav = xr.DataArray(np.zeros((40, 64)), [('pfull', u_av.pfull),
('lat', u_av.lat)])
ududx_tzav = xr.DataArray(np.zeros((40, 64)), [('pfull', u_av.pfull),
('lat', u_av.lat)])
duueddx_tzav = xr.DataArray(np.zeros((40, 64)), [('pfull', u_av.pfull),
('lat', u_av.lat)])
#evaluate a sum of the terms
mom_sum = dphidx_tzav - fv_tzav + ududx_tzav + vdudy_tzav + wdudp_tzav + duueddx_tzav + duveddy_tzav + duweddp_tzav - ddamp_tzav
#produce output dataset
data_out = xr.Dataset({
'dphidx_tzav': dphidx_tzav,
'fv_tzav': fv_tzav,
'ddamp_tzav': ddamp_tzav,
'rdamp_tzav': rdamp_tzav,
'ududx_tzav': ududx_tzav,
'vdudy_tzav': vdudy_tzav,
'wdudp_tzav': wdudp_tzav,
'duueddx_tzav': duueddx_tzav,
'duveddy_tzav': duveddy_tzav,
'duweddp_tzav': duweddp_tzav
})
#Plot up terms to check
plt.figure(1)
dphidx_tzav.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
plt.figure(2)
fv_tzav.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
plt.figure(3)
ududx_tzav.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
plt.figure(4)
vdudy_tzav.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
plt.figure(5)
wdudp_tzav.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
plt.figure(6)
duueddx_tzav.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
plt.figure(7)
duveddy_tzav.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
plt.figure(8)
duweddp_tzav.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
plt.figure(9)
mom_sum.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
plt.figure(10)
ddamp_tzav.plot.contourf(x='lat',
y='pfull',
levels=np.arange(-0.0001, 0.0001, 0.00001),
yincrease=False)
#produce Dima&Wallace plots
#plt.figure(1)
#(-duveddy_tzav).plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
#plt.figure(2)
#(fv_tzav-vdudy_tzav).plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
#plt.figure(3)
#(fv_tzav-vdudy_tzav-wdudp_tzav-duveddy_tzav-duweddp_tzav).plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.show()
return data_out
def mombudg_2d_fn(inp_fol, nproc):
#inp_fol = experiment name
#nproc = no of processors
#load in run data
run_fol = inp_fol + '/np' + nproc
year = 5
rundata = load_year_xr(run_fol, year, pinterp=True)
#Define constants
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi / 180)
f = 2 * omega * np.sin(rundata.lat * np.pi / 180)
#Take first averages: u, v, uv, w, phi, damping terms
u_av = rundata.ucomp.mean(('time'))
v_av = rundata.vcomp.mean(('time'))
w_av = rundata.omega.mean(('time'))
phi_av = rundata.height.mean(('time')) * 9.8
#terms that don't need differentiating etc can be immediately averaged
ddamp_av = rundata.dt_vg_diffusion.mean(('time'))
fu_av = f * rundata.ucomp.mean(('time'))
#look at eddies from chosen average
u_ed = rundata.ucomp - u_av
v_ed = rundata.vcomp - v_av
w_ed = rundata.omega - w_av
#calculate eddy products
uved_av = (u_ed * v_ed).mean(('time'))
vved_av = (v_ed * v_ed).mean(('time')) * coslat * coslat
vwed_av = (v_ed * w_ed).mean(('time'))
#Evaluate gradients needed
dvdx_av = xr.DataArray(cfd(v_av.values, u_av.lon * np.pi / 180, 2),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
udvdx_av = u_av * dvdx_av / coslat / a
dvdy_av = xr.DataArray(
cfd((v_av * coslat).values, u_av.lat * np.pi / 180, 1),
[('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
vdvdy_av = v_av * dvdy_av / coslat / a
dvdp_av = xr.DataArray(cfd(u_av.values, u_av.pfull * 100, 0),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
wdvdp_av = w_av * dvdp_av
duveddx_av = xr.DataArray(cfd(uved_av.values, u_av.lon * np.pi / 180, 2),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
duveddx_av = duveddx_av / coslat / a
dvveddy_av = xr.DataArray(cfd(vved_av.values, u_av.lat * np.pi / 180, 1),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
dvveddy_av = dvveddy_av / coslat / coslat / a
dvweddp_av = xr.DataArray(cfd(vwed_av.values, u_av.pfull * 100, 0),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
dphidy_av = xr.DataArray(cfd(phi_av.values, phi_av.lat * np.pi / 180, 1),
[('pfull', phi_av.pfull), ('lat', phi_av.lat),
('lon', phi_av.lon)])
dphidy_av = dphidy_av / a
#evaluate sums of the terms
mom_eddy = -(duveddx_av + dvveddy_av + dvweddp_av)
mom_mean = -(udvdx_av + vdvdy_av + wdvdp_av)
mom_sum = -fu_av + ddamp_av - dphidy_av + mom_eddy + mom_mean
#Make a test plot to see if local balance appears to be achieved...
plt.figure(1)
(dphidy_av[9, :, :] + fu_av[:, 9, :]).plot.contourf(x='lon',
y='lat',
levels=np.arange(
-0.0002, 0.0002,
0.00001))
plt.ylim(-30, 30)
plt.figure(2)
fu_av[:, 9, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.0002, 0.00001))
plt.ylim(-30, 30)
plt.figure(3)
mom_eddy[9, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.0002, 0.00001))
plt.ylim(-30, 30)
plt.figure(4)
mom_mean[9, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.0002, 0.00001))
plt.ylim(-30, 30)
plt.figure(5)
ddamp_av[9, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.0002, 0.00001))
plt.ylim(-30, 30)
plt.figure(6)
mom_sum[:, 9, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.0002, 0.00001))
plt.ylim(-30, 30)
#plt.figure(7)
#(dphidy_av[2,:,:]+fu_av[:,2,:]).plot.contourf(x='lon', y='lat',levels=np.arange(-0.0001,0.0001,0.00001))
#plt.ylim(-30,30)
#plt.figure(8)
#fu_av[:,2,:].plot.contourf(x='lon', y='lat',levels=np.arange(-0.0001,0.0001,0.00001))
#plt.ylim(-30,30)
#plt.figure(9)
#mom_eddy[2,:,:].plot.contourf(x='lon', y='lat',levels=np.arange(-0.0001,0.0001,0.00001))
#plt.ylim(-30,30)
#plt.figure(10)
#mom_mean[2,:,:].plot.contourf(x='lon', y='lat',levels=np.arange(-0.0001,0.0001,0.00001))
#plt.ylim(-30,30)
#plt.figure(11)
#ddamp_av[2,:,:].plot.contourf(x='lon', y='lat',levels=np.arange(-0.0001,0.0001,0.00001))
#plt.ylim(-30,30)
#plt.figure(12)
#mom_sum[:,2,:].plot.contourf(x='lon', y='lat',levels=np.arange(-0.0001,0.0001,0.00001))
#plt.ylim(-30,30)
plt.show()
return
from netCDF4 import Dataset
import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
from time_av_xr import load_year_xr
#set run name
run_fol = 'bucket_test2/np8'
inp_fol = 'bucket_test2'
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/exp/' + inp_fol + '/input/land.nc'
land = xr.open_dataset(land_file)
rundata = load_year_xr(run_fol, [1])
rundata.coords['month'] = (rundata.time // 30) + 1
#take time and zonal mean
bd_tav = rundata.bucket_depth.groupby('month').mean(('time'))
dbdlh_tav = rundata.bucket_depth_lh.groupby('month').mean(('time'))
dbdcd_tav = rundata.bucket_depth_cond.groupby('month').mean(('time'))
dbdcv_tav = rundata.bucket_depth_conv.groupby('month').mean(('time'))
raincv_tav = rundata.convection_rain.groupby('month').mean(('time'))
raincd_tav = rundata.condensation_rain.groupby('month').mean(('time'))
lhe_tav = rundata.flux_lhe.groupby('month').mean(('time'))
totrain_tav = (raincv_tav + raincd_tav) * 86400.
rundata_month = rundata.groupby('month').mean(('time'))
#Load in u, v, t, w and check lat-lon pictures over a 5 year time average - any waves?
#check also monthly mean latitudinal phi structure
from time_av_xr import load_year_xr
import numpy as np
import pandas as pd
import xarray as xr
import matplotlib.pyplot as plt
#Initialise arrays to load into
run_fol = 'flat_10m_diags/np16'
rundata = load_year_xr(run_fol, 2, pinterp=True)
rundata.coords['month'] = (rundata.time // 30) - 11
mngrp = rundata.ucomp.groupby('month').mean(('time'))
u = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
v = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
w = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
phi = xr.DataArray(np.zeros((12, 17, 64, 128, 5)), [('month', mngrp.month),
('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
def mombudg_2d_an_fn(inp_fol, nproc):
#inp_fol = experiment name
#nproc = no of processors
#set run name
run_fol = inp_fol + '/np' + nproc
year = 2
rundata = load_year_xr(run_fol, year, pinterp=True)
#Define constants
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi / 180)
f = 2 * omega * np.sin(rundata.lat * np.pi / 180)
#Initialise arrays to load into
u = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
v = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
w = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
uu = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
uv = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
uw = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
phi = xr.DataArray(np.zeros((17, 64, 128, 5)), [('pfull', rundata.pfull),
('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
ddamp = xr.DataArray(np.zeros((17, 64, 128, 5)),
[('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', [2, 3, 4, 5, 6])])
for year in range(2, 7):
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) + 1
u[:, :, :, year - 2] = rundata.ucomp.mean(('time'))
v[:, :, :, year - 2] = rundata.vcomp.mean(('time'))
w[:, :, :, year - 2] = rundata.omega.mean(('time'))
uu[:, :, :, year - 2] = rundata.ucomp_sq.mean(('time'))
uv[:, :, :, year - 2] = rundata.ucomp_vcomp.mean(('time'))
uw[:, :, :, year - 2] = rundata.ucomp_omega.mean(('time'))
phi[:, :, :, year - 2] = rundata.height.mean(('time'))
ddamp[:, :, :, year - 2] = rundata.dt_ug_diffusion.mean(('time'))
#Take averages: u, v, uv, w, phi, damping terms
u_av = u.mean(('year'))
v_av = v.mean(('year'))
w_av = w.mean(('year'))
uu_av = uu.mean(('year'))
uv_av = uv.mean(('year'))
uw_av = uw.mean(('year'))
phi_av = phi.mean(('year')) * 9.8
ddamp_av = ddamp.mean(('year'))
fv_av = v_av * f
#Evaluate mean eddy products
uued_av = uu_av - u_av * u_av
uved_av = (uv_av - u_av * v_av) * coslat * coslat
uwed_av = uw_av - u_av * w_av
#Evaluate gradients needed
dudx_av = xr.DataArray(cfd(u_av.values, u_av.lon * np.pi / 180, 2),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
ududx_av = u_av * dudx_av / coslat / a
dudy_av = xr.DataArray(
cfd((u_av * coslat).values, u_av.lat * np.pi / 180, 1),
[('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
vdudy_av = v_av * dudy_av / coslat / a
dudp_av = xr.DataArray(cfd(u_av.values, u_av.pfull * 100, 0),
[('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
wdudp_av = w_av * dudp_av
duueddx_av = xr.DataArray(
cfd(uued_av.values, uued_av.lon * np.pi / 180, 2),
[('pfull', uued_av.pfull), ('lat', uued_av.lat), ('lon', uued_av.lon)])
duueddx_av = duueddx_av / coslat / a
duveddy_av = xr.DataArray(
cfd(uved_av.values, uved_av.lat * np.pi / 180, 1),
[('pfull', uued_av.pfull), ('lat', uued_av.lat), ('lon', uued_av.lon)])
duveddy_av = duveddy_av / coslat / coslat / a
duweddp_av = xr.DataArray(cfd(uwed_av.values, uwed_av.pfull * 100, 0),
[('pfull', uued_av.pfull), ('lat', uued_av.lat),
('lon', uued_av.lon)])
dphidx_av = xr.DataArray(cfd(phi_av.values, phi_av.lon * np.pi / 180, 2),
[('pfull', phi_av.pfull), ('lat', phi_av.lat),
('lon', phi_av.lon)])
dphidx_av = -1 * dphidx_av / coslat / a
#evaluate sums of the terms
mom_eddy = -(duueddx_av + duveddy_av + duweddp_av)
mom_mean = -(ududx_av + vdudy_av + wdudp_av)
mom_sum = fv_av + ddamp_av + dphidx_av + mom_eddy + mom_mean
#Plot up seasonal cycle for each term at 200 and 850hPa and save
dphidx_av[9, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.00021,
0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_dphidx_200.png')
plt.clf()
fv_av[9, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.00021, 0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol + '/an_fv_200.png')
plt.clf()
mom_eddy[9, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.00021,
0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_momeddy_200.png')
plt.clf()
mom_mean[9, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.00021,
0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_mommean_200.png')
plt.clf()
ddamp_av[9, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.00021,
0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_ddamp_200.png')
plt.clf()
mom_sum[9, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0002, 0.00021, 0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_momsum_200.png')
plt.clf()
dphidx_av[2, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0001, 0.00011,
0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_dphidx_850.png')
plt.clf()
fv_av[2, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0001, 0.00011, 0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol + '/an_fv_850.png')
plt.clf()
mom_eddy[2, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0001, 0.00011,
0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_momeddy_850.png')
plt.clf()
mom_mean[2, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0001, 0.00011,
0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_mommean_850.png')
plt.clf()
ddamp_av[2, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0001, 0.00011,
0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_ddamp_850.png')
plt.clf()
mom_sum[2, :, :].plot.contourf(x='lon',
y='lat',
levels=np.arange(-0.0001, 0.00011, 0.00001))
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/an_momsum_850.png')
plt.clf()
data_out = xr.Dataset({
'dphidx_av': dphidx_av,
'fv_av': fv_av,
'ddamp_av': ddamp_av,
'mom_eddy': mom_eddy,
'mom_mean': mom_mean,
'mom_sum': mom_sum
})
return data_out
def mombudg_2d_pd_fn(inp_fol, nproc):
#inp_fol = experiment name
#nproc = no of processors
#set run name
run_fol = inp_fol + '/np' + nproc
year = 2
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
mngrp = rundata.ucomp.groupby('pentad').mean(('time'))
#Define constants
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi / 180)
f = 2 * omega * np.sin(rundata.lat * np.pi / 180)
#Initialise arrays to load into
u = xr.DataArray(np.zeros((72, 17, 64, 128, 5)),
[('pentad', mngrp.pentad), ('pfull', rundata.pfull),
('lat', rundata.lat), ('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
v = xr.DataArray(np.zeros((72, 17, 64, 128, 5)),
[('pentad', mngrp.pentad), ('pfull', rundata.pfull),
('lat', rundata.lat), ('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
w = xr.DataArray(np.zeros((72, 17, 64, 128, 5)),
[('pentad', mngrp.pentad), ('pfull', rundata.pfull),
('lat', rundata.lat), ('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
uu = xr.DataArray(np.zeros((72, 17, 64, 128, 5)),
[('pentad', mngrp.pentad), ('pfull', rundata.pfull),
('lat', rundata.lat), ('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
uv = xr.DataArray(np.zeros((72, 17, 64, 128, 5)),
[('pentad', mngrp.pentad), ('pfull', rundata.pfull),
('lat', rundata.lat), ('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
uw = xr.DataArray(np.zeros((72, 17, 64, 128, 5)),
[('pentad', mngrp.pentad), ('pfull', rundata.pfull),
('lat', rundata.lat), ('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
phi = xr.DataArray(np.zeros((72, 17, 64, 128, 5)),
[('pentad', mngrp.pentad), ('pfull', rundata.pfull),
('lat', rundata.lat), ('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
ddamp = xr.DataArray(np.zeros((72, 17, 64, 128, 5)),
[('pentad', mngrp.pentad), ('pfull', rundata.pfull),
('lat', rundata.lat), ('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
for year in range(2, 7):
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
u[:, :, :, :, year - 2] = rundata.ucomp.groupby('pentad').mean(
('time'))
v[:, :, :, :, year - 2] = rundata.vcomp.groupby('pentad').mean(
('time'))
w[:, :, :, :, year - 2] = rundata.omega.groupby('pentad').mean(
('time'))
uu[:, :, :, :, year - 2] = rundata.ucomp_sq.groupby('pentad').mean(
('time'))
uv[:, :, :, :, year - 2] = rundata.ucomp_vcomp.groupby('pentad').mean(
('time'))
uw[:, :, :, :, year - 2] = rundata.ucomp_omega.groupby('pentad').mean(
('time'))
phi[:, :, :, :, year - 2] = rundata.height.groupby('pentad').mean(
('time'))
ddamp[:, :, :, :,
year - 2] = rundata.dt_ug_diffusion.groupby('pentad').mean(
('time'))
#Take averages: u, v, uv, w, phi, damping terms
u_av = u.mean(('year'))
v_av = v.mean(('year'))
w_av = w.mean(('year'))
uu_av = uu.mean(('year'))
uv_av = uv.mean(('year'))
uw_av = uw.mean(('year'))
phi_av = phi.mean(('year')) * 9.8
ddamp_av = ddamp.mean(('year'))
fv_av = v_av * f
#Evaluate mean eddy products
uued_av = uu_av - u_av * u_av
uved_av = (uv_av - u_av * v_av) * coslat * coslat
uwed_av = uw_av - u_av * w_av
#Evaluate gradients needed
dudx_av = xr.DataArray(cfd(u_av.values, u_av.lon * np.pi / 180, 3),
[('pentad', mngrp.pentad), ('pfull', u_av.pfull),
('lat', u_av.lat), ('lon', u_av.lon)])
ududx_av = u_av * dudx_av / coslat / a
dudy_av = xr.DataArray(
cfd((u_av * coslat).values, u_av.lat * np.pi / 180, 2),
[('pentad', mngrp.pentad), ('pfull', u_av.pfull), ('lat', u_av.lat),
('lon', u_av.lon)])
vdudy_av = v_av * dudy_av / coslat / a
dudp_av = xr.DataArray(cfd(u_av.values, u_av.pfull * 100, 1),
[('pentad', mngrp.pentad), ('pfull', u_av.pfull),
('lat', u_av.lat), ('lon', u_av.lon)])
wdudp_av = w_av * dudp_av
duueddx_av = xr.DataArray(
cfd(uued_av.values, uued_av.lon * np.pi / 180, 3),
[('pentad', mngrp.pentad), ('pfull', uued_av.pfull),
('lat', uued_av.lat), ('lon', uued_av.lon)])
duueddx_av = duueddx_av / coslat / a
duveddy_av = xr.DataArray(
cfd(uved_av.values, uved_av.lat * np.pi / 180, 2),
[('pentad', mngrp.pentad), ('pfull', uued_av.pfull),
('lat', uued_av.lat), ('lon', uued_av.lon)])
duveddy_av = duveddy_av / coslat / coslat / a
duweddp_av = xr.DataArray(cfd(uwed_av.values, uwed_av.pfull * 100, 1),
[('pentad', mngrp.pentad),
('pfull', uued_av.pfull), ('lat', uued_av.lat),
('lon', uued_av.lon)])
dphidx_av = xr.DataArray(cfd(phi_av.values, phi_av.lon * np.pi / 180,
3), [('pentad', mngrp.pentad),
('pfull', phi_av.pfull),
('lat', phi_av.lat),
('lon', phi_av.lon)])
dphidx_av = -1 * dphidx_av / coslat / a
#evaluate sums of the terms
mom_eddy = -(duueddx_av + duveddy_av + duweddp_av)
mom_mean = -(ududx_av + vdudy_av + wdudp_av)
mom_sum = fv_av + ddamp_av + dphidx_av + mom_eddy + mom_mean
data_out = xr.Dataset({
'dphidx_av': dphidx_av,
'fv_av': fv_av,
'ddamp_av': ddamp_av,
'mom_eddy': mom_eddy,
'mom_mean': mom_mean,
'mom_sum': mom_sum
})
return data_out
def onset_diag_fn(inp_fol, nproc):
#set run name
run_fol = inp_fol + '/np' + nproc
year = 1
rundata = load_year_xr(run_fol, [year])
rundata.coords['pentad'] = (rundata.time // 5) + 1
conv_p = rundata.convection_rain.groupby('pentad').mean(('time'))
tot_p = xr.DataArray(np.zeros((72, 64, 128, 5)),
[('pentad', conv_p.pentad), ('lat', rundata.lat),
('lon', rundata.lon), ('year', [2, 3, 4, 5, 6])])
onset_pentad_yearly = xr.DataArray(np.zeros((64, 128, 5)),
[('lat', rundata.lat),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
for year in range(2, 7):
rundata = load_year_xr(run_fol, [year])
rundata.coords['pentad'] = (rundata.time // 5) + 1
conv_p = rundata.convection_rain.groupby('pentad').mean(('time'))
cond_p = rundata.condensation_rain.groupby('pentad').mean(('time'))
tot_p[:, :, :, year - 2] = (conv_p + cond_p) * 86400
tot_p_clim = tot_p.mean(('year'))
#mask values of precip less than 9mm/day
tot_p_clim_masked = np.ma.masked_less(tot_p_clim.values, 9)
#locate first unmasked value along pentad axis
onset_index = np.ma.notmasked_edges(tot_p_clim_masked, axis=0)
onset = np.zeros((64, 128))
onset[:] = np.nan
for i in range(0, len(onset_index[0][1])):
onset[onset_index[0][1][i],
onset_index[0][2][i]] = onset_index[0][0][i] + 1
onset_pentad = xr.DataArray(onset, [('lat', rundata.lat),
('lon', rundata.lon)])
land_file = '/scratch/rg419/GFDL_model/GFDLmoistModel/exp/flat_10m/input/land.nc'
land = xr.open_dataset(land_file)
land_plot = xr.DataArray(land.land_mask.values, [('lat', rundata.lat),
('lon', rundata.lon)])
#plot
onset_pentad.plot.pcolormesh(x='lon',
y='lat',
levels=np.arange(34., 48., 1.))
for i in range(0, len(onset_index[0][1])):
if (onset_pentad.lon[onset_index[0][2][i]] > 60
and onset_pentad.lon[onset_index[0][2][i]] < 180
and onset_pentad.lat[onset_index[0][1][i]] > 0
and onset_pentad.lat[onset_index[0][1][i]] < 30
and onset_index[0][0][i] > 5):
plt.text(onset_pentad.lon[onset_index[0][2][i]],
onset_pentad.lat[onset_index[0][1][i]],
onset_index[0][0][i] + 1,
size=8,
rotation='vertical')
land_plot.plot.contour(x='lon',
y='lat',
levels=np.arange(0., 2., 1.),
colors='k',
add_colorbar=False)
plt.ylim(0, 30)
plt.xlim(60, 180)
plt.xlabel('Longitude')
plt.ylabel('Latitude')
plt.title('Onset pentad')
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol +
'/onset_pentad.png')
plt.clf()
return onset_pentad
def vmombudg_fn(inp_fol,nproc,dim):
#inp_fol = experiment name
#nproc = no of processors
#dim = dimension to take eddies around (time or lon)
#load in run data
run_fol = inp_fol + '/np' + nproc
year = 5
rundata = load_year_xr(run_fol, year, pinterp=True)
#Define constants
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi/180)
f = 2 * omega * np.sin(rundata.lat *np.pi/180)
#Take first averages: u, v, uv, w, phi, damping terms
u_av = rundata.ucomp.mean((dim))
v_av = rundata.vcomp.mean((dim))
w_av = rundata.omega.mean((dim))
phi_av = rundata.height.mean((dim))*9.8
#terms that don't need differentiating etc can be immediately averaged
ddamp_tzav = rundata.dt_vg_diffusion.mean(('time','lon'))
rdamp_tzav = rundata.vdt_rdamp.mean(('time','lon'))
fu_tzav = f*rundata.ucomp.mean(('time','lon'))
#look at eddies from chosen average
u_ed = rundata.ucomp - u_av
v_ed = rundata.vcomp - v_av
w_ed = rundata.omega - w_av
#calculate eddy products
vved_av = (v_ed*v_ed).mean((dim))*coslat*coslat
uved_av = (u_ed*v_ed).mean((dim))
vwed_av = (v_ed*w_ed).mean((dim))
#Evaluate gradients needed
if dim == 'time':
dvdx_av = xr.DataArray( cfd(v_av.values,u_av.lon*np.pi/180,2), [('pfull', u_av.pfull ), ('lat', u_av.lat), ('lon', u_av.lon )])
udvdx_av = u_av*dvdx_av/coslat/a
dvdy_av = xr.DataArray( cfd((v_av*coslat).values,u_av.lat*np.pi/180,1), [('pfull', u_av.pfull ), ('lat', u_av.lat), ('lon', u_av.lon )])
vdvdy_av = v_av*dvdy_av/coslat/a
dvdp_av = xr.DataArray( cfd(v_av.values,u_av.pfull*100,0), [('pfull', u_av.pfull ), ('lat', u_av.lat), ('lon', u_av.lon )])
wdvdp_av = w_av*dvdp_av
duveddx_av = xr.DataArray( cfd(uved_av.values,u_av.lon*np.pi/180,2), [('pfull', u_av.pfull ), ('lat', u_av.lat), ('lon', u_av.lon )])
duveddx_av = duveddx_av/coslat/a
dvveddy_av = xr.DataArray( cfd(vved_av.values ,u_av.lat*np.pi/180,1), [('pfull', u_av.pfull ), ('lat', u_av.lat), ('lon', u_av.lon )])
dvveddy_av = dvveddy_av/coslat/coslat/a
dvweddp_av = xr.DataArray( cfd(vwed_av.values,u_av.pfull*100,0), [('pfull', u_av.pfull ), ('lat', u_av.lat), ('lon', u_av.lon )])
dphidy_av = xr.DataArray( cfd(phi_av.values,phi_av.lat*np.pi/180,1), [('pfull', phi_av.pfull ), ('lat', phi_av.lat), ('lon', phi_av.lon )])
dphidy_av = dphidy_av/a
dphidy_tzav = dphidy_av.mean(('lon'))
udvdx_tzav = udvdx_av.mean(('lon'))
vdvdy_tzav = vdvdy_av.mean(('lon'))
wdvdp_tzav = wdvdp_av.mean(('lon'))
duveddx_tzav = duveddx_av.mean(('lon'))
dvveddy_tzav = dvveddy_av.mean(('lon'))
dvweddp_tzav = dvweddp_av.mean(('lon'))
#evaluate a sum of the terms
mom_sum = dphidy_tzav + fu_tzav + udvdx_tzav + vdvdy_tzav + wdvdp_tzav + duveddx_tzav + dvveddy_tzav + dvweddp_tzav - ddamp_tzav
#Plot up terms to check
plt.figure(1)
(dphidy_tzav+fu_tzav).plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.figure(2)
udvdx_tzav.plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.figure(3)
vdvdy_tzav.plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.figure(4)
wdvdp_tzav.plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.figure(5)
duveddx_tzav.plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.figure(6)
dvveddy_tzav.plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.figure(7)
dvweddp_tzav.plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.figure(8)
mom_sum.plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.figure(9)
ddamp_tzav.plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
#produce Dima&Wallace plots
#plt.figure(1)
#(-duveddy_tzav).plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
#plt.figure(2)
#(fv_tzav-vdudy_tzav).plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
#plt.figure(3)
#(fv_tzav-vdudy_tzav-wdudp_tzav-duveddy_tzav-duweddp_tzav).plot.contourf(x='lat', y='pfull',levels=np.arange(-0.0001,0.0001,0.00001), yincrease=False)
plt.show()
return
def u_hm_fn(inp_fol, nproc):
#set run name
run_fol = inp_fol + '/np' + nproc
year = 1
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) + 1
conv_p = rundata.convection_rain[:, 37, :].groupby('pentad').mean(('time'))
u_low = xr.DataArray(np.zeros((72, 128, 5)), [('pentad', conv_p.pentad),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
u_high = xr.DataArray(np.zeros((72, 128, 5)), [('pentad', conv_p.pentad),
('lon', rundata.lon),
('year', [2, 3, 4, 5, 6])])
for year in range(2, 7):
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) + 1
#take time and zonal mean
u_low[:, :, year - 2] = rundata.ucomp[:, 2,
37, :].groupby('pentad').mean(
('time'))
u_high[:, :, year - 2] = rundata.ucomp[:, 9,
37, :].groupby('pentad').mean(
('time'))
u_low_clim = u_low.mean(('year'))
u_high_clim = u_high.mean(('year'))
#plot
u_low_clim.plot.contourf(x='lon',
y='pentad',
levels=np.arange(-30., 31., 5.),
add_label=False)
plt.plot([float(rundata.lonb[21]),
float(rundata.lonb[21])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[26]),
float(rundata.lonb[26])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[30]),
float(rundata.lonb[30])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[35]),
float(rundata.lonb[35])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[39]),
float(rundata.lonb[39])], [20., 60.], 'k')
plt.ylim(20, 60)
plt.xlim(60, 180)
plt.xlabel('Longitude')
plt.ylabel('Pentad')
plt.title('850 hPa zonal wind speed at 15N, m/s')
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol + '/u850_hm.png')
plt.clf()
u_high_clim.plot.contourf(x='lon',
y='pentad',
levels=np.arange(-30., 31., 5.),
add_label=False)
plt.plot([float(rundata.lonb[21]),
float(rundata.lonb[21])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[26]),
float(rundata.lonb[26])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[30]),
float(rundata.lonb[30])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[35]),
float(rundata.lonb[35])], [20., 60.], 'k')
plt.plot([float(rundata.lonb[39]),
float(rundata.lonb[39])], [20., 60.], 'k')
plt.ylim(20, 60)
plt.xlim(60, 180)
plt.xlabel('Longitude')
plt.ylabel('Pentad')
plt.title('200 hPa zonal wind speed at 15N, m/s')
plt.savefig('/scratch/rg419/workdir_moist/' + inp_fol + '/u200_hm.png')
plt.clf()
#NB may also need to account for damping terms
#Start by doing this for an annual average, can extend after
import sys
sys.path.append('/scratch/rg419/workdir_moist/python_bin')
from time_av_xr import load_year_xr
import numpy as np
import matplotlib.pyplot as plt
from StreamFunction import cal_stream_fn
import xarray as xr
#load in run data
run_fol = 'aquaplanet_10m/np16'
year = 7
rundata = load_year_xr(run_fol, [year])
data = {
'omega': np.ma.asarray((rundata.omega.mean(('lon'))).values),
'vcomp': np.ma.asarray((rundata.vcomp.mean(('lon'))).values),
'lat': rundata.lat.values,
'time': rundata.time.values,
'pfull': rundata.pfull.values
}
psi = cal_stream_fn(data, 1)
psi.shape
plt.figure(1)
plt.contourf(rundata.lat, rundata.pfull, np.mean(psi, 0))
plt.colorbar()