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
