def waqlocal(lon, lat, q, omega=None, sigma=None, flip=True, skip=10):
"""
Return the local finite-amplitude wave activity. See
:cite:`2016:huang` for details.
Parameters
----------
%(params_eqlat)s
%(params_waq)s
References
----------
.. bibliography:: ../bibs/waqlocal.bib
Todo
----
Support using `omega`.
"""
# Graticule considerations
has_omega = omega is not None
has_sigma = sigma is not None
if flip:
lat, q = -np.flipud(lat), -np.flip(q, axis=1)
if has_omega:
omega = -np.flip(omega, axis=1)
if sigma is not None:
sigma = np.flipd(sigma, axis=1)
grid = _LongitudeLatitude(lon, lat)
phib = grid.phib
integral = const.a * phib[None, :]
# Flatten (eqlat can do this, but not necessary here)
arrays = [q]
if has_omega:
arrays.append(omega)
if has_sigma:
arrays.append(sigma)
# Flatten (eqlat can do this, but not necessary here)
with quack._ArrayContext(*arrays, nflat_right=(q.ndim - 2)) as context:
# Get flattened data
if has_omega and has_sigma:
q, omega, sigma = context.data
elif has_omega:
q, omega = context.data
elif has_sigma:
q, sigma = context.data
else:
q = context.data
# Get equivalent latiitudes
bands, q_bands = eqlat(lon, lat, q, sigma=sigma, skip=skip)
L = lon.size
M = lat.size
N = bands.shape[1] # number of equivalent latitudes
K = q.shape[2] # number of extra dimensions
# Get local wave activity measure, as simple line integrals
waq = np.empty((L, M, K))
percent = 0
for k in range(K):
if (k / K) > (0.01 * percent):
print('%d%% finished' % (100 * k / K, ))
percent = percent + 10
# Loop through each contour
waq_k = np.empty((L, N))
for n in range(N):
# Setup, large areas
band = bands[0, n, k] * np.pi / 180
if np.isnan(band): # happens if contours intersect at edge
waq_k[:, n] = np.nan
else:
# Get high anomalies at low latitude (below top graticule) and
# low anomalies at high latitude (above bottom graticule)
anom = q[:, :, k] - q_bands[0, n, k]
f_pos = (anom >= 0) & (phib[None, 1:] < band)
f_neg = (anom < 0) & (phib[None, :-1] >= band)
# See if band is sandwiched between latitudes
# want scalar id (might be zero)
mid, = np.where((phib[:-1] <= band) & (phib[1:] > band))
if mid.size > 0:
# Find longitudes where positive
# Perform partial integrals, positive and negative
f_pos_mid = anom[:, mid] >= 0
p_int = const.a * (band - phib[mid])
m_int = const.a * (phib[mid + 1] - band)
for l in range(L):
# Get individual integrals
integral_pos = (anom[l, f_pos[l, :]] *
integral[:, f_pos[l, :]]).sum()
integral_neg = -( # *minus* a *negative*
anom[l, f_neg[l, :]] *
integral[:, f_neg[l, :]]).sum()
# Get extra bits
if mid.size > 0:
if f_pos_mid[l]:
integral_extra = anom[l, mid] * p_int
else:
integral_extra = -anom[l, mid] * m_int
else:
integral_extra = 0
# Put it all together
waq_k[l,
n] = integral_pos + integral_neg + integral_extra
# Interpolate to actual latitudes
for l in range(L):
waq[l, :, k] = np.interp(lat, bands[0, :, k], waq_k[l, :])
# Replace context data
context.replace_data(waq)
# Return
waq = context.data
if flip:
waq = np.flip(waq, axis=1)
return waq
def waq(lon, lat, q, sigma=None, omega=None, flip=True, skip=10):
"""
Return the finite-amplitude wave activity. See
:cite:`2010:nakamura` for details.
Parameters
----------
%(params_eqlat)s
%(params_waq)s
References
----------
.. bibliography:: ../bibs/waq.bib
"""
# Graticule considerations
has_omega = omega is not None
has_sigma = sigma is not None
if flip:
lat, q = -np.flipud(lat), -np.flip(q, axis=1)
if has_omega:
omega = -np.flip(omega, axis=1)
if has_sigma:
sigma = np.flipd(sigma, axis=1)
grid = _LongitudeLatitude(lon, lat)
areas, dphi, phib = grid.areas, grid.dphi, grid.phib
# Flatten (eqlat can do this, but not necessary here)
arrays = [q]
if has_omega:
arrays.append(omega)
if has_sigma:
arrays.append(sigma)
with quack._ArrayContext(*arrays, nflat_right=(q.ndim - 2)) as context:
# Get flattened data
if has_omega and has_sigma:
q, omega, sigma = context.data
elif has_omega:
q, omega = context.data
elif has_sigma:
q, sigma = context.data
else:
q = context.data
# Get equivalent latiitudes
bands, q_bands = eqlat(lon, lat, q, sigma=sigma, skip=skip)
M = lat.size # number of latitudes onto which we interpolate
N = bands.shape[1] # number of equivalent latitudes
K = q.shape[2] # number of extra dimensions
# Get activity
waq = np.empty((1, M, K))
percent = 0
for k in range(K):
# Status message
if (k / K) > (0.01 * percent):
print('%d%% finished' % (percent, ))
percent = percent + 10
# Loop through each contour
# i, bandki in enumerate(bandk[0,:,k]):
# enumerate(zip(bandk, q_bandk)):
waq_k = np.empty(N)
for n in range(N):
# First, main blocks
band = bands[0, n, k] * np.pi / 180
if np.isnan(band):
waq_k[n] = np.nan
continue
# Work with anomalies? Should be identical, since
# positive/negative regions have same area by construction.
# anom = q[:,:,k] - q_bands[0,n,k]
# f_pos = (anom >= 0) & (phib[None,1:] < band)
# f_neg = (anom < 0) & (phib[None,:-1] >= band)
# integral = (anom[f_pos]*areas[f_pos]).sum() -
# (anom[f_neg]*areas[f_neg]).sum() # minus a negative
# Get the integrand
qk, Qk = q[:, :, k], q_bands[0, n, k]
if has_omega:
qint = q[:, :, k] # the thing being integrated
else:
qint = omega[:, :, k]
# Get high anomalies at low latitude (below top graticule) and
# low anomalies at high latitude (above bottom graticule)
f_pos = (qk >= Qk) & (phib[None, 1:] < band)
f_neg = (qk < Qk) & (phib[None, :-1] >= band)
integral = (
(qint[f_pos] * areas[f_pos]).sum() -
(qint[f_neg] * areas[f_neg]).sum() # minus a negative
)
# Account for tiny pieces along equivalent latitude cells
mid, = np.where((phib[:-1] <= band) & (phib[1:] > band))
integral_extra = 0
if mid:
# Find longitudes where positive
f_pos_mid = qk[:, mid] >= Qk
# Positive high latitude and negative, low latitude
p_dphi = np.cos(
(band + phib[mid]) / 2) * (band - phib[mid])
m_dphi = np.cos(
(band + phib[mid + 1]) / 2) * (phib[mid + 1] - band)
# Extra pieces due to discrete grid
integral_extra = (
qint[f_pos_mid, mid].sum() *
(areas[mid] * m_dphi / dphi[mid]) -
qint[~f_pos_mid, mid].sum() *
(areas[mid] * p_dphi / dphi[mid]) # noqa: E501
)
# Put it all together
waq_k[n] = ((integral + integral_extra) /
(2 * np.pi * const.a * np.cos(band)))
# Interpolate
nanfilt = np.isnan(waq_k)
if sum(~nanfilt) == 0:
warnings._warn_climopy(
f'Warning: No valid waqs calculated for k {k}.')
waq[0, :, k] = np.nan
else:
waq[0, :, k] = np.interp(lat, bands[0, ~nanfilt, k],
waq_k[~nanfilt])
# Reapply data
context.replace_data(waq)
# Return
waq = context.waq
if flip:
waq = np.flip(waq, axis=1)
return waq