def main(cubes):
lon, lat = grid.true_coords(cubes[0])
z300 = convert.calc('altitude',
cubes,
levels=('equivalent_potential_temperature', [300]))[0]
z300.convert_units('km')
mslp = convert.calc('air_pressure_at_sea_level', cubes)
mslp.convert_units('hPa')
# Plot overview
plot.contourf(z300, np.linspace(0, 10, 11))
cs = iplt.contour(mslp, np.linspace(950, 1050, 11), colors='k')
plt.clabel(cs, fmt='%1.0f')
# Warm Sector
warm_sector = z300.data.mask
plim = mslp.data < 1000
loc = np.logical_and(lon > -10, lon < 5)
ws_mask = np.logical_and(np.logical_and(warm_sector, loc), plim)
ws_mask = mslp.copy(data=ws_mask)
iplt.contour(ws_mask, linestyles='--', colors='r')
# Cold Sector
cold_sector = z300.data < 5
loc = np.logical_and(loc, lat < 65)
cs_mask = np.logical_and(np.logical_and(cold_sector, loc), plim)
cs_mask = mslp.copy(data=cs_mask)
iplt.contour(cs_mask, linestyles='--', colors='b')
plt.title(r'$z(\theta_e = 300)$')
return
def create_startf(cubes, theta_level, output_file):
"""Creates a startfile for lagranto for the whole grid on a single
isentropic surface
Reference date 20111129_2300 / Time range 0 min
time lon lat p level
---------------------------------------
0.00 -10.000 50.000 258 320.000
"""
# Find pressure on theta level
P = convert.calc('air_pressure',
cubes,
levels=('air_potential_temperature', theta_level))[0].data
# Get longitude and latitude
lon, lat = grid.true_coords(P)
# Create startfile
with open(output_file, 'w') as output:
output.write('Reference date 20111129_2300 / Time range 0 min\n')
output.write('\n')
output.write(' time lon lat p level\n')
output.write('---------------------------------------\n')
output.write('\n')
ny, nx = P.shape
for j in xrange(ny):
for i in xrange(nx):
output.write(' 0.00 ' + str(round(lon[j, i], 3)) + ' ' +
str(round(lat[j, i], 3)) + ' ' +
str(round(P[j, i] / 100.0, 0)) + ' ' +
str(theta_level) + '\n\n')
def coriolis_parameter(cube):
r"""Calculate the Coriolis parameter on the xy grid of the given cube
:math:`f = 2 \Omega sin(\phi)`
Args:
cube (iris.cube.Cube): Any cube with lon/lat coordinates
Returns:
iris.cube.Cube: Coriolis parameter as function of lon/lat
"""
# Calculate the Coriolis parameter
lat = grid.true_coords(cube)[1]
f = 2 * constants.omega.data * np.sin(np.deg2rad(lat))
# Put the output into a cube
f = iris.cube.Cube(f,
long_name='coriolis_parameter',
units='s-1',
attributes=cube.attributes,
dim_coords_and_dims=[(cube.coord(axis='y',
dim_coords=True), 0),
(cube.coord(axis='x',
dim_coords=True), 1)])
return f
def main(cubes):
theta = convert.calc('equivalent_potential_temperature', cubes)
P = convert.calc('air_pressure', cubes)
P.convert_units('hPa')
mass = convert.calc('mass', cubes)
lon, lat = grid.true_coords(theta)
lonmask = np.logical_or(lon < -15, lon > 5)
latmask = np.logical_or(lat < 47.5, lat > 62.5)
areamask = np.logical_or(lonmask, latmask)
masks = [
np.logical_or(theta.data < theta_front, areamask),
np.logical_or(theta.data > theta_front, areamask)
]
#overview(cubes, areamask)
#plt.savefig(plotdir + 'composite_iop8_24h_overview_750hpa.pdf')
# plt.show()
z_cold, z_warm = bl_heights(cubes, theta, areamask)
# Initialise the plot
fig = plt.figure(figsize=(18, 12))
diags = convert.calc(names, cubes)
masses = []
# Rows are different masks
for n, mask in enumerate(masks):
means = diagnostics.averaged_over(diags, levels, P, mass, mask)
masses.append(convert.calc('mass', means))
# Columns are for different mappings
for m, mapping in enumerate(mappings):
ax = plt.subplot2grid((2, ncol), (n, m))
composite(means, ax, mapping, mass, P)
add_trimmings(ax, n, m)
if m == 0:
ax.set_xlim(0.1, 1.3)
elif m == 1:
ax.set_xlim(-0.6, 0.6)
else:
ax.set_xlim(-1, 1)
multilabel(ax, 3 * n + m, yreversed=True, fontsize=25)
if n == 0:
plt.axhline(z_cold, color='k', linestyle='--')
elif n == 1:
plt.axhline(z_warm, color='k', linestyle='--')
add_figlabels(fig)
fig.savefig(plotdir + 'composite_iop5_24h.pdf')
plt.figure()
for mass in masses:
qplt.plot(mass, mass.coord('air_pressure'))
ax.set_ylim(950, 500)
plt.show()
return
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 forward_trajectories(forecast):
"""Calculate 48 hour forward trajectories from low levels
Start trajectories from all points below 2km
"""
cubes = forecast.set_lead_time(hours=48)
z = convert.calc('altitude', cubes)
theta = convert.calc('air_potential_temperature', cubes)
theta_adv = convert.calc('advection_only_theta', cubes)
pv = convert.calc('ertel_potential_vorticity', cubes)
lon, lat = grid.true_coords(pv)
glon, glat = grid.get_xy_grids(pv)
time = grid.get_datetime(pv)
nz, ny, nx = pv.shape
eqlats = rossby_waves.eqlats
cs = iris.Constraint(time=time)
with iris.FUTURE.context(cell_datetime_objects=True):
eqlat = eqlats.extract(cs)[0]
# Interpolate to the theta and PV surfaces
eqlat = interpolate.main(eqlat, ertel_potential_vorticity=2)
eqlat = interpolate.main(eqlat,
potential_temperature=theta.data.flatten())
# Define the start points
trainp = []
for k in range(nz):
print(k)
for j in range(ny):
for i in range(nx):
if (theta_adv.data[k, j, i] < 300 < theta.data[k, j, i] and
pv.data[k, j, i] < 2 and
lat[j, i] > eqlat.data[k * ny * nx + j * nx + i]):
trainp.append([glon[j, i], glat[j, i], z.data[k, j, i]])
trainp = np.array(trainp)
plot.pcolormesh(pv[33], vmin=0, vmax=10, cmap='cubehelix_r', pv=pv[33])
plt.scatter(trainp[:, 0], trainp[:, 1])
plt.show()
# Calculate the trajectories
tracers = ['air_potential_temperature', 'air_pressure']
traout = caltra.caltra(trainp, mapping, fbflag=-1, tracers=tracers)
# Save the trajectories
traout.save(datadir + 'backward_trajectories.pkl')
def main(cubes, levels, *args, **kwargs):
# Initialise the plot
fig = plt.figure(figsize=(18, 20))
pv = convert.calc('ertel_potential_vorticity', cubes, levels=levels)[0]
theta = convert.calc('equivalent_potential_temperature', cubes,
levels=levels)[0][slices]
rh = convert.calc('relative_humidity', cubes, levels=levels)[0][slices]
lon, lat = grid.true_coords(theta)
lonmask = np.logical_or(lon < -15, lon > 5)
latmask = np.logical_or(lat < 47.5, lat > 62.5)
areamask = np.logical_or(lonmask, latmask)
mask = theta.copy(data=areamask.astype(int))
for n, name in enumerate(names):
row = n / ncols
col = n - row * ncols
print(row, col)
ax = plt.subplot2grid((nrows, ncols), (row, col))
cube = convert.calc(name, cubes, levels=levels)[0] # [slices]
im = iplt.contourf(cube, *args, **kwargs)
#iplt.contour(pv, [2], colors='k', linewidths=2)
iplt.contour(mask, [0.5], colors='k', linestyles='--')
add_map()
plt.title(second_analysis.all_diagnostics[name].symbol)
iplt.contour(theta, [300], colors='k', linewidths=2)
iplt.contour(rh, [0.8], colors='w', linewidths=2)
iplt.contourf(rh, [0.8, 2], colors='None', hatches=['.'])
for n, ax in enumerate(fig.axes):
multilabel(ax, n)
cbar = plt.colorbar(im, ax=fig.axes, orientation='horizontal',
fraction=0.05, spacing='proportional')
cbar.set_label('PVU')
cbar.set_ticks(np.linspace(-2, 2, 17)[::2])
plt.savefig(plotdir + 'iop5_pv_tracers_24h_' +
str(levels[1][0])[0:3] + 'hpa.pdf')
# plt.show()
return
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 composite(cubes):
lon, lat = grid.true_coords(cubes[0])
theta_e = convert.calc('equivalent_potential_temperature', cubes)
mslp = convert.calc('air_pressure_at_sea_level', cubes)
mslp.convert_units('hPa')
mass = convert.calc('mass', cubes)
# Warm Sector
warm_sector = theta_e.data > 300
plim = mslp.data < 1000
loc = np.logical_and(lon > -10, lon < 5)
ws_mask = np.logical_and(np.logical_and(warm_sector, loc), plim)
ws_mask = theta_e.copy(data=ws_mask)
for n in range(20):
c = (n + 1) / 20
iplt.contour(ws_mask[n], linestyles=':', colors=[(c, c, c)])
plt.show()
return
def forecast_theta(cubes):
theta = convert.calc('air_potential_temperature',
cubes,
levels=('ertel_potential_vorticity', [2]))[0]
theta.data = np.ma.masked_where(theta.data > theta_max, theta.data)
iplt.contourf(theta, levels)
plt.gca().coastlines()
cb = plt.colorbar(orientation='horizontal')
cb.set_label('K')
plt.title('(c)'.ljust(30) + r'$\theta (\lambda, \phi, q=2)$'.ljust(60))
lon, lat = grid.true_coords(theta)
lon = theta.copy(data=lon)
lat = theta.copy(data=lat)
add_gridlines(lon, lat)
return theta, lon, lat
def thetapv2_nae(infile):
"""Re-arrange the data as theta(time, grid_lat, grid_lon)
"""
# Load the equivalent latitude data on PV2
eqlats_cube, theta, eqlats = theta_pv2(infile)
# Load the NAE grid
cubes = iris.load(directory + '../iop5/diagnostics_024.nc')
lat = grid.true_coords(cubes[0])[1]
# Initialise the output
nt, ntheta = eqlats.shape
ny, nx = lat.shape
output = np.zeros([nt, ny, nx])
# Search downward for the equivalent latitude value
# Start at theta=400 K
k_start = np.abs(theta - 400).argmin()
for n in range(nt):
print(n)
for j in range(ny):
for i in range(nx):
# Search downward
k = k_start
while eqlats[n, k] < lat[j, i] and k < ntheta:
k += 1
# Linearly interpolate to find theta
alpha = ((lat[j, i] - eqlats[n, k]) /
(eqlats[n, k] - eqlats[n, k - 1]))
output[n, j, i] = (alpha * theta[k] +
(1 - alpha) * theta[k - 1])
output = iris.cube.Cube(
output, long_name='potential_temperature', units='K',
dim_coords_and_dims=[(eqlats_cube.coord('time'), 0),
(cubes[0].coord('grid_latitude'), 1),
(cubes[0].coord('grid_longitude'), 2)])
return output
def main(cubes, theta_value, **kwargs):
"""Plot PV on theta and the equivalent latitude circle
Args:
cubes (iris.cube.CubeList): Contains variables to calculate PV and
potential temperature
"""
pv = convert.calc('ertel_potential_vorticity', cubes,
levels=('air_potential_temperature', [theta_value]))[0]
# Add equivalent latitude
lon, lat = grid.true_coords(pv)
lat = pv.copy(data=lat)
time = grid.get_datetime(pv)[0]
time = PDT(month=time.month, day=time.day, hour=time.hour)
eqlat = rossby_waves.equivalent_latitude(time, theta_value, 2)
# Plot PV on theta
plot.pcolormesh(pv, pv=pv, **kwargs)
iplt.contour(lat, [eqlat.data], colors='r', linewidths=2)
plt.title('')
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 main(cubes, **kwargs):
# Pre-calculate parameters on every plot
mslp = convert.calc('air_pressure_at_sea_level', cubes)
mslp.convert_units('hPa')
theta = convert.calc('equivalent_potential_temperature',
cubes,
levels=levels)[0]
rh = convert.calc('relative_humidity', cubes)
fig = plt.figure(figsize=(18, 15))
lon, lat = grid.true_coords(theta)
lonmask = np.logical_or(lon < -15, lon > 5)
latmask = np.logical_or(lat < 47.5, lat > 62.5)
areamask = np.logical_or(lonmask, latmask)
# Plot overview
ax1 = plt.subplot2grid((2, 2), (0, 0))
iplt.contourf(theta, theta_levs, cmap='coolwarm', extend='both')
add_map()
cs = iplt.contour(mslp, plevs, colors='k', linewidths=2)
plt.clabel(cs, fmt='%1.0f')
mask = theta.copy(data=areamask.astype(int))
iplt.contour(mask, [0.5], colors='k', linestyles='--')
count = 0
for n, (xs, xf, ys, yf) in enumerate(points, start=1):
# Label the cross section points
plt.plot([xs, xf], [ys, yf], '-kx')
plt.text(xs, ys, ascii_uppercase[count], color='k', fontsize=20)
count += 1
plt.text(xf, yf, ascii_uppercase[count], color='k', fontsize=20)
count += 1
multilabel(plt.gca(), 0)
theta = convert.calc('equivalent_potential_temperature', cubes)
# Plot cross sections
titles = ['AB', 'CD', 'EF']
coords = ['grid_longitude', 'altitude']
for n, (xs, xf, ys, yf) in enumerate(points, start=1):
row = n / ncols
col = n - row * ncols
ax = plt.subplot2grid((nrows, ncols), (row, col))
theta_cs = cs_cube(theta, xs, xf, ys, yf)
rh_cs = cs_cube(rh, xs, xf, ys, yf)
im = iplt.contourf(theta_cs,
theta_levs,
coords=coords,
cmap='coolwarm',
extend='both')
iplt.contour(rh_cs, rh_levs, coords=coords, colors='w')
iplt.contourf(rh_cs, [0.8, 2],
coords=coords,
colors='None',
hatches=['.'])
ax.set_ylim(0, 7)
if xs > xf:
ax.invert_xaxis()
multilabel(ax, n, xreversed=True)
else:
multilabel(ax, n)
ax.set_title(titles[n - 1])
ax.set_ylabel('Height (km)')
ax.set_xlabel('Grid Longitude')
# Add colorbar at bottom of figure
cbar = plt.colorbar(im,
ax=fig.axes,
orientation='horizontal',
fraction=0.05,
spacing='proportional')
cbar.set_label('PVU')
fig.savefig(filename)
plt.show()
def main(cubes):
# Read potential vorticity at 320-K isentropic surface
pv = convert.calc('ertel_potential_vorticity', cubes,
levels=('air_potential_temperature', [theta]))[0]
# Set coordinate with north pole in the centre
lon, lat = grid.true_coords(pv)
xp = (90 - lat) * np.sin(lon * np.pi / 180)
yp = - (90 - lat) * np.cos(lon * np.pi / 180)
time = grid.get_datetime(pv)
eqlat = rossby_waves.equivalent_latitude(time, theta, pv_trop).data
# Find the contour of 2 PVU that wraps around the globe
cs = plt.contour(xp, yp, pv.data, [pv_trop])
contours_fct = tropoc.get_contour_verts(cs)
tropopause_contour = tropoc.get_tropopause_contour(contours_fct[0])
if tropoc.is_closed_contour(tropopause_contour):
# if tropoc.is_closed_contour(tropopause_contour):
path = Path(tropopause_contour)
points = np.transpose(np.array([xp.flatten(), yp.flatten()]))
inside = np.reshape(path.contains_points(points), pv.shape)
# Create a category map
vortex = np.logical_and(
inside, np.logical_and(lat > eqlat, pv.data > pv_trop))
ridge = np.logical_and(
np.logical_not(inside), np.logical_and(lat > eqlat,
pv.data < pv_trop))
trough = np.logical_and(
inside, np.logical_and(lat < eqlat, pv.data > pv_trop))
cut_off = np.logical_and(np.logical_not(inside), pv.data > pv_trop)
cut_on = np.logical_and(inside, pv.data < pv_trop)
cat_map = (vortex.astype(int) +
2 * ridge.astype(int) +
3 * trough.astype(int) +
4 * cut_off.astype(int) +
5 * cut_on.astype(int))
plt.clf()
plt.pcolormesh(xp, yp, cat_map, vmin=0, vmax=6, cmap=cmap)
plt.contour(xp, yp, pv.data, [2], colors='k')
plt.plot(tropopause_contour[:, 0], tropopause_contour[:, 1], color='w')
# Add land mask
z = convert.calc('altitude', cubes)
land = z[0].data != 20
plt.contour(xp, yp, land, [0.5], colors='k', linewidths=1)
plt.xlim(-60, 60)
plt.ylim(-60, 60)
plt.savefig(plotdir + 'output/' + 'category_map_pv2_theta' +
str(theta) + '.png')
else:
plt.clf()
plt.plot(tropopause_contour[:, 0], tropopause_contour[:, 1])
print('Open contour')
plt.show()
return ridge
