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