def relvort(u, v, lats, lons):
"""
Compute the relative vertical vorticity of the wind
Args:
u ([type]): 2D arrays of u wind components, in meters per second
v ([type]): 2D arrays of v wind components, in meters per second
Returns:
Returns in units of per second
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(u)
#Compute the gradient of the wind
dudy = calculate.compute_gradient(u, lats, lons)[1]
dvdx = calculate.compute_gradient(v, lats, lons)[0]
#Compute relative vorticity (dv/dx - du/dy)
vort = np.subtract(dvdx, dudy)
#Account for southern hemisphere
tlons, tlats = np.meshgrid(lons, lats)
vort[tlats < 0] = vort[tlats < 0] * -1
#Convert back to xarray dataset, if initially passed as one
if use_xarray == True:
vort = xr.DataArray(vort, coords=[lats, lons], dims=['lat', 'lon'])
return vort
def grid_area_average_degree(prod, deg, lats, lons):
"""
Area averaging a lat/lon grid by a specified radius in degrees (not kilometers)
Args:
prod ([type]): 2D variable to be area-averaged
deg ([type]): Degree radius to smooth over (e.g., 2 for 2 degrees)
lats ([type]): [description]
lons ([type]): [description]
Returns:
[type]: [description]
"""
#Check if input product is an xarray dataarray
use_xarray = arr.check_xarray(prod)
#Determine radius in gridpoint numbers
res = abs(lats[1] - lats[0])
#Perform area-averaging
radius = int(float(deg) / res)
kernel = np.zeros((2 * radius + 1, 2 * radius + 1))
y1, x1 = np.ogrid[-radius:radius + 1, -radius:radius + 1]
mask = x1**2 + y1**2 <= radius**2
kernel[mask] = 1
prod = ndimage.filters.generic_filter(prod, np.average, footprint=kernel)
#Convert back to xarray dataarray, if specified
if use_xarray == 1:
prod = xr.DataArray(prod, coords=[lats, lons], dims=['lat', 'lon'])
#Return product
return prod
def advection(var, u, v, lats, lons):
"""
Compute the magnitude of horizontal advection of a scalar quantity by the wind
Args:
var ([type]): 2D scalar field (e.g. temperature)
u ([type]): 2D arrays of u & v wind components, in meters per second
v ([type]): 2D arrays of u & v wind components, in meters per second
Returns:
Returns in units of (scalar unit) per second
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(var)
#Compute the gradient of the variable
ddx, ddy = calculate.compute_gradient(var, lats, lons)
#Compute advection (-v dot grad var)
#adv = -1 * ((ddx*u) + (ddy*v))
adv = np.add(np.multiply(ddx, u), np.multiply(ddy, v))
adv = np.multiply(-1.0, adv)
#Convert back to xarray dataset, if initially passed as one
if use_xarray == True:
adv = xr.DataArray(adv, coords=[lats, lons], dims=['lat', 'lon'])
return adv
def divergence(u, v, lats, lons):
"""
Compute the horizontal divergence of a vector
Args:
var ([type]): 2D scalar field (e.g. temperature)
u ([type]): 2D arrays of u & v wind components, in meters per second
v ([type]): 2D arrays of u & v wind components, in meters per second
Returns:
Returns in units of per second
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(u)
#Compute the gradient of the wind
dudx = calculate.compute_gradient(u, lats, lons)[0]
dvdy = calculate.compute_gradient(v, lats, lons)[1]
#div = dudx + dvdy #dv/dx - du/dy
div = np.add(dudx, dvdy)
#Convert back to xarray dataset, if initially passed as one
if use_xarray == True:
div = xr.DataArray(div, coords=[lats, lons], dims=['lat', 'lon'])
return div
def ageo(hght, u, v, lats, lons):
"""
Compute the u and v components of the ageostrophic wind
Args:
hght ([type]): 2D scalar geopotential height field (m)
u ([type]): 2D arrays of u components of wind (m/s)
v ([type]): 2D arrays of v components of wind (m/s)
Returns:
Returns in units of meters per second
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(hght)
#Compute the geostrophic wind
ug, vg = geo(hght, lats, lons)
#Compute the ageostrophic wind
ua = u - ug
va = v - vg
#Convert back to xarray dataset, if initially passed as one
if use_xarray == True:
ua = xr.DataArray(ua, coords=[lats, lons], dims=['lat', 'lon'])
va = xr.DataArray(va, coords=[lats, lons], dims=['lat', 'lon'])
return ua, va
def geo(hght, lats, lons):
"""
Compute the u and v components of the geostrophic wind
Args:
hght ([type]): 2D scalar geopotential height field (m)
Returns:
Returns in units of meters per second
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(hght)
#Compute geopotential height gradient on pressure surface
dzdx, dzdy = calculate.compute_gradient(hght, lats, lons)
#2D array of Coriolis parameter for each lat/lon
cor = coriolis(lats, lons)
#Compute the geostrophic wind
ug = (-1.0 * dzdy * g) / cor
vg = (dzdx * g) / cor
#Convert back to xarray dataset, if initially passed as one
if use_xarray == True:
ug = xr.DataArray(ug, coords=[lats, lons], dims=['lat', 'lon'])
vg = xr.DataArray(vg, coords=[lats, lons], dims=['lat', 'lon'])
return ug, vg
def absvort(u, v, lats, lons):
"""
Compute the absolute vertical vorticity of the wind
Args:
u ([type]): 2D arrays of u wind components, in meters per second
v ([type]): 2D arrays of v wind components, in meters per second
Returns:
Returns in units of per second
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(u)
#Compute relative vorticity
vort = relvort(u, v, lats, lons)
#Compute the Coriolis parameter (after converting lat to radians)
cor2d = coriolis(lats, lons)
#Compute absolute vorticity (relative + coriolis parameter)
vort = np.add(vort, cor2d)
#Convert back to xarray dataset, if initially passed as one
if use_xarray == True:
vort = xr.DataArray(vort, coords=[lats, lons], dims=['lat', 'lon'])
return vort
def qvect(temp, hght, lev, lats, lons, smooth, static_stability=1):
"""
Compute the u and v components of the Q-vector
Args:
temp ([type]): 2D scalar temperature field (K)
hght ([type]): 2D scalar geopotential height field (m)
lev ([type]): Pressure level (hPa)
lats ([type]): 1D arrays of lat
lons ([type]): 1D arrays of lon
smooth ([type]): integer representing sigma level of smoothing
static_stability (int, optional): assumed to be 1, unless provided. Defaults to 1.
Returns:
Returns in units of meters per second
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(temp)
#Smooth data
hght = ndimage.gaussian_filter(hght, sigma=smooth, order=0)
temp = ndimage.gaussian_filter(temp, sigma=smooth, order=0)
#Convert pressure to Pa
levPa = lev * 100.0
#Compute the geostrophic wind
ug, vg = geo(hght, lats, lons)
#Compute the constant out front
const = (-1.0 * Rd) / (levPa * static_stability)
#Compute gradient quantities
dtdx, dtdy = calculate.compute_gradient(temp, lats, lons)
dudx, dudy = calculate.compute_gradient(ug, lats, lons)
dvdx, dvdy = calculate.compute_gradient(vg, lats, lons)
#Compute u,v components of Q-vector
Qu = const * ((dudx * dtdx) + (dvdx * dtdy))
Qv = const * ((dudy * dtdx) + (dvdy * dtdy))
#Convert back to xarray dataset, if initially passed as one
if use_xarray == True:
Qu = xr.DataArray(Qu, coords=[lats, lons], dims=['lat', 'lon'])
Qv = xr.DataArray(Qv, coords=[lats, lons], dims=['lat', 'lon'])
return Qu, Qv
def isentropic_transform(temp, u, v, lev, lats, lons, levs, tomask=1):
"""
Transform a variable to isentropic coordinates, specifying a single isentropic level
https://github.com/tomerburg/metlib/blob/master/diagnostics/met_functions.py
Args:
temp ([type]): 3D temperature array (K)
u ([type]): 3D wind array (u)
v ([type]): 3D wind array (v)
lev ([type]): Desired isentropic level (K, scalar)
lats ([type]): 1D lat array
lons ([type]): 1D lon array
levs ([type]): 1D pressure array
tomask (int, optional): mask array values where the desired isentropic surface is below the
# ground. Yes=1, No=0. Default is yes (1). Defaults to 1.
Returns:
Returns a python list with the following quantities:
[0] = 2D pressure array
[1] = u-wind
[2] = v-wind
[3] = d(theta)/dp
[4] = 2D array corresponding to the first k-index of where the
theta threshold is exceeded.
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(temp)
#If using xarray, convert to numpy
if use_xarray == 1:
temp = temp.values
u = u.values
v = v.values
#Subset data values to below 100 hPa
tlev = float(lev)
#Arrange a 3D pressure array of theta
vtheta = np.copy(temp) * 0.0
nvert, nlat, nlon = np.shape(temp)
for k in range(0, nvert):
vtheta[k] = theta(temp[k], levs[k])
#Arrange 2D arrays of other values
tpres = 0
tu = 0
tv = 0
#Eliminate any NaNs to avoid issues
temp = np.nan_to_num(temp)
u = np.nan_to_num(u)
v = np.nan_to_num(v)
vtheta = np.nan_to_num(vtheta)
#==================================================================
#Step 0: Get 3d array of pressure
levs3d = np.copy(temp) * 0.0
nvert, nlat, nlon = np.shape(temp)
for k in range(0, nvert):
levs3d[k] = (vtheta[0] * 0.0) + levs[k]
#------------------------------------------------------------------
#Step 1: find first instances bottom-up of theta exceeding threshold
#Check where the theta threshold is exceeded in the 3D array
check_thres = np.where(vtheta >= tlev)
check_ax1 = check_thres[0]
check_ax2 = check_thres[1]
check_ax3 = check_thres[2]
#This is a 2D array corresponding to the first k-index of where the
#theta threshold is exceeded.
thres_pos = np.copy(vtheta[0]) * 0.0
#Loop through all such positions and only record first instances
for i in range(0, len(check_ax1)):
pres_level = check_ax1[i]
jval = check_ax2[i]
ival = check_ax3[i]
if thres_pos[jval][ival] == 0: thres_pos[jval][ival] = pres_level
#------------------------------------------------------------------
#Step 2: get the theta values corresponding to this axis
#Convert the position of the theta threshold values to something readable
thres_pos = thres_pos.astype('int64')
thres_last = thres_pos - 1
thres_last[thres_last < 0] = 0
#replace NaNs, if any
thres_pos = np.nan_to_num(thres_pos)
thres_last = np.nan_to_num(thres_last)
vtheta = np.nan_to_num(vtheta)
#Get theta values where it's first exceeded and 1 vertical level below it
ktheta = np.ndarray.choose(thres_pos, vtheta)
ltheta = np.ndarray.choose(thres_last, vtheta)
#Get the difference in theta between levels
diffu = np.abs(np.subtract(tlev, ktheta))
diffl = np.abs(np.subtract(tlev, ltheta))
#Percentage from the lower level to the upper one
perc = np.divide(diffl, np.add(diffu, diffl))
#------------------------------------------------------------------
#Step 3: find pressure at this level
valu = np.ndarray.choose(thres_pos, levs3d)
vall = np.ndarray.choose(thres_last, levs3d)
#Adjustment factor
fac = np.multiply(np.subtract(vall, valu), perc)
#New pressure
tpres = np.subtract(vall, fac)
#------------------------------------------------------------------
#Step 3a: get d(theta)/dp array
#d(theta)/dp = ktheta-ltheta / valu-vall
dthetadp = np.divide(np.subtract(ktheta, ltheta), np.subtract(valu, vall))
#Convert to units of K/Pa
dthetadp = np.divide(dthetadp, 100)
#------------------------------------------------------------------
#Step 4: find wind at this level
uu = np.ndarray.choose(thres_pos, u)
ul = np.ndarray.choose(thres_last, u)
vu = np.ndarray.choose(thres_pos, v)
vl = np.ndarray.choose(thres_last, v)
fac = np.multiply(np.subtract(ul, uu), perc)
tu = np.subtract(ul, fac)
fac = np.multiply(np.subtract(vl, vu), perc)
tv = np.subtract(vl, fac)
#==================================================================
# ALL RESUMES HERE
#==================================================================
if tomask == 1:
pres = np.ma.masked_where(thres_pos <= 1.0, tpres)
u = np.ma.masked_where(thres_pos <= 1.0, tu)
v = np.ma.masked_where(thres_pos <= 1.0, tv)
#Convert back to xarray, if specified initially
if use_xarray == 1:
pres = xr.DataArray(pres, coords=[lats, lons], dims=['lat', 'lon'])
u = xr.DataArray(u, coords=[lats, lons], dims=['lat', 'lon'])
v = xr.DataArray(v, coords=[lats, lons], dims=['lat', 'lon'])
return pres, u, v, dthetadp, thres_pos
def integrated_vapor(temp, rh, pressfc, levs, lats, lons):
"""
Compute integrated vapor over a certain pressure layer, assuming the pressure
interval is constant.
Args:
temp ([type]): 3D array (lev,lat,lon) of temperature (K)
rh ([type]): 3D array (lev,lat,lon) of relative humidity (in %)
pressfc ([type]): Surface pressure (hPa)
levs ([type]): 1D array of pressure levels (hPa)
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(temp)
#If using xarray, convert to numpy
if use_xarray == 1:
try:
temp = temp.values
except:
pass
try:
rh = rh.values
except:
pass
#Convert pressure to hPa
levs = levs * 100.0
#determine vertical dz in Pa, assuming levs array is uniform
vint = (levs[1] - levs[0])
#pres,lons0,lats0 = np.meshgrid(rh.lev,lons,lats)
nvert, nlat, nlon = np.shape(rh)
pres = np.copy(rh) * 0.0
#Arrange a 3D pressure array
for k in range(0, nvert):
pres[k] += levs[k]
#Mask by surface pressure
tmp = rh[k]
tmp[pressfc < levs[k]] = 0.0
rh[k] = tmp
#saturated vapor pressure in Pa
es = vapor_pressure(temp) * 100.0
#get e from RH (in decimals) in Pa
e = (rh / 100.0) * es
#Approximate specific humidity q ~ w
w = 0.622 * (e / pres) #used to be 0.622
q = w
#Compute integrated vapor
iv = np.trapz(q, axis=0, dx=vint) / -9.8
#Return ut, vt as xarray DataArrays
if use_xarray == 1:
iv = xr.DataArray(iv, coords=[lats, lons], dims=['lat', 'lon'])
return iv
def ivt(temp, rh, levs, u, v, lats, lons):
"""
Compute integrated vapor transport, assuming the pressure interval is constant
Args:
temp ([type]): 3D array (lev,lat,lon) of temperature (K)
rh ([type]): 3D array (lev,lat,lon) of relative humidity (in %)
levs ([type]): 1D array of pressure levels (hPa)
u ([type]): 3D array u-wind (m/s)
v ([type]): 3D array v-wind (m/s)
"""
#Check if input is an xarray dataarray
use_xarray = arr.check_xarray(temp)
#If using xarray, convert to numpy arrays
if use_xarray == 1:
try:
temp = temp.values
except:
pass
try:
rh = rh.values
except:
pass
try:
u = u.values
except:
pass
try:
v = v.values
except:
pass
#Get list of pressure levels in hPa, convert to Pa
levs = levs * 100.0 #convert pressure to Pa
#determine vertical dz in Pa, assuming levs array is uniform
vint = (levs[1] - levs[0])
nvert, nlat, nlon = np.shape(rh)
pres = np.copy(rh) * 0.0
#Arrange a 3D pressure array
for k in range(0, nvert):
pres[k] += levs[k]
#saturated vapor pressure in Pa
es = vapor_pressure(temp) * 100.0
#get e from RH (in decimals) in Pa
e = (rh / 100.0) * es
#Approximate specific humidity q ~ w
q = 0.622 * (e / pres)
#Compute u and v components of IVT vector
ut = np.trapz(u * q, axis=0, dx=vint) / -9.8
vt = np.trapz(v * q, axis=0, dx=vint) / -9.8
#Compute magnitude of IVT vector
ivt = np.sqrt(np.add(np.square(ut), np.square(vt)))
#Convert back to xarray dataset, if initially passed as one
if use_xarray == True:
ut = xr.DataArray(ut, coords=[lats, lons], dims=['lat', 'lon'])
vt = xr.DataArray(vt, coords=[lats, lons], dims=['lat', 'lon'])
ivt = xr.DataArray(ivt, coords=[lats, lons], dims=['lat', 'lon'])
return ut, vt, ivt