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
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
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) print np.shape(vchi)
