def __init__(self, file_path = "", var_name = "",
bathymetry_path = "/skynet3_rech1/huziy/NEMO_OFFICIAL/dev_v3_4_STABLE_2012/NEMOGCM/CONFIG/GLK_LIM3_Michigan/EXP00/bathy_meter.nc"):
"""
:param file_path:
:param var_name:
:param bathymetry_path: used to mask land points
"""
self.current_time_frame = -1
self.var_name = var_name
self.cube = iris.load_cube(file_path, constraint=iris.Constraint(cube_func=lambda c: c.var_name == var_name))
self.lons, self.lats = cartography.get_xy_grids(self.cube)
lons2d_gl, lats2d_gl = nemo_commons.get_2d_lons_lats_from_nemo(path=bathymetry_path)
mask_gl = nemo_commons.get_mask(path=bathymetry_path)
xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons2d_gl.flatten(), lats2d_gl.flatten())
xt, yt, zt = lat_lon.lon_lat_to_cartesian(self.lons.flatten(), self.lats.flatten())
tree = cKDTree(list(zip(xs, ys, zs)))
dists, indices = tree.query(list(zip(xt, yt, zt)))
self.mask = mask_gl.flatten()[indices].reshape(self.lons.shape)
self.nt = self.cube.shape[0]
assert isinstance(self.cube, Cube)
print(self.nt)
def _map_common(draw_method_name, arg_func, mode, cube, data, *args, **kwargs):
"""
Draw the given cube on a map using its points or bounds.
"Mode" parameter will switch functionality between POINT or BOUND plotting.
"""
# get the 2d x and 2d y from the CS
if mode == iris.coords.POINT_MODE:
x, y = cartography.get_xy_grids(cube)
else:
try:
x, y = cartography.get_xy_contiguous_bounded_grids(cube)
# Exception translation.
except iris.exceptions.CoordinateMultiDimError:
raise ValueError("Could not get XY grid from bounds. "
"X or Y coordinate not 1D.")
except ValueError:
raise ValueError("Could not get XY grid from bounds. "
"X or Y coordinate doesn't have 2 bounds "
"per point.")
# take a copy of the data so that we can make modifications to it
data = data.copy()
# If we are global, then append the first column of data the array to the
# last (and add 360 degrees) NOTE: if it is found that this block of code
# is useful in anywhere other than this plotting routine, it may be better
# placed in the CS.
x_coord = cube.coord(axis="X")
if getattr(x_coord, 'circular', False):
_, direction = iris.util.monotonic(x_coord.points,
return_direction=True)
y = np.append(y, y[:, 0:1], axis=1)
x = np.append(x, x[:, 0:1] + 360 * direction, axis=1)
data = ma.concatenate([data, data[:, 0:1]], axis=1)
# Replace non-cartopy subplot/axes with a cartopy alternative.
cs = cube.coord_system('CoordSystem')
if cs:
cartopy_proj = cs.as_cartopy_projection()
else:
cartopy_proj = cartopy.crs.PlateCarree()
ax = _get_cartopy_axes(cartopy_proj)
draw_method = getattr(ax, draw_method_name)
# Set the "from transform" keyword.
# NB. While cartopy doesn't support spherical contours, just use the
# projection as the source CRS.
assert 'transform' not in kwargs, 'Transform keyword is not allowed.'
kwargs['transform'] = cartopy_proj
if arg_func is not None:
new_args, kwargs = arg_func(x, y, data, *args, **kwargs)
else:
new_args = (x, y, data) + args
# Draw the contour lines/filled contours.
return draw_method(*new_args, **kwargs)
def mslp(theta, Pi, w, lapse_rate=0.0065, npmsl_height=500.0):
"""Calculate mean sea-level pressure
Args:
theta (iris.cube.Cube): Potential temperature at theta-points
Pi (iris.cube.Cube): Exner pressure at rho-points
w (iris.cube.Cube): Vertical velocity at theta-points, including
surface level
lapse_rate (float): (default is 0.0065)
npmsl_height (float): (default is 500.0)
Returns:
iris.cube.Cube:
"""
# Extract height coordinate of each grid position
z_theta = constants.earth_avg_radius.data + w.coord('altitude').points
z_rho = constants.earth_avg_radius.data + Pi.coord('altitude').points[:-1]
level_height = theta.coord('level_height').points
# Calculate grid variables
grid_lons, grid_lats = cartography.get_xy_grids(theta)
cos_theta_lat = np.cos(grid_lats)
delta_lambda = grid_lons[0, 1] - grid_lons[0, 0]
delta_phi = grid_lats[1, 0] - grid_lats[0, 0]
# Calculate pressure and exner on theta levels
Pi = interpolate.interpolate(Pi,
level_height=w.coord('level_height').points)
P = pressure(Pi)
# Extract surface pressure and remove from other variables
P_star = P[0]
Pi = Pi[1:]
P = P[1:]
# Calculate mean sea-level pressure using hacked MetUM routine
P_msl = fvariable.calc_pmsl(theta.data, Pi.data, P.data, P_star.data,
z_theta, z_rho, cos_theta_lat, level_height,
constants.Cp_d.data, constants.g.data,
constants.Rd.data, lapse_rate,
constants.earth_avg_radius.data, delta_lambda,
delta_phi, npmsl_height)
# Copy data into a new cube
P_msl = P[0].copy(data=P_msl)
P_msl.rename('air_pressure_at_mean_sea_level')
return P_msl
def true_coords(cube):
"""Extract latitude and longitude from a cube with rotated coordinates
Args:
cube (iris.cube.Cube):
Returns:
tuple (numpy.ndarray, numpy.ndarray):
N-dimensional arrays of longitude and latitude
"""
lon, lat = cartography.get_xy_grids(cube)
cs = cube.coord_system()
if cs.grid_mapping_name == 'rotated_latitude_longitude':
lon, lat = cartography.unrotate_pole(lon, lat,
cs.grid_north_pole_longitude,
cs.grid_north_pole_latitude)
return lon, lat
def polar_coords(cube):
"""Get spherical polar coordinates from a cube
Args:
cube (iris.cube.Cube): An iris cube with the relevant coordinate
information. XY grids and the altitude coordinate are required.
Returns:
tuple (numpy.ndarray, numpy.ndarray, numpy.ndarray):
The zonal and meridional co-ordinates in radians and the vertical
co-ordinate in metres
"""
rho = constants.earth_avg_radius.data + cube.coord('altitude').points
x, y = cartography.get_xy_grids(cube)
theta = x * (np.pi / 180)
phi = (90 - y) * (np.pi / 180)
return theta, phi, rho
def __init__(
self,
file_path="",
var_name="",
bathymetry_path="/skynet3_rech1/huziy/NEMO_OFFICIAL/dev_v3_4_STABLE_2012/NEMOGCM/CONFIG/GLK_LIM3_Michigan/EXP00/bathy_meter.nc"
):
"""
:param file_path:
:param var_name:
:param bathymetry_path: used to mask land points
"""
self.current_time_frame = -1
self.var_name = var_name
self.cube = iris.load_cube(
file_path,
constraint=iris.Constraint(
cube_func=lambda c: c.var_name == var_name))
self.lons, self.lats = cartography.get_xy_grids(self.cube)
lons2d_gl, lats2d_gl = nemo_commons.get_2d_lons_lats_from_nemo(
path=bathymetry_path)
mask_gl = nemo_commons.get_mask(path=bathymetry_path)
xs, ys, zs = lat_lon.lon_lat_to_cartesian(lons2d_gl.flatten(),
lats2d_gl.flatten())
xt, yt, zt = lat_lon.lon_lat_to_cartesian(self.lons.flatten(),
self.lats.flatten())
tree = cKDTree(list(zip(xs, ys, zs)))
dists, indices = tree.query(list(zip(xt, yt, zt)))
self.mask = mask_gl.flatten()[indices].reshape(self.lons.shape)
self.nt = self.cube.shape[0]
assert isinstance(self.cube, Cube)
print(self.nt)
def rdf_pv(data_in, data_out, name, pole_lon, pole_lat, domain):
# Load the trajectories
trajectories = trload.raw(data_in)
# Check which trajectories are in the domain
for trajectory in trajectories:
in_domain(trajectory, pole_lon, pole_lat, domain)
# Calculate variable at the start of all trajectories
rlons, rlats, values, cube = rdf(name, trajectories, pole_lon, pole_lat)
# Save the preliminary data
with open(data_out + '.pkl', 'w') as output:
pickle.dump((rlons, rlats, values), output)
# Put the values on the normal grid
lons, lats = get_xy_grids(cube)
field = regular_grid(values, rlons, rlats, lons, lats)
cube = cube[0].copy(data=field)
# Save the field
files.save([cube], data_out + '.nc')
return
def mass_integrals(cubes, lon, lat, theta_level, dtheta, dlambda, dphi):
"""
Args:
cubes (iris.cube.CubeList):
lon (np.Array): Circuit longitudes (degrees)
lat (np.Array): Circuit latitudes (degrees)
theta_level: Isentropic level of the circuit
dtheta (int): Isentrope spacing used to calculate volume integrals
dlambda (float): Longitudinal grid spacing in radians
dphi (float): Latitudinal grid spacing in radians
Returns:
iris.cube.Cube: A tuple containing four `iris.cube.Cube`s. Area, volume, mass
and circulation calculated from the integrals.
"""
# Get height at [theta - dtheta/2, theta, theta + dtheta/2]
levels = ('air_potential_temperature', [
theta_level - dtheta / 2.0, theta_level, theta_level + dtheta / 2.0
])
zth = convert.calc('altitude', cubes, levels=levels)
# Extract model output grid
glon, glat = cartography.get_xy_grids(zth)
gridpoints = np.array([glon.flatten(), glat.flatten()]).transpose()
# Mask all points that are not contained in the circuit
points = np.array([lon, lat]).transpose()
pth = Path(points)
mask = np.logical_not(pth.contains_points(gridpoints).reshape(glat.shape))
# Area = r**2 cos(phi) dlambda dphi
area = (a + zth[1])**2 * np.cos(np.deg2rad(glat)) * dlambda * dphi
area.units = 'm^2'
area.rename('area')
total_area = integrate(area, mask)
# Volume = area * dz
volume = area * (zth[2] - zth[0])
# theta coordinate disappears here
volume.add_aux_coord(area.coord("air_potential_temperature"))
volume.rename('volume')
total_volume = integrate(volume, mask)
# Mass = density * volume
levels = ('air_potential_temperature', [theta_level])
density = convert.calc('air_density', cubes, levels=levels)[0]
mass = density * volume
mass.rename('mass')
total_mass = integrate(mass, mask)
# Circulation = \int rho.pv.dv / dtheta
pv = convert.calc('ertel_potential_vorticity', cubes, levels=levels)[0]
pv.convert_units('m^2 K s-1 kg-1')
pv_substance = pv * mass
circulation = integrate(pv_substance, mask) / dtheta
circulation.rename('mass_integrated_circulation')
return total_area, total_volume, total_mass, circulation