Source code for wrapper

"""
------------------------------------------
File name: wrapper.py
Author: Clare Huang
"""
import numpy as np
from falwa.constant import EARTH_RADIUS
from falwa import basis


[docs]def barotropic_eqlat_lwa(ylat, vort, area, dmu, n_points, planet_radius=EARTH_RADIUS): """ Compute equivalent latitude and wave activity on a barotropic sphere. Parameters ---------- ylat : numpy.array 1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat. vort : numpy.ndarray 2-d numpy array of vorticity values; dimension = (nlat, nlon). area : numpy.ndarray 2-d numpy array specifying differential areal element of each grid point; dimension = (nlat, nlon). dmu: numpy.array 1-d numpy array of latitudinal differential length element (e.g. dmu = planet_radius * cos(lat) d(lat)). Size = nlat. n_points : int, default None Analysis resolution to calculate equivalent latitude. If input as None, it will be initialized as *nlat*. planet_radius : float, default 6.378e+6 radius of spherical planet of interest consistent with input 'area'. Returns ------- qref : numpy.array 1-d numpy array of value Q(y) where latitude y is given by ylat; dimension = (nlat). lwa_result : numpy.ndarray 2-d numpy array of local wave activity values; dimension = [nlat_s x nlon] Examples -------- :doc:`Example Notebook <notebooks/Example_barotropic>` """ nlat = vort.shape[0] nlon = vort.shape[1] if n_points is None: n_points = nlat qref, fawa = basis.eqvlat_fawa( ylat, vort, area, n_points) lwa_result, _ = basis.lwa(nlon, nlat, vort, qref, dmu) return qref, lwa_result
[docs]def barotropic_input_qref_to_compute_lwa(ylat, qref, vort, area, dmu, planet_radius=EARTH_RADIUS): """ This function computes LWA based on a *prescribed* Qref instead of Qref obtained from the vorticity field on a barotropic sphere. Parameters ---------- ylat : numpy.array 1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat. qref : numpy.array 1-d numpy array of prescribed reference value of vorticity at each latitude; dimension = nlat. vort : numpy.ndarray 2-d numpy array of vorticity values; dimension = (nlat, nlon). area : numpy.ndarray 2-d numpy array specifying differential areal element of each grid point; dimension = (nlat, nlon). dmu: numpy.array 1-d numpy array of latitudinal differential length element (e.g. dmu = planet_radius * cos(lat) d(lat)). Size = nlat. planet_radius : float, default 6.378e+6 radius of spherical planet of interest consistent with input 'area'. Returns ------- lwa_result : numpy.ndarray 2-d numpy array of local wave activity values; dimension = [nlat_s x nlon] """ nlat = vort.shape[0] nlon = vort.shape[1] lwa_result = basis.lwa(nlon, nlat, vort, qref, dmu) return lwa_result
[docs]def eqvlat_hemispheric(ylat, vort, area, nlat_s=None, n_points=None, planet_radius=EARTH_RADIUS): """ Compute equivalent latitude in a hemispheric domain. Parameters ---------- ylat : numpy.array 1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat. vort : numpy.ndarray 2-d numpy array of vorticity values; dimension = (nlat, nlon). area : numpy.ndarray 2-d numpy array specifying differential areal element of each grid point; dimension = (nlat, nlon). nlat_s : int, default None The index of grid point that defines the extent of hemispheric domain from the pole. If input as None, it will be initialize as nlat // 2. n_points : int, default None Analysis resolution to calculate equivalent latitude. If input as None, it will be initialized as *nlat_s*. planet_radius : float, default 6.378e+6 radius of spherical planet of interest consistent with input 'area'. Returns ------- q_part : numpy.ndarray 1-d numpy array of value Q(y) where latitude y is given by ylat; dimension = (nlat). """ nlat = vort.shape[0] qref = np.zeros(nlat) if nlat_s is None: nlat_s = nlat // 2 if n_points is None: n_points = nlat_s # --- Southern Hemisphere --- qref1, _ = basis.eqvlat_fawa(ylat[:nlat_s], vort[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) qref[:nlat_s] = qref1 # --- Northern Hemisphere --- vort2 = -vort[::-1, :] qref2, _ = basis.eqvlat_fawa(ylat[:nlat_s], vort2[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) qref[-nlat_s:] = qref2[::-1] return qref
[docs]def eqvlat_bracket_hemispheric(ylat, vort, area, nlat_s=None, n_points=None, planet_radius=EARTH_RADIUS, vgrad=None): """ Compute equivalent latitude and <...>_Q in Nakamura and Zhu (2010) in a hemispheric domain. Parameters ---------- ylat : numpy.array 1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat. vort : numpy.ndarray 2-d numpy array of vorticity values; dimension = (nlat, nlon). area : numpy.ndarray 2-d numpy array specifying differential areal element of each grid point; dimension = (nlat, nlon). nlat_s : int, default None The index of grid point that defines the extent of hemispheric domain from the pole. If input as None, it will be initialize as nlat // 2. n_points : int, default None Analysis resolution to calculate equivalent latitude. If input as None, it will be initialized as *nlat_s*. planet_radius : float, default 6.378e+6 radius of spherical planet of interest consistent with input 'area'. vgrad: numpy.ndarray, optional 2-d numpy array of laplacian (or higher-order laplacian) values; dimension = (nlat, nlon) Returns ------- q_part : numpy.ndarray 1-d numpy array of value Q(y) where latitude y is given by ylat; dimension = (nlat). brac : numpy.ndarray or None 1-d numpy array of averaged vgrad in the square bracket. If *vgrad* = None, *brac* = None. """ nlat = vort.shape[0] qref = np.zeros(nlat) brac = np.zeros(nlat) if nlat_s is None: nlat_s = nlat // 2 if n_points is None: n_points = nlat_s # --- Southern Hemisphere --- qref1, brac1 = basis.eqvlat_fawa(ylat[:nlat_s], vort[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius, vgrad=vgrad) qref[:nlat_s] = qref1 # --- Northern Hemisphere --- vort2 = -vort[::-1, :] qref2, brac2 = basis.eqvlat_fawa(ylat[:nlat_s], vort2[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius, vgrad=vgrad) qref[-nlat_s:] = qref2[::-1] brac[:nlat_s] = brac1 brac[-nlat_s:] = brac2[::-1] return qref, brac
[docs]def qgpv_eqlat_lwa(ylat, vort, area, dmu, nlat_s=None, n_points=None, planet_radius=EARTH_RADIUS): """ Compute equivalent latitutde *qref* and local wave activity *lwa_result* based on Quasi-geostrophic potential vorticity field *vort* at a pressure level as outlined in Huang and Nakamura (2017). Parameters ---------- ylat : numpy.array 1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat. vort : numpy.ndarray 2-d numpy array of Quasi-geostrophic potential vorticity field; dimension = (nlat, nlon). area : numpy.ndarray 2-d numpy array specifying differential areal element of each grid point; dimension = (nlat, nlon). dmu: numpy.array 1-d numpy array of latitudinal differential length element (e.g. dmu = planet_radius * cos(lat) d(lat)). Size = nlat. nlat_s : int, default None The index of grid point that defines the extent of hemispheric domain from the pole. If input as None, it will be initialize as nlat // 2. n_points : int, default None Analysis resolution to calculate equivalent latitude. If input as None, it will be initialized as *nlat_s*. planet_radius : float, default 6.378e+6 radius of spherical planet of interest consistent with input 'area'. Returns ------- qref : numpy.ndarray 1-d numpy array of value Q(y) where latitude y is given by ylat; dimension = (nlat). lwa_result : numpy.ndarray 2-d numpy array of local wave activity values; dimension = [nlat_s x nlon] Examples -------- :doc:`Example Notebook <notebooks/Example_qgpv>` """ nlat = vort.shape[0] nlon = vort.shape[1] if nlat_s is None: nlat_s = nlat // 2 if n_points is None: n_points = nlat_s qref = np.zeros(nlat) lwa_result = np.zeros((nlat, nlon)) # --- Southern Hemisphere --- qref1, _ = basis.eqvlat_fawa( ylat[:nlat_s], vort[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) qref[:nlat_s] = qref1 lwa_result[:nlat_s, :], _ = basis.lwa(nlon, nlat_s, vort[:nlat_s, :], qref1, dmu[:nlat_s]) # --- Northern Hemisphere --- vort2 = -vort[::-1, :] # Added the minus sign, but gotta see if NL_North is affected qref2, _ = basis.eqvlat_fawa( ylat[:nlat_s], vort2[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) qref[-nlat_s:] = -qref2[::-1] lwa_result[-nlat_s:, :], _ = basis.lwa(nlon, nlat_s, vort[-nlat_s:, :], qref[-nlat_s:], dmu[-nlat_s:]) return qref, lwa_result
[docs]def qgpv_eqlat_lwa_ncforce(ylat, vort, ncforce, area, dmu, nlat_s=None, n_points=None, planet_radius=EARTH_RADIUS): """ Compute equivalent latitutde *qref*, local wave activity *lwa_result* and non-conservative force on wave activity *capsigma* based on Quasi- geostrophic potential vorticity field *vort* at a pressure level as outlined in Huang and Nakamura (2017). Parameters ---------- ylat : numpy.array 1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat. vort : numpy.ndarray 2-d numpy array of Quasi-geostrophic potential vorticity field; dimension = (nlat, nlon). ncforce: numpy.ndarray 2-d numpy array of non-conservative force field (i.e. theta in NZ10(a) in equation (23a) and (23b)); dimension = (nlat, nlon). area : numpy.ndarray 2-d numpy array specifying differential areal element of each grid point; dimension = (nlat, nlon). dmu: numpy.array 1-d numpy array of latitudinal differential length element (e.g. dmu = planet_radius * cos(lat) d(lat)). Size = nlat. nlat_s : int, default None The index of grid point that defines the extent of hemispheric domain from the pole. If input as None, it will be initialize as nlat // 2. n_points : int, default None Analysis resolution to calculate equivalent latitude. If input as None, it will be initialized as *nlat_s*. planet_radius : float, default 6.378e+6 radius of spherical planet of interest consistent with input 'area'. Returns ------- qref : numpy.ndarray 1-d numpy array of value Q(y) where latitude y is given by ylat; dimension = (nlat). lwa_result : numpy.ndarray 2-d numpy array of local wave activity values; dimension = (nlat, nlon). capsigma: numpy.ndarray 2-d numpy array of non-conservative force contribution value; dimension = (nlat, nlon). """ nlat = vort.shape[0] nlon = vort.shape[1] if nlat_s is None: nlat_s = nlat // 2 if n_points is None: n_points = nlat_s qref = np.zeros(nlat) lwa_result = np.zeros((nlat, nlon)) capsigma = np.zeros((nlat, nlon)) # --- Southern Hemisphere --- qref1, _ = basis.eqvlat_fawa(ylat[:nlat_s], vort[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) qref[:nlat_s] = qref1 lwa_result[:nlat_s, :], \ capsigma[:nlat_s, :] = basis.lwa(nlon, nlat_s, vort[:nlat_s, :], qref1, dmu[:nlat_s], ncforce=ncforce[:nlat_s, :]) # --- Northern Hemisphere --- vort2 = -vort[::-1, :] # Added the minus sign, but gotta see if NL_North is affected qref2, _ = basis.eqvlat_fawa(ylat[:nlat_s], vort2[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) qref[-nlat_s:] = -qref2[::-1] lwa_result[-nlat_s:, :], \ capsigma[-nlat_s:, :] = basis.lwa(nlon, nlat_s, vort[-nlat_s:, :], qref[-nlat_s:], dmu[-nlat_s:], ncforce=ncforce[-nlat_s:, :]) return qref, lwa_result, capsigma
[docs]def qgpv_eqlat_lwa_options(ylat, vort, area, dmu, nlat_s=None, n_points=None, vgrad=None, ncforce=None, planet_radius=EARTH_RADIUS): """ Compute equivalent latitutde *qref*, local wave activity *lwa_result* and non-conservative force on wave activity *capsigma* based on Quasi- geostrophic potential vorticity field *vort* at a pressure level as outlined in Huang and Nakamura (2017). Parameters ---------- ylat : numpy.array 1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat. vort : numpy.ndarray 2-d numpy array of Quasi-geostrophic potential vorticity field; dimension = (nlat, nlon). ncforce: numpy.ndarray 2-d numpy array of non-conservative force field (i.e. theta in NZ10(a) in equation (23a) and (23b)); dimension = (nlat, nlon). area : numpy.ndarray 2-d numpy array specifying differential areal element of each grid point; dimension = (nlat, nlon). dmu: numpy.array 1-d numpy array of latitudinal differential length element (e.g. dmu = planet_radius * cos(lat) d(lat)). Size = nlat. nlat_s : int, default None The index of grid point that defines the extent of hemispheric domain from the pole. If input as None, it will be initialize as nlat // 2. n_points : int, default None Analysis resolution to calculate equivalent latitude. If input as None, it will be initialized as *nlat_s*. planet_radius : float, default 6.378e+6 radius of spherical planet of interest consistent with input 'area'. Returns ------- return_dict : dictionary A dictionary that consist of the 4 computed outputs listed below. (1) qref : numpy.ndarray 1-d numpy array of value Q(y) where latitude y is given by ylat; dimension = (nlat). (2) brac_result: numpy.ndarray 1-d numpy array of <...>_Q(y) in NZ10 where latitude y is given by ylat; dimension = (nlat). (3) lwa_result : numpy.ndarray 2-d numpy array of local wave activity values; dimension = (nlat, nlon). (4) capsigma: numpy.ndarray 2-d numpy array of non-conservative force contribution value; dimension = (nlat, nlon). """ nlat = vort.shape[0] nlon = vort.shape[1] if nlat_s is None: nlat_s = nlat // 2 if n_points is None: n_points = nlat_s qref = np.zeros(nlat) lwa_result = np.zeros((nlat, nlon)) if ncforce is not None: capsigma = np.zeros((nlat, nlon)) if vgrad is not None: brac_result = np.zeros(nlat) # --- Southern Hemisphere --- if vgrad is None: qref1, brac = basis.eqvlat_fawa(ylat[:nlat_s], vort[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) else: qref1, brac = basis.eqvlat_fawa(ylat[:nlat_s], vort[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius, vgrad=vgrad[:nlat_s, :]) qref[:nlat_s] = qref1 if vgrad is not None: brac_result[:nlat_s] = brac if ncforce is not None: lwa_result[:nlat_s, :], capsigma[:nlat_s, :] = \ basis.lwa(nlon, nlat_s, vort[:nlat_s, :], qref1, dmu[:nlat_s], ncforce=ncforce[:nlat_s, :]) else: lwa_result[:nlat_s, :], _ = basis.lwa(nlon, nlat_s, vort[:nlat_s, :], qref1, dmu[:nlat_s]) # --- Northern Hemisphere --- vort2 = -vort[::-1, :] # Added the minus sign, but gotta see if NL_North is affected if vgrad is None: qref2, brac = basis.eqvlat_fawa(ylat[:nlat_s], vort2[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) else: vgrad2 = -vgrad[::-1, :] # Is this correct? qref2, brac = basis.eqvlat_fawa(ylat[:nlat_s], vort2[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius, vgrad=vgrad2[:nlat_s, :]) qref[-nlat_s:] = -qref2[::-1] if vgrad is not None: brac_result[-nlat_s:] = -brac[::-1] if ncforce is not None: lwa_result[-nlat_s:, :], \ capsigma[-nlat_s:, :] = basis.lwa(nlon, nlat_s, vort[-nlat_s:, :], qref[-nlat_s:], dmu[-nlat_s:], ncforce=ncforce[-nlat_s:, :]) else: lwa_result[-nlat_s:, :], _ = basis.lwa(nlon, nlat_s, vort[-nlat_s:, :], qref[-nlat_s:], dmu[-nlat_s:]) # Return things as dictionary return_dict = dict() return_dict['qref'] = qref return_dict['lwa_result'] = lwa_result if vgrad is not None: return_dict['brac_result'] = brac_result if ncforce is not None: return_dict['capsigma'] = capsigma return return_dict
[docs]def qgpv_input_qref_to_compute_lwa(ylat, qref, vort, area, dmu, nlat_s=None, planet_radius=EARTH_RADIUS): """ Compute equivalent latitutde *qref* and local wave activity *lwa_result* based on Quasi-geostrophic potential vorticity field *vort* at a pressure level as outlined in Huang and Nakamura (2017). This function computes lwa based on a prescribed *qref* instead of *qref* obtained from the QGPV field. Parameters ---------- ylat : numpy.array 1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat. qref : numpy.ndarray 1-d numpy array of value Q(y) where latitude y is given by ylat; dimension = (nlat). vort : numpy.ndarray 2-d numpy array of Quasi-geostrophic potential vorticity field; dimension = (nlat, nlon). area : numpy.ndarray 2-d numpy array specifying differential areal element of each grid point; dimension = (nlat, nlon). dmu: numpy.array 1-d numpy array of latitudinal differential length element (e.g. dmu = planet_radius * cos(lat) d(lat)). Size = nlat. nlat_s : int, default None The index of grid point that defines the extent of hemispheric domain from the pole. If input as None, it will be initialize as nlat // 2. planet_radius : float, default 6.378e+6 radius of spherical planet of interest consistent with input 'area'. Returns ------- lwa_result : numpy.ndarray 2-d numpy array of local wave activity values; dimension = (nlat, nlon). """ nlat = vort.shape[0] nlon = vort.shape[1] if nlat_s is None: nlat_s = nlat // 2 lwa_result = np.zeros((nlat, nlon)) # --- Southern Hemisphere --- lwa_result[:nlat_s, :], dummy = basis.lwa(nlon, nlat_s, vort[:nlat_s, :], qref[:nlat_s], dmu[:nlat_s]) # --- Northern Hemisphere --- lwa_result[-nlat_s:, :], dummy = basis.lwa(nlon, nlat_s, vort[-nlat_s:, :], qref[-nlat_s:], dmu[-nlat_s:]) return lwa_result
[docs]def theta_lwa(ylat, theta, area, dmu, nlat_s=None, n_points=None, planet_radius=EARTH_RADIUS): """ Compute the surface wave activity *B* based on surface potential temperature. See Nakamura and Solomon (2010a) for details. Parameters ---------- ylat : numpy.array 1-d numpy array of latitude (in degree) with equal spacing in ascending order; dimension = nlat. theta : numpy.ndarray 2-d numpy array of surface potential temperature field; dimension = (nlat, nlon). area : numpy.ndarray 2-d numpy array specifying differential areal element of each grid point; dimension = (nlat, nlon). dmu: numpy.array 1-d numpy array of latitudinal differential length element (e.g. dmu = planet_radius * cos(lat) d(lat)). Size = nlat. nlat_s : int, default None The index of grid point that defines the extent of hemispheric domain from the pole. If input as None, it will be initialize as nlat // 2. planet_radius : float, default 6.378e+6 radius of spherical planet of interest consistent with input 'area'. Returns ------- qref : numpy.array 1-d numpy array of value reference potential temperature *Theta(y)* approximated by box counting method, where latitude y is given by ylat; dimension = (nlat). lwa_result : numpy.ndarray 2-d numpy array of local surface wave activity values; dimension = (nlat, nlon). """ nlat = theta.shape[0] nlon = theta.shape[1] if nlat_s is None: nlat_s = nlat // 2 if n_points is None: n_points = nlat_s qref = np.zeros(nlat) lwa_result = np.zeros((nlat, nlon)) # --- southern Hemisphere --- qref[:nlat_s], brac = basis.eqvlat_fawa(ylat[:nlat_s], theta[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) lwa_result[:nlat_s, :], dummy = basis.lwa(nlon, nlat_s, theta[:nlat_s, :], qref[:nlat_s], dmu[:nlat_s]) # --- northern Hemisphere --- theta2 = theta[::-1, :] # Added the minus sign, but gotta see if NL_north is affected qref2, brac = basis.eqvlat_fawa(ylat[:nlat_s], theta2[:nlat_s, :], area[:nlat_s, :], n_points, planet_radius=planet_radius) qref[-nlat_s:] = qref2[::-1] lwa_result[-nlat_s:, :], dummy = basis.lwa(nlon, nlat_s, theta[-nlat_s:, :], qref[-nlat_s:], dmu[-nlat_s:]) return qref, lwa_result