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 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 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 main(cubes, dz, name, **kwargs):
"""Produces cross section plots above and below the tropopause
"""
# Extract cube to be plotted
cube = convert.calc(name, cubes)
# Find the height of the tropopause
ztrop, fold_t, fold_b = tropopause.height(cubes)
# Produce a new co-ordinate above and below the tropopause
ny, nx = ztrop.shape
new_coord = np.zeros([3, ny, nx])
new_coord[0, :, :] = ztrop.data - dz
new_coord[1, :, :] = ztrop.data
new_coord[2, :, :] = ztrop.data + dz
# Interpolate the cubes to be plotted to the coordinate above and below the
# tropopause
plotcube = interpolate.to_level(cube, altitude=new_coord)
# Plot the cross sections
for n in xrange(3):
plt.figure()
plot.pcolormesh(plotcube[n], **kwargs)
plt.show()
def calc(names, cubelist, levels=None):
"""Wrapper function to ``calculate`` with the option to interpolate to a
level
Args:
names (str or list[str]): The CF standard name(s) of the
variable(s) to be calculated.
cubelist (iris.cube.CubeList): Contains either the requested variable or
the variables required to calculate the requested variable.
levels (tuple or None): The name and levels to interpolate the
requested variable to. Default is None.
Returns:
iris.cube.Cube or iris.cube.CubeList: The variable(s) requested.
Optionally interpolated to the given level
"""
if type(names) == str:
names = [names]
# Call calculate to get the cube
cubes = [calculate(name, cubelist) for name in names]
# Return the cube if interpolation is not requested
if levels is None:
if len(names) == 1:
return cubes[0]
else:
return iris.cube.CubeList(cubes)
else:
output = iris.cube.CubeList()
# Extract the coordinate requested
coord_name, values = levels
for cube in cubes:
# Extract the coordinate from the cubes
try:
coord = cube.coord(coord_name)
# Alternatively use a cube from the cubelist
except iris.exceptions.CoordinateNotFoundError:
coord = calculate(coord_name, cubelist)
coord = grid.make_coord(coord)
cube.add_aux_coord(coord, range(cube.ndim))
# Interpolate to the requested coordinate levels
if coord.points.ndim == 1:
result = cube.interpolate([(coord_name, values)],
iris.analysis.Linear())
else:
result = interpolate.to_level(cube, **{coord_name: values})
output.append(result)
if len(names) == 1:
return output[0]
else:
return output
def isentropic_circulation(pv, pressure, mask=None):
r"""Calculate the circulation associated with a given PV anomaly
:math:`C_{\theta} = \iint_R \sigma Q dx dy`
args:
pv (iris.cube.Cube or iris.cube.CubeList): Potential vorticity on
isentropic levels
pressure (iris.cube.Cube): Pressure at gridpoints
mask (numpy.ndarray): Mask for area to integrate over
returns:
iris.cube.CubeList: The circulation from each PV in q
"""
# Make sure to iterate over list or cubelist
if type(pv) == iris.cube.Cube:
pv = [pv]
# Calculate pressure halfway between theta levels
theta_levs = pv[0].coord('air_potential_temperature').points
dtheta = theta_levs[1] - theta_levs[0]
theta_levs = list(theta_levs - dtheta / 2)
theta_levs.append(theta_levs[-1] + dtheta)
p_theta = interpolate.to_level(pressure,
air_potential_temperature=theta_levs)
# Calculate isentropic theta gradient at theta levels
dp_dtheta = calculus.differentiate(p_theta, 'air_potential_temperature')
# Calculate isentropic density
sigma = -1 * dp_dtheta / constants.g
# Apply the mask to sigma as it is used in each calculation
if mask is not None:
sigma.data = np.ma.masked_where(mask, sigma.data)
# Extract the xy coordinate names to collapse the cube over
coords = [sigma.coord(axis=axis).name() for axis in ['x', 'y']]
# Calculate the weights to perform area integration
# weights = iris.analysis.cartography.area_weights(sigma)
weights = grid.volume(pressure)[-1].data * np.ones_like(pv[0].data)
circulation = iris.cube.CubeList()
for pv_i in pv:
# Calculate the circulation
c_i = (sigma * pv_i).collapsed(coords,
iris.analysis.MEAN,
weights=weights)
# Give the cube a meaningful name
c_i.rename('circulation_due_to_' + pv_i.name())
circulation.append(c_i)
return circulation
def profile(x, surface, dz, mask=None, aggregator=MEAN, **kwargs):
"""Calculate an averaged vertical profile relative to a specified surface
Args:
x (iris.cube.Cube or iris.cube.CubeList): The variable(s) to produce
profiles of.
surface (iris.cube.Cube): 2d cube specifying the altitude of the
surface to interpolate relative to.
dz: Iterable container of floats prescribing the distances from the
given surface to interpolate to.
mask (array, optional): A binary array with the same shape as the cubes
data. Default is None.
aggregator (iris.analysis.Aggregator): The aggregator used to collapse
the cube on each level. Default is MEAN.
**kwargs: Additional keywords to be supplied to the
:py:func:`iris.cube.Cube.collapsed` method (e.g. weights).
Returns:
iris.cube.CubeList: A vertical profile for each variable in x
"""
# Check that inputs are correct
if type(x) == iris.cube.Cube:
x = [x.copy()]
elif type(x) == iris.cube.CubeList or type(x) == list:
x = [cube.copy() for cube in x]
else:
raise TypeError('Must specify cube or cubelist')
# Loop over each cube
cubes = iris.cube.CubeList()
for cube in x:
# Create a new coordinate for the cube
newcoord = cube.coord('altitude').points - surface.data
coord_name = 'distance_from_' + surface.name()
newcoord = AuxCoord(newcoord, long_name=coord_name, units='m')
cube.add_aux_coord(newcoord, [0, 1, 2])
# Interpolate the cube relative to the surface
newcube = interpolate.to_level(cube, **{coord_name: dz})
# Apply the mask to the cube
if mask is not None:
newcube.data = np.ma.masked_where(
mask * np.ones_like(newcube.data), newcube.data)
# Collapse the cube along the horizontal dimensions
xcoord = cube.coord(axis='X', dim_coords=True)
ycoord = cube.coord(axis='Y', dim_coords=True)
cubes.append(newcube.collapsed([xcoord, ycoord], aggregator, **kwargs))
return cubes
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 bl_heights(cubes, theta, areamask):
# Interpolate theta to bl height
z_bl = convert.calc('boundary_layer_height', cubes)
theta_bl = interpolate.to_level(theta, altitude=z_bl.data[None, :, :])[0]
# Theta on bl defines cold and warm sectors
cold_mask = np.logical_or(areamask, theta_bl.data > 300)
warm_mask = np.logical_or(areamask, theta_bl.data < 300)
# Calculate average bl height in cold and warm sectors
p = convert.calc('air_pressure',
cubes,
levels=('altitude', z_bl.data[None, :, :]))[0]
p.convert_units('hPa')
cold_z = np.ma.masked_where(cold_mask, p.data)
warm_z = np.ma.masked_where(warm_mask, p.data)
return cold_z.mean(), warm_z.mean()
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 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