def calculate_relvor(self, lev=850):
from windspharm.standard import VectorWind
# Load the full U and V fields
u_full = self['uzonal_{}hPa'.format(lev)].values
v_full = self['umeridional_{}hPa'.format(lev)].values
relvor = np.zeros(np.shape(u_full))
# Loop through all the valid times and ensemble members
for i1 in range(np.shape(u_full)[0]):
for i2 in range(np.shape(u_full)[1]):
# Create the spherical harmonics vector object
u = u_full[i1, i2, ::-1, :] # lats must go from N to S
v = v_full[i1, i2, ::-1, :] # lats must go from N to S
wnd = VectorWind(u, v, gridtype='regular')
# Calculate the relative vorticity
relvor[i1, i2, :, :] = wnd.vorticity()[::-1, :]
# Assign relative vorticity as a new variable
vorvar = xarray.DataArray(relvor,
dims=self['uzonal_{}hPa'.format(lev)].dims)
assignvar = {'relvor_{}hPa'.format(lev): vorvar}
self.update(self.assign(**assignvar))
def windsfarm(): #Prepare wund for vector w = VectorWind(uwnd, vwnd) #Calculating basic variables eta = w.absolutevorticity() div = w.divergence() uchi, vchi = w.irrotationalcomponent() etax, etay = w.gradient(eta) planvor = w.planetaryvorticity S = -eta * div - uchi * etax - vchi * etay S = recover_data(S, uwnd_info) return(w, eta, div, uchi, vchi, etax, etay, planvor, S)
def rws(uwnd, vwnd):
from windspharm.standard import VectorWind
uwnd = np.transpose(uwnd, (1, 2, 0))
vwnd = np.transpose(vwnd, (1, 2, 0))
w = VectorWind(uwnd, vwnd)
eta = w.absolutevorticity()
div = w.divergence()
uchi, vchi = w.irrotationalcomponent()
etax, etay = w.gradient(eta)
s1 = -eta * div
s2 = - (uchi * etax + vchi * etay)
s = s1 + s2
s = np.mean(s, axis=2)
s1 = np.mean(s1, axis=2)
s2 = np.mean(s2, axis=2)
return s, s1, s2
def windprops(ua, va, lat, lon, exps=['All-Nat', 'SST-Nat']):
from windspharm.standard import VectorWind
def diverg(u, v, lat, lon):
meshlon, meshlat = np.meshgrid(lon, lat)
phi = meshlat * np.pi / 180
cs_phi = np.cos(phi)
dphi = lat * np.pi / 180
dtheta = lon * np.pi / 180
a = 6.371e6
v_y = np.gradient(v * cs_phi, dphi, axis=1) / (a * cs_phi)
u_x = np.gradient(u, dtheta, axis=2) / (a * cs_phi)
div = u_x + v_y
return div
u = ua[exps[0]]
v = va[exps[0]]
u_ghg = ua[exps[1]]
v_ghg = va[exps[1]]
div = diverg(u, v, lat, lon)
div_ghg = diverg(u_ghg, v_ghg, lat, lon)
u = np.transpose(u, (1, 2, 0))
v = np.transpose(v, (1, 2, 0))
u_ghg = np.transpose(u_ghg, (1, 2, 0))
v_ghg = np.transpose(v_ghg, (1, 2, 0))
w = VectorWind(u, v)
w_ghg = VectorWind(u_ghg, v_ghg)
vp = w.velocitypotential()
vp_ghg = w_ghg.velocitypotential()
vp = np.transpose(vp, (2, 0, 1))
vp_ghg = np.transpose(vp_ghg, (2, 0, 1))
delta_vp = vp_ghg - vp
delta_div = div_ghg - div
return div, delta_div, vp, delta_vp
#--------------- Calcul du potentiel de vitesses et de son gradient ----------------------- #--------- = partie divergente du champ = partie par laquelle se fait l export --------------- #--------------- MOYENNE ----------------------- #--- The standard interface requires that latitude and longitude be the leading ---- #--- dimensions of the input wind components, and that wind components must be ---- #--- either 2D or 3D arrays. The data read in is 3D and has latitude and ---- #--- longitude as the last dimensions. ---- #--- It is also required that the latitude dimension is north-to-south. Again the ---- #--- bundled tools make this easy. ---- lats_r, uwnd, vwnd = order_latdim(lats,UE_mean_pin,VE_mean_pin) #- Create a VectorWind instance to handle the computation of streamfunction and velocity potential- w = VectorWind(uwnd, vwnd) #--- fonction de courant (sf ; streamfunction) ---- #--- potentiel de vitesse (vp ; velocity potential) ---- sf, vp = w.sfvp() #--- partie divergente du champ = gradient du potentiel de vitesses ---- grad_vp = w.irrotationalcomponent() #--- on masque les continents parce que les valeurs sur les continents sont trop fortes ---- pourcter_ma = ma.array(pourcter,mask=pourcter>0.1) #--- attention on a re-inverse les latitudes pour pouvoir utiliser VectorWind pour calculer ----
#--------------- MOYENNE ----------------------- #--- The standard interface requires that latitude and longitude be the leading ---- #--- dimensions of the input wind components, and that wind components must be ---- #--- either 2D or 3D arrays. The data read in is 3D and has latitude and ---- #--- longitude as the last dimensions. ---- #--- It is also required that the latitude dimension is north-to-south. Again the ---- #--- bundled tools make this easy. ---- UE_mean_pin_o = np.mean(UE_mean_pin, axis=0) VE_mean_pin_o = np.mean(VE_mean_pin, axis=0) lats_r, uwnd, vwnd = order_latdim(lats, UE_mean_pin_o, VE_mean_pin_o) #- Create a VectorWind instance to handle the computation of streamfunction and velocity potential- w = VectorWind(uwnd, vwnd) #--- fonction de courant (sf ; streamfunction) ---- #--- potentiel de vitesse (vp ; velocity potential) ---- sf, vp = w.sfvp() #--- partie divergente du champ = gradient du potentiel de vitesses ---- grad_vp = w.irrotationalcomponent() #--- on masque les continents parce que les valeurs sur les continents sont trop fortes ---- pourcter_ma = ma.array(pourcter, mask=pourcter > 0.1) #--- attention on a re-inverse les latitudes pour pouvoir utiliser VectorWind pour calculer ----
def rws_fprime(ua200, va200):
from windspharm.standard import VectorWind
from netcdfread import ncread, ncsave
from scipy import interpolate
lon42 = ncread(
'/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc',
'lon')
lat42 = ncread(
'/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc',
'lat')
lon = ncread(
'/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/All-Hist/mon/tas/item3236_monthly_mean_a011_2006-01_2016-12.nc',
'longitude0')
lat = ncread(
'/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/Plus15-Future_LCO2/day/ua/item15201_daily_mean_a00b_2090-01_2100-12.nc',
'latitude1')
u = ua200['All-Nat'] #[1]
v = va200['All-Nat'] #[1]
u_sst = ua200['GHG-Nat'] #[1]
v_sst = va200['GHG-Nat'] #[1]
u_ghg = ua200['SST-Nat'] #[1]
v_ghg = va200['SST-Nat'] #[1]
v = w5rem(v)
u = w5rem(u)
uwnd = np.zeros((64, 128, len(u)))
vwnd = np.zeros((64, 128, len(v)))
uwnd_sst = np.zeros((64, 128, len(u_sst)))
vwnd_sst = np.zeros((64, 128, len(u_sst)))
uwnd_ghg = np.zeros((64, 128, len(u_ghg)))
vwnd_ghg = np.zeros((64, 128, len(u_ghg)))
u = np.transpose(u, (1, 2, 0))
v = np.transpose(v, (1, 2, 0))
u_sst = np.transpose(u_sst, (1, 2, 0))
v_sst = np.transpose(v_sst, (1, 2, 0))
u_ghg = np.transpose(u_ghg, (1, 2, 0))
v_ghg = np.transpose(v_ghg, (1, 2, 0))
for i in range(np.ma.size(uwnd, axis=2)):
h = interpolate.interp2d(lon, lat[::-1], u[:, :, i])
uwnd[:, :, i] = h(lon42, lat42[::-1])
g = interpolate.interp2d(lon, lat[::-1], v[:, :, i])
vwnd[:, :, i] = g(lon42, lat42[::-1])
for i in range(np.ma.size(uwnd_sst, axis=2)):
h = interpolate.interp2d(lon, lat[::-1], u_sst[:, :, i])
uwnd_sst[:, :, i] = h(lon42, lat42[::-1])
g = interpolate.interp2d(lon, lat[::-1], v_sst[:, :, i])
vwnd_sst[:, :, i] = g(lon42, lat42[::-1])
for i in range(np.ma.size(uwnd_ghg, axis=2)):
h = interpolate.interp2d(lon, lat[::-1], u_ghg[:, :, i])
uwnd_ghg[:, :, i] = h(lon42, lat42[::-1])
g = interpolate.interp2d(lon, lat[::-1], v_ghg[:, :, i])
vwnd_ghg[:, :, i] = g(lon42, lat42[::-1])
w = VectorWind(uwnd, vwnd)
w_sst = VectorWind(uwnd_sst, vwnd_sst)
w_ghg = VectorWind(uwnd_ghg, vwnd_ghg)
eta = w.absolutevorticity()
eta_sst = w_sst.absolutevorticity()
eta_ghg = w_ghg.absolutevorticity()
div = w.divergence()
div_sst = w_sst.divergence()
div_ghg = w_ghg.divergence()
uchi, vchi = w.irrotationalcomponent()
uchi_sst, vchi_sst = w_sst.irrotationalcomponent()
uchi_ghg, vchi_ghg = w_ghg.irrotationalcomponent()
etax, etay = w.gradient(eta)
# etax_sst, etay_sst = w_sst.gradient(eta_sst)
# etax_ghg, etay_ghg = w_ghg.gradient(eta_ghg)
eta = np.transpose(eta, (2, 0, 1))
eta_sst = np.transpose(eta_sst, (2, 0, 1))
eta_ghg = np.transpose(eta_ghg, (2, 0, 1))
div = np.transpose(div, (2, 0, 1))
div_sst = np.transpose(div_sst, (2, 0, 1))
div_ghg = np.transpose(div_ghg, (2, 0, 1))
etax = np.transpose(etax, (2, 0, 1))
etay = np.transpose(etay, (2, 0, 1))
uchi = np.transpose(uchi, (2, 0, 1))
uchi_sst = np.transpose(uchi_sst, (2, 0, 1))
uchi_ghg = np.transpose(uchi_ghg, (2, 0, 1))
vchi = np.transpose(vchi, (2, 0, 1))
vchi_sst = np.transpose(vchi_sst, (2, 0, 1))
vchi_ghg = np.transpose(vchi_ghg, (2, 0, 1))
#f_ghg = -eta.mean(axis=0)*(div_ghg-div.mean(axis=0))-(uchi_ghg-uchi.mean(axis=0))*etax.mean(axis=0)-(vchi_ghg-vchi.mean(axis=0))*etay.mean(axis=0)
#f_sst = -eta.mean(axis=0)*(div_sst-div.mean(axis=0))-(uchi_sst-uchi.mean(axis=0))*etax.mean(axis=0)-(vchi_sst-vchi.mean(axis=0))*etay.mean(axis=0)
f_ghg = -eta * (div_ghg - div) - ((uchi_ghg - uchi) * etax +
(vchi_ghg - vchi) * etay)
f_sst = -eta * (div_sst - div) - ((uchi_sst - uchi) * etax +
(vchi_sst - vchi) * etay)
meshlon, meshlat = np.meshgrid(lon42, lat42)
'''
ncsave('/home/bakerh/Downloads/vort200_control', lat42, lon42, eta.mean(axis=0)-2*np.sin(meshlat*np.pi/180)*7.2921e-5, 'vorticity')
ncsave('/home/bakerh/Downloads/vort200_sst', lat42, lon42, eta_sst.mean(axis=0)-2*np.sin(meshlat*np.pi/180)*7.2921e-5, 'vorticity')
ncsave('/home/bakerh/Downloads/vort200_ghg', lat42, lon42, eta_ghg.mean(axis=0)-2*np.sin(meshlat*np.pi/180)*7.2921e-5, 'vorticity')
ncsave('/home/bakerh/Downloads/forcing_ghg', lat42, lon42, f_ghg.mean(axis=0), 'forcing')
ncsave('/home/bakerh/Downloads/forcing_sst', lat42, lon42, f_sst.mean(axis=0), 'forcing')
'''
month = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
ncsave('/home/bakerh/Downloads/vort200_control', month, lat42, lon42,
eta - 2 * np.sin(meshlat * np.pi / 180) * 7.2921e-5, 'vorticity')
ncsave('/home/bakerh/Downloads/vort200_sst', month, lat42, lon42,
eta_sst - 2 * np.sin(meshlat * np.pi / 180) * 7.2921e-5,
'vorticity')
ncsave('/home/bakerh/Downloads/vort200_ghg', month, lat42, lon42,
eta_ghg - 2 * np.sin(meshlat * np.pi / 180) * 7.2921e-5,
'vorticity')
ncsave('/home/bakerh/Downloads/forcing_ghg', month, lat42, lon42, f_ghg,
'forcing')
ncsave('/home/bakerh/Downloads/forcing_sst', month, lat42, lon42, f_sst,
'forcing')
def ming_forcing(ua, va, psl):
from windspharm.standard import VectorWind
from netcdfread import ncread, ncsave
from scipy import interpolate
lon42 = ncread(
'/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc',
'lon')
lat42 = ncread(
'/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc',
'lat')
lon = ncread(
'/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/All-Hist/mon/tas/item3236_monthly_mean_a011_2006-01_2016-12.nc',
'longitude0')
lat = ncread(
'/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/Plus15-Future_LCO2/day/ua/item15201_daily_mean_a00b_2090-01_2100-12.nc',
'latitude1')
lat145 = ncread(
'/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/All-Hist/mon/tas/item3236_monthly_mean_a011_2006-01_2016-12.nc',
'latitude0')
u = ua['All-Nat'][1]
v = va['All-Nat'][1]
u_sst = ua['GHG-Nat'][1]
v_sst = va['GHG-Nat'][1]
u_ghg = ua['SST-Nat'][1]
v_ghg = va['SST-Nat'][1]
p = psl['All-Nat'][1]
p_sst = psl['GHG-Nat'][1]
p_ghg = psl['SST-Nat'][1]
uwnd = np.zeros((64, 128, len(u)))
vwnd = np.zeros((64, 128, len(v)))
uwnd_sst = np.zeros((64, 128, len(u_sst)))
vwnd_sst = np.zeros((64, 128, len(u_sst)))
uwnd_ghg = np.zeros((64, 128, len(u_ghg)))
vwnd_ghg = np.zeros((64, 128, len(u_ghg)))
ps = np.zeros((64, 128, len(p)))
ps_sst = np.zeros((64, 128, len(p_sst)))
ps_ghg = np.zeros((64, 128, len(p_ghg)))
u = np.transpose(u, (1, 2, 0))
v = np.transpose(v, (1, 2, 0))
u_sst = np.transpose(u_sst, (1, 2, 0))
v_sst = np.transpose(v_sst, (1, 2, 0))
u_ghg = np.transpose(u_ghg, (1, 2, 0))
v_ghg = np.transpose(v_ghg, (1, 2, 0))
p = np.transpose(p, (1, 2, 0))
p_sst = np.transpose(p_sst, (1, 2, 0))
p_ghg = np.transpose(p_ghg, (1, 2, 0))
for i in range(np.ma.size(uwnd, axis=2)):
h = interpolate.interp2d(lon, lat[::-1], u[:, :, i])
uwnd[:, :, i] = h(lon42, lat42[::-1])
g = interpolate.interp2d(lon, lat[::-1], v[:, :, i])
vwnd[:, :, i] = g(lon42, lat42[::-1])
j = interpolate.interp2d(lon, lat145[::-1], p[:, :, i])
ps[:, :, i] = j(lon42, lat42[::-1])
for i in range(np.ma.size(uwnd_sst, axis=2)):
h = interpolate.interp2d(lon, lat[::-1], u_sst[:, :, i])
uwnd_sst[:, :, i] = h(lon42, lat42[::-1])
g = interpolate.interp2d(lon, lat[::-1], v_sst[:, :, i])
vwnd_sst[:, :, i] = g(lon42, lat42[::-1])
j = interpolate.interp2d(lon, lat145[::-1], p_sst[:, :, i])
ps_sst[:, :, i] = j(lon42, lat42[::-1])
for i in range(np.ma.size(uwnd_ghg, axis=2)):
h = interpolate.interp2d(lon, lat[::-1], u_ghg[:, :, i])
uwnd_ghg[:, :, i] = h(lon42, lat42[::-1])
g = interpolate.interp2d(lon, lat[::-1], v_ghg[:, :, i])
vwnd_ghg[:, :, i] = g(lon42, lat42[::-1])
j = interpolate.interp2d(lon, lat145[::-1], p_ghg[:, :, i])
ps_ghg[:, :, i] = j(lon42, lat42[::-1])
w = VectorWind(uwnd, vwnd)
w_sst = VectorWind(uwnd_sst, vwnd_sst)
w_ghg = VectorWind(uwnd_ghg, vwnd_ghg)
psx, psy = w.gradient(ps)
psx_sst, psy_sst = w_sst.gradient(ps_sst)
psx_ghg, psy_ghg = w_ghg.gradient(ps_ghg)
omega = np.mean(uwnd * psx + vwnd * psy, axis=2)
omega_sst = uwnd_sst * psx_sst + vwnd_sst * psy_sst
omega_ghg = uwnd_ghg * psx_ghg + vwnd_ghg * psy_ghg
omega_sst = np.transpose(omega_sst, (2, 0, 1))
omega_ghg = np.transpose(omega_ghg, (2, 0, 1))
oforc_sst = np.mean(omega_sst - omega, axis=0)
oforc_ghg = np.mean(omega_ghg - omega, axis=0)
return omega, oforc_sst, oforc_ghg
# extract data f=Dataset(model, mode='r') uwnd = f.variables['ua'][0,0,:,:] vwnd = f.variables['va'][0,0,:,:] uwnd2 = np.array(uwnd) vwnd2 = np.array(vwnd) lons = f.variables['lon'][:] lats = f.variables['lat'][:] f.close() plots=plt.subplot(2,3, count+1) print count # velocity potential calculation - requires some data formatting using Windspharm's tools uwnd, uwnd_info = prep_data(uwnd, 'yx') vwnd, vwnd_info = prep_data(vwnd, 'yx') lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd) w = VectorWind(uwnd, vwnd, gridtype='gaussian') vp = w.velocitypotential() vp = recover_data(vp, uwnd_info) vp = vp*1e-06 # divergence calculation divg = w.divergence() divg = recover_data(divg, uwnd_info) divg = divg*1e6 #divergent wind compotnents calculation uchi,vchi = w.irrotationalcomponent() uchi = np.squeeze(uchi) vchi = np.squeeze(vchi) # print some data dimensions - useful for colourbar range etc print np.shape(lons) print np.shape(lats) print np.shape(uchi)
ncv.close() # The standard interface requires that latitude and longitude be the leading # dimensions of the input wind components, and that wind components must be # either 2D or 3D arrays. The data read in is 3D and has latitude and # longitude as the last dimensions. The bundled tools can make the process of # re-shaping the data a lot easier to manage. uwnd, uwnd_info = prep_data(uwnd, 'tyx') vwnd, vwnd_info = prep_data(vwnd, 'tyx') # It is also required that the latitude dimension is north-to-south. Again the # bundled tools make this easy. lats, uwnd, vwnd = order_latdim(lats, uwnd, 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) # Combine the components to form the Rossby wave source term. Re-shape the # Rossby wave source array to the 4D shape of the wind components as they were # read off files. S = -eta * div - (uchi * etax + vchi * etay) S = recover_data(S, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is
def mass_stream_func(uu, vv, lat, lvl, Global=True):
'''
Name:
MASS_STREAM_FUNC
Purpose:
An IDL procedure to compute the zonal-mean meridional stream
function
Inputs:
uu : 3-D array of zonal winds [nlvl, nlat, nlon]
vv : 3-D array of meridional winds [nlvl, nlat, nlon]
lat : A 1- or 2-D array of latitude values (degrees) [nlat]
lvl : A 1- or 2-D array of pressure levels (Pa) [nlvl]
Outputs:
psi : The meridional mass stream function (kg s**-1)
p : Pressures for the stream function (Pa)
Keywords:
Global : Sets if computing for global or for a
zonal subset using method of Zhang and Wang (2013)
Author and history:
Kyle R. Wodzicki
'''
if (not Global) and (not VectorWind):
raise Exception('Failed to import windspharm.standard.VectorWind!!!')
elif (not Global):
uu, uu_info = prep_data(uu, 'pyx') # Prepare data for VectorWind class
vv, vv_info = prep_data(vv, 'pyx') # Prepare data for VectorWind class
lat, uu, vv = order_latdim(lat, uu, vv) # Fix latitude order
VW = VectorWind(uu, vv) # Initialize vector wind class
uchi, vchi = VW.irrotationalcomponent() # Get irrotational components
vv = recover_data(
vchi, vv_info) # Convert v-irrot component back to input order
vv = np.nanmean(vv, axis=2) # Average over longitude (last dimension)
dims = vv.shape # Get dimensions of VV
if (lat.ndim == 1): # If the latitude is only 1-D
if (dims[0] == lat.size):
vv = vv.transpose()
dims = vv.shape
lat = np.repeat(lat.reshape(1, dims[1]), dims[0],
axis=0) # Reshape to match vv array
lat = lat[:-1, :]
if (lvl.ndim == 1): # If the level is only 1-D
lvl = np.repeat(lvl.reshape(dims[0], 1), dims[1],
axis=1) # Reshape to match vv array
revFlat = False
if (lvl[0, 0] > lvl[-1, 0]): # If pressure levels are decending
revFlat = True
lvl = lvl[::-1, :] # Reverse to ascending
vv = vv[::-1, :] # Reverse vv too
dp = lvl[1:, :] - lvl[:-1, :] # Compute change in pressure between levels
dv = (vv[1:, :] + vv[:-1, :]) / 2.0 # Compute mean wind for level
p = ((lvl[1:, 0] + lvl[:-1, 0]) /
2.0).flatten() # Reform mean pressure for level for output
psi = 2.0 * np.pi * R_e * np.cos(
np.radians(lat)) / g # Compute scalar for psi equation
psi *= np.cumsum(dv * dp,
axis=0) # Multiply scalar by the integeral of vv * dp
return psi, p # Return stream function and pressure
# The standard interface requires that latitude and longitude be the leading
# dimensions of the input wind components, and that wind components must be
# either 2D or 3D arrays. The data read in is 3D and has latitude and
# longitude as the last dimensions. The bundled tools can make the process of
# re-shaping the data a lot easier to manage.
uwnd1, uwnd_info = prep_data(np.squeeze(uwnd[:,ilev,:,:]), 'tyx')
vwnd1, vwnd_info = prep_data(np.squeeze(vwnd[:,ilev,:,:]), 'tyx')
# It is also required that the latitude dimension is north-to-south. Again the
# bundled tools make this easy.
lats, uwnd1, vwnd1 = order_latdim(lats, uwnd1, vwnd1)
# Create a VectorWind instance to handle the computation of streamfunction and
# velocity potential.
w = VectorWind(uwnd1, vwnd1)
# Compute the streamfunction and velocity potential. Also use the bundled
# tools to re-shape the outputs to the 4D shape of the wind components as they
# were read off files.
sf, vp = w.sfvp()
sf = recover_data(sf, uwnd_info)
vp = recover_data(vp, uwnd_info)
#print sf.shape
#print sf[1,1,1]
print("Streamfunction done")
#---NetCDF write---------------------------------------------------------------
print("Start NetCDF writing")
U0 = scipy.sparse.linalg.spsolve(A, F0)
U[i, :] = U0
test = np.fft.ifft(U)
from windspharm.standard import VectorWind
from windspharm.tools import prep_data, recover_data, order_latdim
uwnd, uwnd_info = prep_data(u_n[4:-4, 4:-4], 'xyp')
vwnd, vwnd_info = prep_data(v_n[4:-4, 4:-4], 'xyp')
lats, uwnd, vwnd = order_latdim(phi, uwnd, vwnd)
w = VectorWind(uwnd, vwnd)
sf, vp = w.sfvp()
sf = recover_data(sf, uwnd_info)
vp = recover_data(vp, vwnd_info)
uchi, vchi, upsi, vpsi = w.helmholtz()
uchi = recover_data(uchi, uwnd_info)
vchi = recover_data(vchi, uwnd_info)
upsi = recover_data(upsi, uwnd_info)
vpsi = recover_data(vpsi, uwnd_info)
s = w.s
#convert vort + div from A-grid model to u,v using Spharm
vrt, vrt_info = prep_data(vort_n[4:-4, 4:-4], 'xyp')
div, div_info = prep_data(div_n[4:-4, 4:-4], 'xyp')
lats, vrt, div = order_latdim(phi, vrt, div)
latn, uwnd_neg, vwnd_neg = order_latdim(lat, uwnd_neg, vwnd_neg) hituwnd_pos, hituwndinfo_pos = prep_data(hitum_pos, 'tyx') hitvwnd_pos, hitvwndinfo_pos = prep_data(hitvm_pos, 'tyx') hituwnd_neg, hituwndinfo_neg = prep_data(hitum_neg, 'tyx') hitvwnd_neg, hitvwndinfo_neg = prep_data(hitvm_neg, 'tyx') latn, hituwnd_pos, hitvwnd_pos = order_latdim(lat, hituwnd_pos, hitvwnd_pos) latn, hituwnd_neg, hitvwnd_neg = order_latdim(lat, hituwnd_neg, hitvwnd_neg) ### Change lat/lon to mesh lon2n, lat2n = np.meshgrid(lon, latn) ### Calculate VelocityWind object wpos = VectorWind(uwnd_pos, vwnd_pos) wneg = VectorWind(uwnd_neg, vwnd_neg) hitwpos = VectorWind(hituwnd_pos, hitvwnd_pos) hitwneg = VectorWind(hituwnd_neg, hitvwnd_neg) ### Calculate Absolute Velocity avrt_pos = wpos.absolutevorticity() avrt_neg = wneg.absolutevorticity() hitavrt_pos = hitwpos.absolutevorticity() hitavrt_neg = hitwneg.absolutevorticity() ### Change dimensions fictapos = recover_data(avrt_pos, uwndinfo_pos) fictaneg = recover_data(avrt_neg, uwndinfo_neg)
def calc_quantity(uwnd, vwnd, quantity, lat_axis, lon_axis, axis_order):
"""Calculate a single wind quantity.
Args:
uwnd (numpy.ndarray): Zonal wind
vwnd (numpy.ndarray): Meridional wind
quantity (str): Quantity to be calculated
lat_axis (list): Latitude axis values
lon_axis (list): Longitude axis values
axis_order (str): e.g. tyx
Design:
windsparm requires the input data to be on a global grid
(due to the spherical harmonic representation used),
latitude and longitude to be the leading axes and the
latitude axis must go from 90 to -90. The cdms2 interface
is supposed to adjust for these things but I've found that
things come back upsidedown if the lat axis isn't right, so
I've just used the standard interface here instead.
Reference:
ajdawson.github.io/windspharm
"""
check_global(lat_axis, lon_axis)
# Make latitude and longitude the leading coordinates
uwnd, uwnd_info = prep_data(numpy.array(uwnd), axis_order)
vwnd, vwnd_info = prep_data(numpy.array(vwnd), axis_order)
# Make sure latitude dimension is north-to-south
lats, uwnd, vwnd = order_latdim(lat_axis, uwnd, vwnd)
flip_lat = False if lats[0] == lat_axis[0] else True
w = VectorWind(uwnd, vwnd)
data_out = {}
if quantity == 'rossbywavesource':
eta = w.absolutevorticity()
div = w.divergence()
uchi, vchi = w.irrotationalcomponent()
etax, etay = w.gradient(eta)
data_out['rws1'] = (-eta * div) / (1.e-11)
data_out['rws2'] = (-(uchi * etax + vchi * etay)) / (1.e-11)
data_out['rws'] = data_out['rws1'] + data_out['rws2']
elif quantity == 'magnitude':
data_out['spd'] = w.magnitude()
elif quantity == 'vorticity':
data_out['vrt'] = w.vorticity()
elif quantity == 'divergence':
div = w.divergence()
data_out['div'] = div / (1.e-6)
elif quantity == 'absolutevorticity':
avrt = w.absolutevorticity()
data_out['avrt'] = avrt / (1.e-5)
elif quantity == 'absolutevorticitygradient':
avrt = w.absolutevorticity()
ugrad, vgrad = w.gradient(avrt)
avrtgrad = numpy.sqrt(numpy.square(ugrad) + numpy.square(vgrad))
data_out['avrtgrad'] = avrtgrad / (1.e-5)
elif quantity == 'planetaryvorticity':
data_out['pvrt'] = w.planetaryvorticity()
elif quantity == 'irrotationalcomponent':
data_out['uchi'], data_out['vchi'] = w.irrotationalcomponent()
elif quantity == 'nondivergentcomponent':
data_out['upsi'], data_out['vpsi'] = w.nondivergentcomponent()
elif quantity == 'streamfunction':
sf = w.streamfunction()
data_out['sf'] = sf / (1.e+6)
elif quantity == 'velocitypotential':
vp = w.velocitypotential()
data_out['vp'] = vp / (1.e+6)
else:
sys.exit('Wind quantity not recognised')
# Return data to its original shape
for key in data_out.keys():
data_out[key] = recover_structure(data_out[key], flip_lat, uwnd_info)
return data_out
# The standard interface requires that latitude and longitude be the leading # dimensions of the input wind components, and that wind components must be # either 2D or 3D arrays. The data read in is 3D and has latitude and # longitude as the last dimensions. The bundled tools can make the process of # re-shaping the data a lot easier to manage. uwnd, uwnd_info = prep_data(uwnd, 'tyx') vwnd, vwnd_info = prep_data(vwnd, 'tyx') # It is also required that the latitude dimension is north-to-south. Again the # bundled tools make this easy. lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd) # Create a VectorWind instance to handle the computation of streamfunction and # velocity potential. w = VectorWind(uwnd, vwnd) # Compute the streamfunction and velocity potential. Also use the bundled # tools to re-shape the outputs to the 4D shape of the wind components as they # were read off files. sf, vp = w.sfvp() sf = recover_data(sf, uwnd_info) vp = recover_data(vp, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is # for plotting purposes). sf_dec, lons_c = addcyclic(sf[11], lons) vp_dec, lons_c = addcyclic(vp[11], lons) # Plot streamfunction. m = Basemap(projection='cyl', resolution='c', llcrnrlon=0, llcrnrlat=-90,
um=u/coslat[:,None]
vm=v/coslat[:,None]
# um checked!!!
for j in range(0,np.size(um,axis=0)) :
print lats[j], ym[j],u[j,i],um[j,i]
# Create a VectorWind instance to handle the computations.
uwnd, uwnd_info = prep_data(uwnd, 'tyx')
vwnd, vwnd_info = prep_data(vwnd, 'tyx')
#w = VectorWind(uwnd, vwnd)
# Compute absolute vorticity
w = VectorWind(u, v)
q = w.absolutevorticity()
qbar = q
#qbar = np.average(q[:,:,nt[nt>0]],axis=2)
print "qbar(4,0)=",qbar[4,0]
#qbar checked!!!
print "------------------------------"
print "gradients"
print np.version.version
print " "
print "----- wind and q gradients ---------"
# umx, umy = w.gradient(u)
#coslat[0]=0 # a very small number is used instead #coslat[-1]=0 # ----"""---- # velocity in the Mercator projection um = u / coslat[:, None] vm = v / coslat[:, None] # um checked!!! # Create a VectorWind instance to handle the computations. uwnd, uwnd_info = prep_data(uwnd, 'tyx') vwnd, vwnd_info = prep_data(vwnd, 'tyx') #w = VectorWind(uwnd, vwnd) # Compute absolute vorticity w = VectorWind(u, v) q = w.absolutevorticity() #qbar = q #qbar checked!!! #----BetaM--------------------------------------------------------------------- print 'Calculate BetaM' cos2 = coslat * coslat # dum, cosuy=w.gradient(u*cos2[:,None]) # dum, cosuyy = w.gradient(cosuy/coslat[:,None]) # # tmp = 2*e_omega *cos2/radius
ncv.close() # The standard interface requires that latitude and longitude be the leading # dimensions of the input wind components, and that wind components must be # either 2D or 3D arrays. The data read in is 3D and has latitude and # longitude as the last dimensions. The bundled tools can make the process of # re-shaping the data a lot easier to manage. uwnd, uwnd_info = prep_data(uwnd, "tyx") vwnd, vwnd_info = prep_data(vwnd, "tyx") # It is also required that the latitude dimension is north-to-south. Again the # bundled tools make this easy. lats, uwnd, vwnd = order_latdim(lats, uwnd, 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) # Combine the components to form the Rossby wave source term. Re-shape the # Rossby wave source array to the 4D shape of the wind components as they were # read off files. S = -eta * div - (uchi * etax + vchi * etay) S = recover_data(S, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is
def rws_fprime_co2(ua200, va200):
from windspharm.standard import VectorWind
from netcdfread import ncread, ncsave
from scipy import interpolate
lon42 = ncread(
'/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc',
'lon')
lat42 = ncread(
'/network/aopp/hera/mad/bakerh/BTmodel_COR/main/inputs/forcing_ghg.nc',
'lat')
# HadAM3P
lon = ncread(
'/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/All-Hist/mon/tas/item3236_monthly_mean_a011_2006-01_2016-12.nc',
'longitude0')
lat = ncread(
'/network/aopp/hera/mad/bakerh/HAPPI/HadAM3P-N96/Plus15-Future_LCO2/day/ua/item15201_daily_mean_a00b_2090-01_2100-12.nc',
'latitude1')
# MIROC5
# lat = ncread('/network/aopp/hera/mad/bakerh/HAPPI/MIROC5/All-Hist/day/tas/tas_Aday_MIROC5_All-Hist_est1_v2-0_run001_20060101-20161231.nc','lat')
# lon = np.arange(0, 360, 360/256)
# CAM4
#lat = np.linspace(-90, 90, 96)
#lon = np.linspace(0, 357.5, 144)
u = ua200['Plus15-Future_LCO2']
v = va200['Plus15-Future_LCO2']
u_ghg = ua200['Plus15-Future_HCO2']
v_ghg = va200['Plus15-Future_HCO2']
uwnd = np.zeros((64, 128, len(u)))
vwnd = np.zeros((64, 128, len(v)))
uwnd_ghg = np.zeros((64, 128, len(u_ghg)))
vwnd_ghg = np.zeros((64, 128, len(u_ghg)))
u = np.transpose(u, (1, 2, 0))
v = np.transpose(v, (1, 2, 0))
u_ghg = np.transpose(u_ghg, (1, 2, 0))
v_ghg = np.transpose(v_ghg, (1, 2, 0))
for i in range(np.ma.size(uwnd, axis=2)):
h = interpolate.interp2d(lon, lat[::-1], u[:, :, i])
uwnd[:, :, i] = h(lon42, lat42[::-1])
g = interpolate.interp2d(lon, lat[::-1], v[:, :, i])
vwnd[:, :, i] = g(lon42, lat42[::-1])
for i in range(np.ma.size(uwnd_ghg, axis=2)):
h = interpolate.interp2d(lon, lat[::-1], u_ghg[:, :, i])
uwnd_ghg[:, :, i] = h(lon42, lat42[::-1])
g = interpolate.interp2d(lon, lat[::-1], v_ghg[:, :, i])
vwnd_ghg[:, :, i] = g(lon42, lat42[::-1])
w = VectorWind(uwnd, vwnd)
w_ghg = VectorWind(uwnd_ghg, vwnd_ghg)
eta = w.absolutevorticity()
eta_ghg = w_ghg.absolutevorticity()
div = w.divergence()
div_ghg = w_ghg.divergence()
uchi, vchi = w.irrotationalcomponent()
uchi_ghg, vchi_ghg = w_ghg.irrotationalcomponent()
etax, etay = w.gradient(eta)
# etax_ghg, etay_ghg = w_ghg.gradient(eta_ghg)
eta = np.transpose(eta, (2, 0, 1))
eta_ghg = np.transpose(eta_ghg, (2, 0, 1))
div = np.transpose(div, (2, 0, 1))
div_ghg = np.transpose(div_ghg, (2, 0, 1))
etax = np.transpose(etax, (2, 0, 1))
etay = np.transpose(etay, (2, 0, 1))
uchi = np.transpose(uchi, (2, 0, 1))
uchi_ghg = np.transpose(uchi_ghg, (2, 0, 1))
vchi = np.transpose(vchi, (2, 0, 1))
vchi_ghg = np.transpose(vchi_ghg, (2, 0, 1))
f_ghg = -eta * (div_ghg - div) - ((uchi_ghg - uchi) * etax +
(vchi_ghg - vchi) * etay)
meshlon, meshlat = np.meshgrid(lon42, lat42)
'''
ncsave('/home/bakerh/Downloads/vort200_control', lat42, lon42, eta.mean(axis=0)-2*np.sin(meshlat*np.pi/180)*7.2921e-5, 'vorticity')
ncsave('/home/bakerh/Downloads/forcing_ghg', lat42, lon42, f_ghg.mean(axis=0), 'forcing')
'''
month = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12])
ncsave('/home/bakerh/Downloads/vort200_control', month, lat42, lon42,
eta - 2 * np.sin(meshlat * np.pi / 180) * 7.2921e-5, 'vorticity')
ncsave('/home/bakerh/Downloads/forcing_ghg', month, lat42, lon42, f_ghg,
'forcing')
# The standard interface requires that latitude and longitude be the leading # dimensions of the input wind components, and that wind components must be # either 2D or 3D arrays. The data read in is 3D and has latitude and # longitude as the last dimensions. The bundled tools can make the process of # re-shaping the data a lot easier to manage. uwnd, uwnd_info = prep_data(uwnd, 'tyx') vwnd, vwnd_info = prep_data(vwnd, 'tyx') # It is also required that the latitude dimension is north-to-south. Again the # bundled tools make this easy. lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd) # Create a VectorWind instance to handle the computation of streamfunction and # velocity potential. w = VectorWind(uwnd, vwnd) # Compute the streamfunction and velocity potential. Also use the bundled # tools to re-shape the outputs to the 4D shape of the wind components as they # were read off files. sf, vp = w.sfvp() sf = recover_data(sf, uwnd_info) vp = recover_data(vp, uwnd_info) # Pick out the field for December and add a cyclic point (the cyclic point is # for plotting purposes). sf_dec, lons_c = add_cyclic_point(sf[11], lons) vp_dec, lons_c = add_cyclic_point(vp[11], lons) # Plot streamfunction. ax1 = plt.axes(projection=ccrs.PlateCarree(central_longitude=180))
def compute_rws(ds_u,
ds_v,
lat_coord='lat',
lon_coord='lon',
time_coord='time'):
"""
Computation of absolute vorticity, divergence, and Rossby wave source.
Outputs xarray datasets of each.
Args:
ds_u (xarray data array): Zonal (u) wind (m/s).
ds_v (xarray data array): Meridional (v) wind (m/s).
lat_coord (str): Latitude coordinate. Defaults to ``lat``.
lon_coord (str): Longitude coordinate. Defaults to ``lon``.
time_coord (str): Time coordinate. Defaults to ``time``.
Returns:
Xarray datasets for absolute vorticity, divergence, and Rossby wave source.
"""
from windspharm.standard import VectorWind
from windspharm.tools import prep_data, recover_data, order_latdim
# grab lat and lon coords
lats = ds_u.coords[lat_coord].values
lons = ds_u.coords[lon_coord].values
time = ds_u.coords[time_coord].values
_, wnd_info = prep_data(ds_u.values, 'tyx')
# reorder dims into lat, lon, time
uwnd = ds_u.transpose(lat_coord, lon_coord, time_coord).values
vwnd = ds_v.transpose(lat_coord, lon_coord, time_coord).values
# reorder lats to north-south direction
lats, uwnd, vwnd = order_latdim(lats, uwnd, vwnd)
# initialize wind vector instance
w = VectorWind(uwnd, vwnd)
# Absolute vorticity (sum of relative and planetary vorticity).
eta = w.absolutevorticity()
# Horizontal divergence.
div = w.divergence()
# Irrotational (divergent) component of the vector wind.
uchi, vchi = w.irrotationalcomponent()
# Computes the vector gradient of a scalar field on the sphere.
etax, etay = w.gradient(eta)
# Compute rossby wave source
S = -eta * div - (uchi * etax + vchi * etay)
# recover data shape
S = recover_data(S, wnd_info)
div = recover_data(div, wnd_info)
eta = recover_data(eta, wnd_info)
# assemble xarray datasets
data_rws = xr.Dataset({
'rws': (['time', 'lat', 'lon'], S),
},
coords={
'time': (['time'], time),
'lat': (['lat'], lats),
'lon': (['lon'], lons)
},
attrs={'long_name': 'Rossby wave source'})
data_div = xr.Dataset(
{
'div': (['time', 'lat', 'lon'], div),
},
coords={
'time': (['time'], time),
'lat': (['lat'], lats),
'lon': (['lon'], lons)
},
attrs={'long_name': 'Horizontal divergence (300-mb)'})
data_eta = xr.Dataset(
{
'eta': (['time', 'lat', 'lon'], eta),
},
coords={
'time': (['time'], time),
'lat': (['lat'], lats),
'lon': (['lon'], lons)
},
attrs={
'long_name':
'Absolute vorticity (sum of relative and planetary vorticity)'
})
return data_eta, data_div, data_rws