def main(cubes, **kwargs):
"""Calculate and plot the distance from the tropopause to cloud
"""
pv = convert.calc('ertel_potential_vorticity', cubes)
cl = convert.calc('mass_fraction_of_cloud', cubes)
# Get altitude as a cube
z = grid.make_cube(pv, 'altitude')
# Add pv and cloud as coordinates to the altitude cube
pv = grid.make_coord(pv)
z.add_aux_coord(pv, [0, 1, 2])
cl = grid.make_coord(cl)
z.add_aux_coord(cl, [0, 1, 2])
# Calculate the height of the cloud top
z_cloud = interpolate.to_level(z, mass_fraction_of_cloud=[1e-5])[0]
# Calculate the height of the tropopause
z_pv2 = interpolate.to_level(z, ertel_potential_vorticity=[2])[0]
# Calculate the distance from tropopause to cloud
dz = z_pv2.data - z_cloud.data
dz = z_pv2.copy(data=dz)
# Plot
plot.pcolormesh(dz, **kwargs)
plt.show()
def create_startpoints(cubes, coordinate, values):
"""Create an array of startpoints
Output array has shape Nx3 where N is the number of startpoints and the 3
indices are for longitude, latitude and altitude.
Args:
cubes (iris.cube.CubeList)
Returns:
train (np.array): The input array to trajectory calculations
"""
# Get the values of altitude at the given coordinate
P = convert.calc(coordinate, cubes)
z = grid.make_cube(P, 'altitude')
z.add_aux_coord(grid.make_coord(P), [0, 1, 2])
z500 = interpolate.to_level(z, **{coordinate: values})[0].data
lon, lat = grid.get_xy_grids(P)
# Convert the 2d arrays to an Nx3 array
train = np.array([lon.flatten(),
lat.flatten(),
z500.flatten()]).transpose()
return train
def get_startpoints(cubes):
"""Find all points within 1km of the tropopause that have a negative PV
anomaly
"""
# Find the tropopause height
pv = convert.calc('ertel_potential_vorticity', cubes)
grid.add_hybrid_height(pv)
z = grid.make_cube(pv, 'altitude').data
zpv2 = get_tropopause_height(pv).data * np.ones_like(z)
# Find point below the tropopause
tropopause_relative = np.logical_and(z < zpv2, z > zpv2 - 2000)
# Find negative pv anomalies
diff = convert.calc('total_minus_advection_only_pv', cubes)
negative_pv = diff.data < -1
# Combine criteria
criteria = np.logical_and(negative_pv, tropopause_relative)
# Get indices where criterion is met
indices = np.where(criteria)
# Convert to lon, lat, z
longitude = pv.coord('grid_longitude').points
lons = np.array([longitude[idx] for idx in indices[2]])
latitude = pv.coord('grid_latitude').points
lats = np.array([latitude[idx] for idx in indices[1]])
height = pv.coord('level_height').points
altitude = np.array([height[idx] for idx in indices[0]])
# Make start points array
start_points = np.array([lons, lats, altitude]).transpose()
return start_points
def brunt_vaisala_squared(theta, zcoord='altitude'):
r"""Calculate the squared Brunt-Vaisala frequency
:math:`N^2 = \frac{g}{\theta} \frac{d\theta}{dz}`
Uses a Centred difference to calculate the vertical gradient of potential
temperature
Args:
theta (iris.cube.Cube): Potential temperature
zcoord (str): The name of the vertical (z) coordinate
Returns:
iris.cube.Cube: Squared Brunt-Vaisala frequency at theta points
"""
# Differentiate theta with respect to altitude
z = grid.make_cube(theta, zcoord)
dtheta_dz = calculus.multidim(theta, z, 'z')
# Interpolate derivatives back to original levels
zdim = theta.coord(axis='z')
dtheta_dz = interpolate.interpolate(dtheta_dz,
**{zdim.name(): zdim.points})
# Calculate Brunt-Vaisala frequency squared
N_sq = (dtheta_dz / theta) * constants.g
# Correct units because iris is shit sometimes
N_sq.units = 's-2'
return N_sq
def Eady(theta,
u,
P,
pressure_levels=(85000, 40000),
theta_0=iris.cube.Cube(300, units='K')):
r"""Calculate the Eady growth rate
:math:`\sigma = 0.31 \frac{f}{N} |\frac{dU}{dz}|`
Where :math:`f = 2 \Omega sin(\phi)`
and :math:`N^2 = \frac{g}{\theta_0} \frac{d \theta}{dz}`
Args:
theta (iris.cube.Cube): Air potential temperature
u (iris.cube.Cube): Zonal wind speed
P (iris.cube.Cube): Air pressure
pressure_levels:
theta_0:
Returns:
iris.cube.Cube: Eady growth rate
"""
# Calculate Coriolis parameter
f = coriolis_parameter(theta)
# Extract altitude as a cube
z = grid.make_cube(P, 'altitude')
# Add pressure as a coordinate to the cubes
P = grid.make_coord(P)
for cube in (theta, u, z):
cube.add_aux_coord(P, [0, 1, 2])
# Interpolate variables to pressure levels
theta = interpolate.to_level(theta, air_pressure=pressure_levels)
u = interpolate.to_level(u, air_pressure=pressure_levels)
z = interpolate.to_level(z, air_pressure=pressure_levels)
# Calculate the Brunt-Vaisala frequency
dtheta_dz = calculus.multidim(theta, z, 'z')
N_sq = dtheta_dz * (constants.g / theta_0)
N_sq.units = 's-2'
N = N_sq**0.5
# Calclate the wind shear
du_dz = calculus.multidim(u, z, P.name())
du_dz.data = np.abs(du_dz.data)
# Calculate the Eady index
sigma = 0.31 * (du_dz / N) * f
return sigma
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 height(cube):
"""Extract height coordinate as a cube
Args:
cube (iris.cube.Cube):
Returns:
iris.cube.Cube:
"""
z = grid.make_cube(cube, 'altitude')
if z.shape != cube.shape:
z = grid.broadcast_to_cube(z, cube)
return z
def multidim(y, x, axis):
"""Calculate a derivative against a multi dimensional coordinate
The function :py:func:`iris.analysis.calculus.differentiate` is a useful and
efficient function for calculating a centred finite difference derivative of
a multidimensional array with respect to a coordinate spanning a single
dimension.
This function is a way of extending that function to a coordinate spanning
more than one dimension. An example is if the cube is in terms of model
levels but we want the vertical derivative with respect to pressure. Simply
by calculating the derivative of the cube and pressure with respect to
model levels then applying the chain rule we get the derivative of the cube
with respect to pressure at the midpoints in the vertical for all
grid-points.
args:
y, x (iris.cube.Cube): The two multidimensional coordinates to
calculate the derivative. Must be on the same grid and have a
one dimensional coordinate for the specified axis. Alternatively x
can be a string corresponding to a 3d coordinate in y
axis (str): A letter marking the axis to differentiate with respect to
(x,y,z,t) or the name of a coordinate on one of these axes.
returns:
iris.cube.Cube: The input cube y differentiated with respect to x along
the chosen axis.
"""
# Extract x from y as cube if it is specified as a string
if type(x) == str:
x = grid.make_cube(y, x)
# Get the coordinate for the axis to be differentiated over
try:
k = y.coord(axis=axis, dim_coords=True)
except ValueError:
# Allow referencing the coordinate by name
k = y.coord(axis)
# Calculate derivatives with respect to single dimensional coordinate
dy_dk = differentiate(y, k)
dx_dk = differentiate(x, k)
# Apply the chain rule
dy_dx = dy_dk / dx_dk
return dy_dx
def get_tropopause_height(pv):
"""Find the altitude where pv=2 is crossed
Args:
pv (iris.cube.Cube):
Returns:
zpv2 (iris.cube.Cube): 2-dimensional cube with the altitude of the
highest 2 PVU surface
"""
z = grid.make_cube(pv, 'altitude')
pv = grid.make_coord(pv)
z.add_aux_coord(pv, [0, 1, 2])
zpv2 = interpolate.to_level(z, ertel_potential_vorticity=[2])[0]
return zpv2
def surface_height(cube):
"""Extract height coordinate as a cube
Args:
cube (iris.cube.Cube)
Returns:
iris.cube.Cube:
"""
z_s = grid.make_cube(cube[0], 'surface_altitude')
for aux_factory in z_s.aux_factories:
z_s.remove_aux_factory(aux_factory)
for coord in z_s.aux_coords:
z_s.remove_coord(coord)
return z_s
def polar_horizontal(cube, axis):
r"""Differentiate a cube with respect to lon/lat and convert to x/y
:math:`dx = r cos(\phi) d\lambda`
:math:`dy = r d\phi`
args:
cube (iris.cube.Cube):
axis (str): A letter marking the axis to differentiate with respect to
(x,y).
returns:
iris.cube.Cube: The derivative of the cube along the chosen polar axis
"""
# Calculate derivative with respect to polar axis in radians
diff = diff_by_axis(cube, axis) / constants.radians_per_degree
# Convert the differential to x/y by considering the distance in lon/lat
# Coordinates
# Calculate radius relative to Earth centre
radius = grid.make_cube(diff, 'altitude')
radius.data += constants.earth_avg_radius.data
if radius.ndim == 1:
radius = iris.util.broadcast_to_shape(radius.data, diff.shape, [0])
radius = diff.copy(data=radius)
radius.rename("altitude")
radius.units = "m"
if axis.lower() == 'x':
lat = (diff.coord(axis='y').points * constants.radians_per_degree.data)
lat = np.outer(lat, np.ones(diff.shape[-1]))
metres_per_radian = radius * np.cos(lat)
elif axis.lower() == 'y':
metres_per_radian = radius
else:
raise ValueError('Can only specify x or y axis')
metres_per_radian.units = 'm radian-1'
diff = diff / metres_per_radian
return diff
def main(cubes, dz):
# Calculate dp/dz (Hydrostatic balance dp/dz = -\rho g
P = convert.calc('air_pressure', cubes)
z = grid.make_cube(P, 'altitude')
dP_dz = calculus.multidim(P, z, 'z')
dP_dz = remap_3d(dP_dz, P)
# Calculate absolute vorticity
u = convert.calc('x_wind', cubes)
v = convert.calc('y_wind', cubes)
du_dy = calculus.diff_by_axis(u, 'y')
du_dy = remap_3d(du_dy, P)
dv_dx = calculus.diff_by_axis(v, 'x')
dv_dx = remap_3d(dv_dx, P)
vort = du_dy.data - dv_dx.data
lat = grid.true_coords(P)[1]
abs_vort = 2 * convert.omega.data * np.cos(lat) * vort
# Calculate Nsq for each PV tracer
tracers = convert.calc(names, cubes)
nsq_0 = variable.N_sq(convert.calc('air_potential_temperature', cubes))
nsq = []
nsq.append(nsq_0)
for tracer in tracers:
nsq_i = -1 * tracer * dP_dz / abs_vort
nsq_i.rename(tracer.name())
nsq.append(nsq_i)
# Create an average profile
thetapv2 = convert.calc('air_potential_temperature',
cubes,
levels=('ertel_potential_vorticity', [2]))
ridges, troughs = rossby_waves.make_nae_mask(thetapv2)
pv = grid.make_coord(convert.calc('advection_only_pv', cubes))
z.add_aux_coord(pv, [0, 1, 2])
zpv = interpolate.to_level(z, advection_only_pv=[3.5])[0]
y = diagnostics.profile(nsq, zpv, dz, mask=troughs)
# Plot nsq
plot.multiline(y)
plt.show()
def pv_invert(pv,
boundary_theta,
thrs1=0.1,
relaxation_parameter=0.4,
max_iterations=1999,
max_cycles=999,
omega_s=1.7,
omega_h=1.4):
"""Inverts pv to find balanced geopotential and streamfunction
Args:
pv (iris.cube.Cube):
boundary_theta (iris.cube.Cube):
"""
threshold = gravity * thrs1 / tho
lon, lat = grid.true_coords(pv)
laplacian_coefficients = coefficients(len(lat))
coriolis_parameter = 2 * omega * np.sin(lat)
cos_latitude = np.cos(lat)
geopotential_height, streamfunction = first_guess(pv)
# Extract the vertical coordinate from the cube
pressure = grid.make_cube(pv, 'air_pressure')
exner = variable.exner(pressure)
# Perform the PV inversion
geopotential_height, streamfunction, pv_out, boundary_theta = \
pvi.balnc(coriolis_parameter, cos_latitude, laplacian_coefficients,
geopotential_height, streamfunction, pv.data,
boundary_theta.data, exner, relaxation_parameter, threshold,
max_iterations, max_cycles, omega_s, omega_h)
return pv_out, geopotential_height, streamfunction
def geopotential_height(P, levels):
"""Calculate height above sea level on pressure levels
Args:
P (iris.cube.Cube):
levels (list or numpy.ndarray): The values of pressure to output the
geopotential height on
Returns:
iris.cube.Cube: Height on pressure levels
"""
# Create a height cube
z = grid.make_cube(P, 'altitude')
# Add pressure as a coordinate to the height
pressure_coord = grid.make_coord(P)
z.add_aux_coord(pressure_coord, range(z.ndim))
# Interpolate the height on to pressure levels
geopotential = interpolate.to_level(z, **{P.name(): levels})
geopotential.rename('Geopotential_height')
return geopotential
