def div(u, v, w):
r"""Calculate divergence of the input vector
:math:`\nabla \cdot \mathbf{x}`
Args:
u,v,w (iris.cube.Cube): i,j,k components of the input vector
Returns:
iris.cube.Cube: The divergence
"""
# Calculate individual components
du_dx = polar_horizontal(u, 'x')
du_dx = interpolate.remap_3d(du_dx, w)
dv_dy = polar_horizontal(v, 'y')
dv_dy = interpolate.remap_3d(dv_dy, w)
dw_dz = multidim(w, 'altitude', 'z')
dw_dz = interpolate.remap_3d(dw_dz, w)
# Sum to divergence
divergence = du_dx.data + dv_dy.data + dw_dz.data
# Copy to a cube
divergence = dw_dz.copy(data=divergence)
return divergence
def potential_vorticity(u, v, w, theta, rho):
r"""Calculate PV
.. math::
q = \frac{1}{\rho} (\nabla \times \mathbf{u} + 2 \boldsymbol{\Omega})
\cdot \nabla \theta
Args:
u (iris.cube.Cube): Zonal velocity
v (iris.cube.Cube): Meridional velocity
w (iris.cube.Cube): Vertical velocity
theta (iris.cube.Cube): Potential temperature
rho (iris.cube.Cube): Density
Returns:
iris.cube.Cube: PV in PVU
"""
# Relative Vorticity
xterm, yterm, zterm = vorticity(u, v, w)
# Absolute vorticity
lat = grid.true_coords(theta)[1]
f = 2 * constants.omega.data * np.sin(lat * np.pi / 180)
zterm.data = zterm.data + f
# Grad(theta)
dtheta_dx = calculus.polar_horizontal(theta, 'x')
dtheta_dx = interpolate.remap_3d(dtheta_dx, theta)
dtheta_dy = calculus.polar_horizontal(theta, 'y')
dtheta_dy = interpolate.remap_3d(dtheta_dy, theta)
z = grid.make_cube(theta, 'altitude')
dtheta_dz = calculus.multidim(theta, z, 'z')
dtheta_dz = interpolate.remap_3d(dtheta_dz, theta)
# PV components
PV_x = xterm * dtheta_dx
PV_y = yterm * dtheta_dy
PV_z = zterm * dtheta_dz
# Full PV
rho_theta = interpolate.remap_3d(rho, theta)
epv = (PV_x + PV_y + PV_z) / rho_theta
epv.rename('ertel_potential_vorticity')
epv.convert_units('PVU')
return epv
def grad(cube):
r"""Calculate gradient of the input field
:math:`\nabla x`
Args:
cube (iris.cube.Cube): Input field
Returns:
tuple (iris.cube.Cube, iris.cube.Cube, iris.cube.Cube):
Three cubes of the different components of the vector gradient
"""
d_dx = polar_horizontal(cube, 'x')
d_dx = interpolate.remap_3d(d_dx, cube)
d_dy = polar_horizontal(cube, 'y')
d_dy = interpolate.remap_3d(d_dy, cube)
d_dz = multidim(cube, 'altitude', 'z')
d_dz = interpolate.remap_3d(d_dz, cube)
return d_dx, d_dy, d_dz
def redo_cubes(cubes, basis_cube, stash_maps=[], slices=None, time=None):
# Coordinates to copy to analyses
z = grid.extract_dim_coord(basis_cube, 'z')
# Define attributes of custom variables by stash mapping
for stash_map in stash_maps:
stash_map.remap_cubelist(cubes)
newcubelist = CubeList()
for cube in cubes:
print(cube)
newcube = cube.copy()
# Remove unneccessary time coordinates
for coord in ['forecast_period', 'forecast_reference_time']:
try:
newcube.remove_coord(coord)
except iris.exceptions.CoordinateNotFoundError:
pass
if newcube.ndim == 3:
# Use the hybrid height coordinate as the main vertical coordinate
try:
newcube.remove_coord('model_level_number')
except iris.exceptions.CoordinateNotFoundError:
pass
newcube.coord('level_height').rename(z.name())
iris.util.promote_aux_coord_to_dim_coord(newcube, z.name())
# Remap the cubes to theta points
newcube = interpolate.remap_3d(newcube, basis_cube, z.name())
else:
# Regrid in the horizontal
newcube = newcube.regrid(basis_cube, Linear())
# Convert the main time coordinate
if time is not None:
newcube.coord('time').convert_units(time)
newcubelist.append(newcube)
return newcubelist
def curl(u, v, w):
r"""Calculate curl of the input vector
:math:`\nabla \times \mathbf{x}`
Args:
u,v,w (iris.cube.Cube): i,j,k components of the input vector
Returns:
tuple (iris.cube.Cube, iris.cube.Cube, iris.cube.Cube):
Three cubes of the different components of the vector curl
"""
# Calculate individual gradients
dw_dx = polar_horizontal(w, 'x')
dw_dx = interpolate.remap_3d(dw_dx, w)
dw_dy = polar_horizontal(w, 'y')
dw_dy = interpolate.remap_3d(dw_dy, w)
du_dz = multidim(u, 'altitude', 'z')
du_dz = interpolate.remap_3d(du_dz, w)
dv_dz = multidim(v, 'altitude', 'z')
dv_dz = interpolate.remap_3d(dv_dz, w)
du_dy = polar_horizontal(u, 'y')
du_dy = interpolate.remap_3d(du_dy, w)
dv_dx = polar_horizontal(v, 'x')
dv_dx = interpolate.remap_3d(dv_dx, w)
# Calculate the components of vorticity
xi_i = dw_dy.data - dv_dz.data
xi_i = dw_dx.copy(data=xi_i)
xi_j = du_dz.data - dw_dx.data
xi_j = dw_dx.copy(data=xi_j)
xi_k = dv_dx.data - du_dy.data
xi_k = dw_dx.copy(data=xi_k)
return xi_i, xi_j, xi_k