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 create_startpoints(cubes, levels, stride=1):
"""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
z = convert.calc('altitude', cubes, levels=levels)
# Get the grid latitude and longitude as 2d arrays
lon, lat = grid.get_xy_grids(z)
# Convert the 2d arrays to an Nx3 array
nz = z.shape[0]
lon = np.tile(lon[::stride, ::stride].flatten(), nz)
lat = np.tile(lat[::stride, ::stride].flatten(), nz)
z = z.data[:, ::stride, ::stride].flatten()
train = np.array([lon, lat, z]).transpose()
return train
def track_position(mslp, idx_0):
# Set the start point of the cyclone manually
lon, lat = grid.get_xy_grids(mslp)
current_x, current_y = lon[idx_0], lat[idx_0]
nt = mslp.shape[0]
track = [[current_x - 360, current_y]]
indices = [idx_0]
# Take the nearest local minima in sea-level pressure at each timestep
for n in range(nt):
print(n)
local_minima = find_local_minima(mslp[n].data)
new_x, new_y, idx = find_nearest_minima(current_x, current_y,
local_minima, lon, lat)
track.append([new_x - 360, new_y])
indices.append(idx)
current_x, current_y = new_x, new_y
# Save the data as a numpy array s
track = np.array(track)
indices = np.array(indices)
return track, indices
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 get_points(dtheta, pv):
# Smooth the theta_adv and pv data first
dtheta.data = filters.gaussian_filter(dtheta.data, sigma=10, truncate=4)
pv.data = filters.median_filter(pv.data, size=20)
# Extract the contour surrounding the outflow region
criteria = np.logical_and(pv.data < 2, dtheta.data > 0)
pv.data = criteria.astype(int)
cs = iplt.contour(pv, [0.5])
contours = tropopause_contour.get_contour_verts(cs)
closed_loop = tropopause_contour.get_tropopause_contour(contours[0])
path = mpl.path.Path(closed_loop)
# Create an array containing all the grid points within the outflow region
lon, lat = grid.get_xy_grids(dtheta)
points = np.transpose(np.array([lon.flatten(), lat.flatten()]))
points = points[np.where(path.contains_points(points))]
return closed_loop, points
def calc_circulation(trajectories, forecast, theta_level, dtheta):
# Select an individual theta level
trajectories = trajectories.select(
'air_potential_temperature', '==', theta_level)
print(len(trajectories))
levels = ('air_potential_temperature', [theta_level])
results = iris.cube.CubeList()
for n, cubes in enumerate(forecast):
print(n)
if n == 0:
# Load grid parameters
example_cube = convert.calc('upward_air_velocity', cubes,
levels=levels)
# Create a 1d array of points for determining which gridpoints are
# contained in the trajectory circuit when performing volume
# integrals
glon, glat = grid.get_xy_grids(example_cube)
gridpoints = np.array([glon.flatten(), glat.flatten()]).transpose()
cs = example_cube.coord_system()
# Load trajectory positions -(n+2) because the trajectories are
# backwards in time. +2 to skip the analysis which is not in the
# forecast object (i.e. n=0 corresponds to idx=-2 in the trajectories)
x = trajectories.x[:, -(n + 2)]
y = trajectories.y[:, -(n + 2)]
z = trajectories['altitude'][:, -(n + 2)]
u = trajectories['x_wind'][:, -(n + 2)]
v = trajectories['y_wind'][:, -(n + 2)]
w = trajectories['upward_air_velocity'][:, -(n + 2)]
# Integrals are invalid once trajectories leave the domain but we don't
# want to stop the script so just print out the number of trajectories
# that have left the domain
leftflag = (trajectories['air_pressure'][:, -(n+2)] < 0).astype(int)
leftcount = np.count_nonzero(leftflag)
print(leftcount)
# Calculate enclosed area integrals
integrals = mass_integrals(cubes, x, y, glat, gridpoints,
theta_level, dtheta)
for icube in integrals:
# Set integrals to zero if trajectories have left the domain
if leftcount > 0:
icube.data = 0.
results.append(icube)
# Convert to global coordinates in radians
u, v, lon, lat = get_geographic_coordinates(u, v, x, y, cs)
# Unrotated coordinates in radians
lon = np.deg2rad(lon)
lat = np.deg2rad(lat)
# Calculate the velocity due to Earth's rotation
u_abs = Omega.data * (a+z) * np.cos(lat)
u += u_abs
# Integrate around the circuit
if leftcount > 0:
circulation = 0
else:
circulation = circuit_integral_rotated(u, v, w, lon, lat, z)
ccube = icube.copy(data=circulation)
ccube.rename('circulation')
ccube.units = 's-1'
results.append(ccube)
iris.save(results.merge(),
datadir + 'circulations_' + str(theta_level) + 'K.nc')
return