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