def uv_partitioning(data, start_index):
u_tav = data.ucomp[start_index:start_index +
4, :, :, :].mean('xofyear')
u_ztav = u_tav.sel(lon=lons).mean('lon')
v_tav = data.vcomp[start_index:start_index +
4, :, :, :].mean('xofyear')
v_ztav = v_tav.sel(lon=lons).mean('lon')
uv_trans = (data.ucomp_vcomp[start_index:start_index +
4, :, :, :].mean('xofyear') -
u_tav * v_tav).sel(lon=lons).mean('lon')
uv_stat = (u_tav * v_tav).sel(lon=lons).mean('lon') - u_ztav * v_ztav
uv_conv_trans = xr.DataArray(
cfd((uv_trans * coslat * coslat).values, data.lat * np.pi / 180,
1), [('pfull', data.pfull), ('lat', data.lat)])
uv_conv_trans = -86400. * uv_conv_trans / coslat / coslat / a
uv_conv_stat = xr.DataArray(
cfd((uv_stat * coslat * coslat).values, data.lat * np.pi / 180, 1),
[('pfull', data.pfull), ('lat', data.lat)])
uv_conv_stat = -86400. * uv_conv_stat / coslat / coslat / a
return uv_conv_trans, uv_conv_stat, u_ztav
def wind_comp_uv(run, months, filename='atmos_pentad', timeav='pentad', period_fac=1.):
data = time_means(run, months, filename=filename, timeav=timeav,period_fac=period_fac)
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi/180)
f = 2 * omega * np.sin(data.lat *np.pi/180)
dudy = xr.DataArray( cfd((data.ucomp*coslat).values,data.lon*np.pi/180.,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dvdy = xr.DataArray( cfd((data.vcomp*coslat).values,data.lon*np.pi/180.,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dlat=xr.DataArray(data.latb.diff('latb').values*np.pi/180, coords=[('lat', data.lat)])
v_div = np.cumsum(dvdy * dlat, axis=2)/coslat
u_vort = np.cumsum(dudy * dlat, axis=2)/coslat
u_div = data.ucomp - u_vort
v_vort = data.vcomp - v_div
plt.figure(1)
v_vort[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(2)
v_div[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(3)
data.vcomp[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(4)
u_vort[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.figure(5)
u_div[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.figure(6)
data.ucomp[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.show()
def calc_eddy_heat(ut, vt, wt):
#Calculate momentum advection by eddies
a = 6376.0e3 #radius used in model
coslat = np.cos(ut.lat * np.pi / 180.)
vt = vt * coslat
duteddx = xr.DataArray(cfd(ut.values, vt.lon * np.pi / 180., 3),
[('xofyear', vt.xofyear), ('pfull', vt.pfull),
('lat', vt.lat), ('lon', vt.lon)])
duteddx = duteddx / coslat / a
dvteddy = xr.DataArray(cfd(vt.values, vt.lat * np.pi / 180., 2),
[('xofyear', vt.xofyear), ('pfull', vt.pfull),
('lat', vt.lat), ('lon', vt.lon)])
dvteddy = dvteddy / coslat / a
dwteddp = xr.DataArray(cfd(wt.values, vt.pfull * 100., 1),
[('xofyear', vt.xofyear), ('pfull', vt.pfull),
('lat', vt.lat), ('lon', vt.lon)])
dwteddp = dwteddp
out = duteddx + dvteddy + dwteddp
return dvteddy, dwteddp
def u_gradients(data):
data['dudx'] = (('xofyear','pfull','lat','lon'), cfd(data.ucomp.values,data.lon*np.pi/180.,3) )
data['dudx'] = (('xofyear','pfull','lat','lon'), data.dudx/coslat/a )
data['dudy'] = (('xofyear','pfull','lat','lon'), cfd((data.ucomp*coslat).values,data.lat*np.pi/180.,2) )
data['dudy'] = (('xofyear','pfull','lat','lon'), data.dudy/coslat/a )
data['dudp'] = (('xofyear','pfull','lat','lon'), cfd(data.ucomp.values,data.pfull*100.,1) )
def u_transients(data,lons):
uu_ed = data.ucomp_sq - data.ucomp*data.ucomp
uv_ed = data.ucomp_vcomp - data.ucomp*data.vcomp
uw_ed = data.ucomp_omega - data.ucomp*data.omega
data['duudx_ed'] = (('xofyear','pfull','lat','lon'), cfd(uu_ed.values,data.lon*np.pi/180.,3) )
data['duudx_ed'] = (('xofyear','pfull','lat','lon'), data.duudx_ed/coslat/a )
data['duvdy_ed'] = (('xofyear','pfull','lat','lon'), cfd((uv_ed*coslat*coslat).values,data.lat*np.pi/180.,2) )
data['duvdy_ed'] = (('xofyear','pfull','lat','lon'), data.duvdy_ed/coslat/coslat/a )
data['duwdp_ed'] = (('xofyear','pfull','lat','lon'), cfd(uw_ed.values,data.pfull*100.,1) )
data['mom_trans'] = (('xofyear','pfull','lat'), -1.*(data.duudx_ed[:,:,:,lons] + data.duvdy_ed[:,:,:,lons] + data.duwdp_ed[:,:,:,lons]).mean('lon') )
def ddp(field):
field_dp = xr.DataArray(cfd(field.values, field.pfull * 100.,
1), [('xofyear', field.xofyear),
('pfull', field.pfull),
('lat', field.lat)])
field_dp = -86400. * field_dp
return field_dp
def mombudg_closure_fn(data):
#Evaluate momentum budget
#Define constants
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi/180)
f = 2 * omega * np.sin(data.lat *np.pi/180)
mom_div = -1.*calc_eddy_mom(data.ucomp_sq, data.ucomp_vcomp, data.ucomp_omega)
dphidx = xr.DataArray( cfd(data.height.values,data.lon*np.pi/180,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dphidx = -1.*dphidx/coslat/a
fv = data.vcomp*f
mom_sum = (fv + dphidx*9.8 + mom_div)*10000.
mom_sum[45,9,:,:].plot.contourf(x='lon', y='lat',levels=np.arange(-2.,2.1,0.2))
plt.show()
return
def fv_and_dphidx(data):
data['dphidx'] = (('xofyear', 'pfull', 'lat', 'lon'),
cfd(data.height.values, data.lon * np.pi / 180., 3))
data['dphidx'] = (('xofyear', 'pfull', 'lat', 'lon'),
data.dphidx / coslat / a * -9.8)
data['fv'] = (('xofyear', 'pfull', 'lat', 'lon'), data.vcomp * f)
def ddy(field, vector=True, uv=False, a=6376.0e3):
"""Calculate d/dy of a given DataArray. DataArray must include lat dimension.
kwargs: vector - specify if input field is vector or scalar
prod - if a vector, is the field uv?"""
try:
field.coords['lat']
except:
raise NameError('Coord lat not found')
coslat = np.cos(field.lat * np.pi / 180)
if vector and uv:
cosfac = coslat**2
elif vector:
cosfac = coslat
else:
cosfac = 1.
field_dy = cfd((field * cosfac).values, field.lat.values * np.pi / 180,
field.get_axis_num('lat'))
field_dy = xr.DataArray(field_dy, dims=field.dims, coords=field.coords)
field_dy = field_dy / cosfac / a
return field_dy
def mombudg_fn(data):
#Evaluate momentum budget
#Define constants
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi/180)
f = 2 * omega * np.sin(data.lat *np.pi/180)
uu_ms, uu_stat, uu_trans = eddy_decomp(data.ucomp, data.ucomp, data.ucomp_sq)
uv_ms, uv_stat, uv_trans = eddy_decomp(data.ucomp, data.vcomp, data.ucomp_vcomp)
uw_ms, uw_stat, uw_trans = eddy_decomp(data.ucomp, data.omega, data.ucomp_omega)
print 'eddy decomposition done'
mom_mean = calc_mean_mom( uv_ms, uw_ms)
print 'mean advective terms done'
mom_trans = calc_eddy_mom(uu_trans, uv_trans, uw_trans)
print 'transient terms done'
mom_stat = calc_eddy_mom(uu_stat, uv_stat, uw_stat)
print 'stationary terms done'
dphidx = xr.DataArray( cfd(data.height.values,data.lon*np.pi/180,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dphidx = -1*dphidx/coslat/a
fv = data.vcomp*f
data_out = xr.Dataset({'fv': fv, 'ddamp': data.dt_ug_diffusion, 'rdamp': data.udt_rdamp, 'dphidx': dphidx,
'mom_mean': mom_mean, 'mom_trans': mom_trans, 'mom_stat': mom_stat})
return data_out
def mombudg_closure_fn(data):
#Evaluate momentum budget
#Define constants
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi/180)
f = 2 * omega * np.sin(data.lat *np.pi/180)
mom_div = -1.*calc_eddy_mom(data.ucomp_sq, data.ucomp_vcomp, data.ucomp_omega)
dphidx = xr.DataArray( cfd(data.height.values,data.lon*np.pi/180,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dphidx = -1.*dphidx/coslat/a
fv = data.vcomp*f
mom_sum = (fv + dphidx*9.8 + mom_div + data.dt_ug_diffusion)*10000.
plt.figure(1)
mom_sum.mean(('xofyear','lon')).plot.contourf(x='lat', y='pfull', yincrease=False, levels=np.arange(-0.6,0.65,0.05))
plt.figure(2)
(mom_div*10000.).mean(('xofyear','lon')).plot.contourf(x='lat', y='pfull', yincrease=False, levels=np.arange(-0.6,0.65,0.05))
plt.show()
return
def calc_mean_mom(uv,uw):
#Calculate momentum advection by time and zonal mean
a= 6376.0e3 #radius used in model
coslat = np.cos(uv.lat * np.pi/180)
uv = uv*coslat*coslat
duvdy = xr.DataArray( cfd(uv.values ,uv.lat*np.pi/180,2), [('xofyear', uv.xofyear ), ('pfull', uv.pfull ), ('lat', uv.lat)])
duvdy = duvdy/coslat/coslat/a
duwdp = xr.DataArray( cfd(uw.values,uv.pfull*100,1), [('xofyear', uv.xofyear ), ('pfull', uv.pfull ), ('lat', uv.lat)])
duwdp = duwdp
out = duvdy + duwdp
return out
def ddx(field):
a = 6376.0e3 #radius used in model
coslat = np.cos(field.lat * np.pi / 180)
field_dx = xr.DataArray(cfd(field.values, field.lon * np.pi / 180, 2),
[('xofyear', field.xofyear), ('lat', field.lat),
('lon', field.lon)])
field_dx = -86400. * field_dx / coslat / a
return field_dx
def calc_eddy_mom(uu,uv,uw):
#Calculate momentum advection by eddies
a= 6376.0e3 #radius used in model
coslat = np.cos(uu.lat * np.pi/180.)
uv = uv*coslat*coslat
duueddx = xr.DataArray( cfd(uu.values,uu.lon*np.pi/180,3), [('xofyear', uu.xofyear ), ('pfull', uu.pfull ), ('lat', uu.lat), ('lon', uu.lon )])
duueddx = duueddx/coslat/a
duveddy = xr.DataArray( cfd(uv.values ,uu.lat*np.pi/180,2), [('xofyear', uu.xofyear ), ('pfull', uu.pfull ), ('lat', uu.lat), ('lon', uu.lon )])
duveddy = duveddy/coslat/coslat/a
duweddp = xr.DataArray( cfd(uw.values,uu.pfull*100,1), [('xofyear', uu.xofyear ), ('pfull', uu.pfull ), ('lat', uu.lat), ('lon', uu.lon )])
duweddp = duweddp
out = duueddx + duveddy + duweddp
return out
def ddp(field, pname='pfull'):
"""Calculate d/dp of a given DataArray. DataArray must include a pfull dimension"""
try:
field.coords[pname]
except:
raise NameError('Coord ' + pname + ' not found')
field_dp = cfd( field.values, field[pname].values*100., field.get_axis_num(pname) )
field_dp = xr.DataArray( field_dp, dims = field.dims, coords = field.coords )
return field_dp
def calc_mean_heat(vt, wt):
#Calculate momentum advection by time and zonal mean
a = 6376.0e3 #radius used in model
coslat = np.cos(vt.lat * np.pi / 180)
vt = vt * coslat
dvtdy = xr.DataArray(cfd(vt.values, vt.lat * np.pi / 180, 2),
[('xofyear', vt.xofyear), ('pfull', vt.pfull),
('lat', vt.lat)])
dvtdy = duvdy / coslat / a
dwtdp = xr.DataArray(cfd(uw.values, vt.pfull * 100, 1),
[('xofyear', vt.xofyear), ('pfull', vt.pfull),
('lat', vt.lat)])
dwtdp = dwtdp
out = dvtdy + dwtdp
return out
def ddy(field, prod=True):
a = 6376.0e3 #radius used in model
coslat = np.cos(field.lat * np.pi / 180)
cosfac = coslat**2. if prod == True else coslat
field_cosfac = (field * cosfac)
field_dy = xr.DataArray(
cfd(field_cosfac.values, field.lat * np.pi / 180, 1),
[('xofyear', field.xofyear), ('lat', field.lat)])
field_dy = -86400. * field_dy / cosfac / a
return field_dy
def ddt(field, timedir = 'xofyear', secperunit = 5.*86400., cyclic=True):
"""Calculate d/dt in unit/s of a given DataArray. DataArray must include a time dimension
Define seconds per unit time using secperunit. Default calc is for pentads"""
try:
field.coords[timedir]
except:
raise NameError('Coord ' + timedir + ' not found')
field_dt = cfd( field.values, field.coords[timedir].values*secperunit, field.get_axis_num(timedir), cyclic=cyclic)/secperunit/2.
field_dt = xr.DataArray( field_dt, dims = field.dims, coords = field.coords )
return field_dt
def ddp(field):
"""Calculate d/dp of a given DataArray. DataArray must include pfull dimension"""
try:
field.coords['pfull']
except:
raise NameError('Coord pfull not found')
field_dp = cfd(field.values, field.pfull.values * 100.,
field.get_axis_num('pfull'))
field_dp = xr.DataArray(field_dp, dims=field.dims, coords=field.coords)
return field_dp
def wind_components(run, months, filename='atmos_pentad', timeav='pentad', period_fac=1.):
data = time_means(run, months, filename=filename, timeav=timeav,period_fac=period_fac)
omega = 7.2921150e-5
a= 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi/180)
f = 2 * omega * np.sin(data.lat *np.pi/180)
dudx = xr.DataArray( cfd(data.ucomp.values,data.lon*np.pi/180.,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dudx = dudx #/coslat/a Don't bother with geometric factor as will just be multiplied by in integration
dvdx = xr.DataArray( cfd(data.vcomp.values,data.lon*np.pi/180.,3), [('xofyear', data.xofyear ), ('pfull', data.pfull ), ('lat', data.lat), ('lon', data.lon )])
dvdx = dvdx #/coslat/a
dvdy_div = data.div - dudx
dudy_vort = dvdx - data.vor
dlat=xr.DataArray(data.latb.diff('latb').values*np.pi/180, coords=[('lat', data.lat)])
v_div = np.cumsum(dvdy_div * dlat, axis=2)/coslat
u_vort = np.cumsum(dudy_vort * dlat, axis=2)/coslat
u_div = data.ucomp - u_vort
v_vort = data.vcomp - v_div
plt.figure(1)
v_vort[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(2)
v_div[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(3)
data.vcomp[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-10,11,1))
plt.figure(4)
u_vort[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.figure(5)
u_div[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.figure(6)
data.ucomp[40,15,:,:].plot.contourf(x='lon', y='lat', levels=np.arange(-80,81,10))
plt.show()
def ddx(field, a = 6376.0e3, latname='lat', lonname='lon'):
"""Calculate d/dx of a given DataArray. DataArray must include a lat and lon dimensions"""
try:
field.coords[lonname]
except:
raise NameError('Coord ' + lonname + ' not found')
try:
field.coords[latname]
except:
raise NameError('Coord ' + latname + ' not found')
coslat = np.cos(field[latname] * np.pi/180)
field_dx = cfd( field.values, field[lonname].values*np.pi/180, field.get_axis_num(lonname) )
field_dx = xr.DataArray( field_dx, dims = field.dims, coords = field.coords )
field_dx = field_dx/coslat/a
return field_dx
def ddx(field, a=6376.0e3):
"""Calculate d/dx of a given DataArray. DataArray must include lat and lon dimensions"""
try:
field.coords['lon']
except:
raise NameError('Coord lon not found')
try:
field.coords['lat']
except:
raise NameError('Coord lat not found')
coslat = np.cos(field.lat * np.pi / 180)
field_dx = cfd(field.values, field.lon.values * np.pi / 180,
field.get_axis_num('lon'))
field_dx = xr.DataArray(field_dx, dims=field.dims, coords=field.coords)
field_dx = field_dx / coslat / a
return field_dx
def uv_partitioning(data, start_index, lons):
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
u_ztav = data.ucomp[start_index:start_index +
4, :, :, :].sel(lon=lons).mean(('xofyear', 'lon'))
v_ztav = data.vcomp[start_index:start_index +
4, :, :, :].sel(lon=lons).mean(('xofyear', 'lon'))
uv_eddy = data.ucomp_vcomp[start_index:start_index +
4, :, :, :].sel(lon=lons).mean(
('xofyear', 'lon')) - u_ztav * v_ztav
uv_conv = xr.DataArray(
cfd((uv_eddy * coslat * coslat).values, data.lat * np.pi / 180, 1),
[('pfull', data.pfull), ('lat', data.lat)])
uv_conv = -86400. * uv_conv / coslat / coslat / a
return uv_conv, u_ztav
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
def mom_adv_terms_fn(run_fol,years, trange):
year = years[0]
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
mngrp = rundata.ucomp.groupby('pentad').mean(('time'))
trl = trange[1]-trange[0]
#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((trl,17,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('pfull', rundata.pfull), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
v = xr.DataArray(np.zeros((trl,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
w = xr.DataArray(np.zeros((trl,17,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('pfull', rundata.pfull), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
uu = xr.DataArray(np.zeros((trl,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
uv = xr.DataArray(np.zeros((trl,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
uw = xr.DataArray(np.zeros((trl,17,64,128,len(years))), [('pentad', range(trange[0]+1,trange[1]+1) ), ('pfull', rundata.pfull), ('lat', rundata.lat), ('lon', rundata.lon), ('year', years )])
for year in years:
print year
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
u[:,:,:,:,year-years[0]] = rundata.ucomp.groupby('pentad').mean(('time'))[trange[0]:trange[1],:,:,:]
v[:,:,:,year-years[0]] = rundata.vcomp.groupby('pentad').mean(('time'))[trange[0]:trange[1],9,:,:]
w[:,:,:,:,year-years[0]] = rundata.omega.groupby('pentad').mean(('time'))[trange[0]:trange[1],:,:,:]
uu[:,:,:,year-years[0]] = rundata.ucomp_sq.groupby('pentad').mean(('time'))[trange[0]:trange[1],9,:,:]
uv[:,:,:,year-years[0]] = rundata.ucomp_vcomp.groupby('pentad').mean(('time'))[trange[0]:trange[1],9,:,:]
uw[:,:,:,:,year-years[0]] = rundata.ucomp_omega.groupby('pentad').mean(('time'))[trange[0]:trange[1],:,:,:]
#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'))
#Evaluate mean eddy products
uued_av = uu_av - u_av[:,9,:,:]*u_av[:,9,:,:]
uved_av = (uv_av - u_av[:,9,:,:]*v_av)*coslat*coslat
uwed_av = uw_av - u_av*w_av
#Evaluate gradients needed
dudx_av = xr.DataArray( cfd(u_av[:,9,:,:].values,u_av.lon*np.pi/180,2), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', u_av.lat), ('lon', u_av.lon )])
ududx_av = u_av[:,9,:,:]*dudx_av/coslat/a
dudy_av = xr.DataArray( cfd((u_av[:,9,:,:]*coslat).values,u_av.lat*np.pi/180,1), [('pentad', range(trange[0]+1,trange[1]+1) ), ('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', range(trange[0]+1,trange[1]+1) ), ('pfull', u_av.pfull ), ('lat', u_av.lat), ('lon', u_av.lon )])
wdudp_av = w_av[:,9,:,:]*dudp_av[:,9,:,:]
duueddx_av = xr.DataArray( cfd(uued_av.values,uued_av.lon*np.pi/180,2), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', u_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), [('pentad', range(trange[0]+1,trange[1]+1) ), ('lat', uved_av.lat), ('lon', uved_av.lon )])
duveddy_av = duveddy_av/coslat/coslat/a
duweddp_av = xr.DataArray( cfd(uwed_av.values,uwed_av.pfull*100,1), [('pentad', range(trange[0]+1,trange[1]+1) ), ('pfull', uwed_av.pfull ), ('lat', u_av.lat), ('lon', uwed_av.lon )])
duweddp_av = duweddp_av[:,9,:,:]
data_out = xr.Dataset({'ududx_av': ududx_av, 'vdudy_av': vdudy_av, 'wdudp_av': wdudp_av,
'duueddx_av': duueddx_av, 'duveddy_av': duveddy_av, 'duweddp_av': duweddp_av, 'u_av': u_av[:,9,:,:], 'v_av': v_av })
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 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 heatbudg_2d_pd_fn(run_fol, years, trange, heatbudg=True, humbudg=False):
year = years[0]
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
mngrp = rundata.ucomp.groupby('pentad').mean(('time'))
trl = trange[1] - trange[0]
#Define constants
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(rundata.lat * np.pi / 180)
#Initialise arrays to load into
u = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
v = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
w = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
if heatbudg == True:
t = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
if humbudg == True:
q = xr.DataArray(np.zeros((trl, 17, 64, 128, len(years))),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', rundata.pfull), ('lat', rundata.lat),
('lon', rundata.lon), ('year', years)])
if heatbudg == False and humbudg == False:
print 'No budget specified'
return
for year in years:
print year
rundata = load_year_xr(run_fol, year, pinterp=True)
rundata.coords['pentad'] = (rundata.time // 5) - 71
u[:, :, :, :, year - years[0]] = rundata.ucomp.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
v[:, :, :, :, year - years[0]] = rundata.vcomp.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
w[:, :, :, :, year - years[0]] = rundata.omega.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
if heatbudg == True:
t[:, :, :, :,
year - years[0]] = rundata.temp.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
if humbudg == True:
q[:, :, :, :,
year - years[0]] = rundata.sphum.groupby('pentad').mean(
('time'))[trange[0]:trange[1], :, :, :]
#Take averages: u, v, uv, w, phi, damping terms
u_av = u.mean(('year'))
v_av = v.mean(('year'))
w_av = w.mean(('year'))
#Evaluate eddy quantities
u_ed = u - u_av
v_ed = v - v_av
w_ed = w - w_av
if heatbudg == True:
t_av = t.mean(('year'))
t_ed = t - t_av
ut_ed = (u_ed * t_ed).mean(('year'))
vt_ed = (v_ed * t_ed).mean(('year')) * coslat
wt_ed = (w_ed * t_ed).mean(('year'))
#Evaluate gradients needed
dtdx_av = xr.DataArray(
cfd(t_av.values, t_av.lon * np.pi / 180, 3),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
udtdx_av = u_av * dtdx_av / coslat / a
dtdy_av = xr.DataArray(
cfd(t_av.values, t_av.lat * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
vdtdy_av = v_av * dtdy_av / a
dtdp_av = xr.DataArray(
cfd(t_av.values, t_av.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
wdtdp_av = w_av * dtdp_av
duteddx_av = xr.DataArray(
cfd(ut_ed.values, ut_ed.lon * np.pi / 180, 3),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
duteddx_av = duteddx_av / coslat / a
dvteddy_av = xr.DataArray(
cfd(vt_ed.values, vt_ed.lat * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
dvteddy_av = dvteddy_av / coslat / a
dwteddp_av = xr.DataArray(
cfd(wt_ed.values, wt_ed.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
#evaluate sums of the terms
heat_eddy = -(duteddx_av + dvteddy_av + dwteddp_av)
heat_mean = -(udtdx_av + vdtdy_av + wdtdp_av)
heat_sum = heat_eddy + heat_mean
if humbudg == True:
q_av = q.mean(('year'))
q_ed = q - q_av
uq_ed = (u_ed * q_ed).mean(('year'))
vq_ed = (v_ed * q_ed).mean(('year')) * coslat
wq_ed = (w_ed * q_ed).mean(('year'))
#Evaluate gradients needed
dqdx_av = xr.DataArray(
cfd(q_av.values, q_av.lon * np.pi / 180, 3),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
udqdx_av = u_av * dqdx_av / coslat / a
dqdy_av = xr.DataArray(
cfd(q_av.values, q_av.lat * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
vdqdy_av = v_av * dqdy_av / a
dqdp_av = xr.DataArray(
cfd(q_av.values, q_av.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
wdqdp_av = w_av * dqdp_av
duqeddx_av = xr.DataArray(
cfd(uq_ed.values, uq_ed.lon * np.pi / 180, 3),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
duqeddx_av = duqeddx_av / coslat / a
dvqeddy_av = xr.DataArray(
cfd(vq_ed.values, vq_ed.lat * np.pi / 180, 2),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
dvqeddy_av = dvqeddy_av / coslat / a
dwqeddp_av = xr.DataArray(
cfd(wq_ed.values, wq_ed.pfull * 100, 1),
[('pentad', range(trange[0] + 1, trange[1] + 1)),
('pfull', u_av.pfull), ('lat', u_av.lat), ('lon', u_av.lon)])
#evaluate sums of the terms
hum_eddy = -(duqeddx_av + dvqeddy_av + dwqeddp_av)
hum_mean = -(udqdx_av + vdqdy_av + wdqdp_av)
hum_sum = hum_eddy + hum_mean
if heatbudg == True and humbudg == True:
data_out = xr.Dataset({
'hum_eddy': hum_eddy,
'hum_mean': hum_mean,
'hum_sum': hum_sum,
'heat_eddy': heat_eddy,
'heat_mean': heat_mean,
'heat_sum': heat_sum,
'u_av': u_av,
'v_av': v_av
})
elif heatbudg == True:
data_out = xr.Dataset({
'heat_eddy': heat_eddy,
'heat_mean': heat_mean,
'heat_sum': heat_sum,
'u_av': u_av,
'v_av': v_av
})
elif humbudg == True:
data_out = xr.Dataset({
'hum_eddy': hum_eddy,
'hum_mean': hum_mean,
'hum_sum': hum_sum,
'u_av': u_av,
'v_av': v_av
})
return data_out
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 vort_av(run,
months,
filename='atmos_pentad',
timeav='pentad',
period_fac=1.):
data = time_means(run,
months,
filename=filename,
timeav=timeav,
period_fac=period_fac)
omega = 7.2921150e-5
a = 6376.0e3 #radius used in model
coslat = np.cos(data.lat * np.pi / 180)
f = 2 * omega * np.sin(data.lat * np.pi / 180)
abs_vort = data.vor + f
dvortdx = xr.DataArray(cfd(abs_vort.values, data.lon * np.pi / 180., 3),
[('xofyear', data.xofyear), ('pfull', data.pfull),
('lat', data.lat), ('lon', data.lon)])
dvortdx = dvortdx / coslat / a
dvortdy = xr.DataArray(cfd(abs_vort.values, data.lat * np.pi / 180., 2),
[('xofyear', data.xofyear), ('pfull', data.pfull),
('lat', data.lat), ('lon', data.lon)])
dvortdy = dvortdy / a
dvortdp = xr.DataArray(cfd(abs_vort.values, data.pfull * 100., 1),
[('xofyear', data.xofyear), ('pfull', data.pfull),
('lat', data.lat), ('lon', data.lon)])
vort_adv = data.ucomp * dvortdx + data.vcomp * dvortdy
vort_vert = data.omega * dvortdp
vort_div = abs_vort * data.div
precip = (data.convection_rain + data.condensation_rain) * 86400.
plt.figure(1)
vort_adv[:, 15, :, :].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-4.5e-10, 4.6e-10,
0.5e-10))
plt.figure(2)
vort_div[:, 15, :, :].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-4.5e-10, 4.6e-10,
0.5e-10))
plt.figure(3)
vort_vert[:, 15, :, :].mean('lon').plot.contourf(x='xofyear', y='lat')
plt.figure(4)
abs_vort[:, 15, :, :].mean('lon').plot.contourf(x='xofyear',
y='lat',
levels=np.arange(
-2.4e-4, 2.5e-4,
0.2e-4))
precip.mean('lon').plot.contour(x='xofyear',
y='lat',
levels=np.arange(6., 31., 6.),
add_label=False,
add_colorbar=False,
colors='k')
plt.show()
