# -*- coding: utf-8 -*-
"""Miscellaneous."""
import warnings
import iris
import numpy as np
from .const import get_planet_radius
from .coord_utils import UM_LATLON, nearest_coord_value
from .exceptions import AeolusWarning
__all__ = (
"vertical_cross_section_area",
"horizontal_fluxes_through_region_boundaries",
"net_horizontal_flux_to_region",
)
[docs]def vertical_cross_section_area(cube2d, r_planet=None):
"""Create a cube of vertical cross-section areas in metres."""
cube2d = cube2d.copy()
if r_planet is None:
r = get_planet_radius(cube2d)
else:
r = r_planet
m_per_deg = (np.pi / 180) * r
if iris.util.guess_coord_axis(cube2d.dim_coords[1]) == "X":
m_per_deg *= np.cos(np.deg2rad(cube2d.coord(axis="Y").points[0]))
for dim_coord in cube2d.dim_coords:
if not dim_coord.has_bounds():
dim_coord.guess_bounds()
x_bounds = cube2d.coord(cube2d.dim_coords[1]).bounds
z_bounds = cube2d.coord(cube2d.dim_coords[0]).bounds
vc_area = cube2d.copy(
data=(z_bounds[:, 1] - z_bounds[:, 0])[:, None]
* ((x_bounds[:, 1] - x_bounds[:, 0])[None, :] * m_per_deg)
)
vc_area.units = "m**2"
vc_area.rename("vertical_section_area")
for dim_coord in vc_area.dim_coords:
dim_coord.bounds = None
return vc_area
def horizontal_fluxes_through_region_boundaries(
scalar_cube, region, u, v, r_planet=None, vertical_constraint=None, warn_thresh=10
):
"""Calculate horizontal fluxes of `scalar_cube` through planes of a rectangular region."""
perpendicular_wind_cmpnts = {UM_LATLON[1]: u, UM_LATLON[0]: v}
if r_planet is None:
r = get_planet_radius(scalar_cube)
else:
r = r_planet
total_h_fluxes = iris.cube.CubeList()
for bound in region:
this_coord = bound["coord"]
other_coord, (other_min, other_max) = region._perpendicular_side_limits(bound["name"])
nearest = nearest_coord_value(scalar_cube, this_coord, bound["value"])
if abs(nearest - bound["value"]) >= warn_thresh:
warnings.warn(
f"Nearest value is {np.round(nearest - bound['value'], 2)} deg away"
f" from the given value of {this_coord}",
AeolusWarning,
)
vcross_cnstr = iris.Constraint(**{this_coord: nearest})
vcross_cnstr &= vertical_constraint
if other_max >= other_min:
vcross_cnstr &= iris.Constraint(**{other_coord: lambda x: other_min <= x <= other_max})
cube = scalar_cube.extract(vcross_cnstr)
else:
vcross_cnstr &= iris.Constraint(
**{other_coord: lambda x: (other_max >= x) or (other_min <= x)}
)
cube = scalar_cube.extract(vcross_cnstr)
cube_slice = next(cube.slices([cube.coord(axis="z").name(), other_coord]))
vcross_area = vertical_cross_section_area(cube_slice, r_planet=r)
# Calculate energy flux (2d)
cube = perpendicular_wind_cmpnts[this_coord].extract(vcross_cnstr) * cube * vcross_area
cube.rename(f"{scalar_cube.name()}_flux_through_{bound['name']}_boundary")
# Total flux
collapsible_dims = [
i for i in cube.dim_coords if iris.util.guess_coord_axis(i) in ["Z", "Y", "X"]
]
cube_total = cube.collapsed(collapsible_dims, iris.analysis.SUM)
total_h_fluxes.append(cube_total)
return total_h_fluxes
[docs]def net_horizontal_flux_to_region(
scalar_cube, region, u, v, r_planet=None, vertical_constraint=None
):
"""Calculate horizontal fluxes of `scalar_cube` quantity and add them to get the net result."""
total_h_fluxes = horizontal_fluxes_through_region_boundaries(
scalar_cube, region, u, v, r_planet=r_planet, vertical_constraint=vertical_constraint
)
net_flux = (
total_h_fluxes.extract_strict(
iris.Constraint(cube_func=lambda x: "through_west" in x.name())
)
- total_h_fluxes.extract_strict(
iris.Constraint(cube_func=lambda x: "through_east" in x.name())
)
+ total_h_fluxes.extract_strict(
iris.Constraint(cube_func=lambda x: "through_south" in x.name())
)
- total_h_fluxes.extract_strict(
iris.Constraint(cube_func=lambda x: "through_north" in x.name())
)
)
net_flux.rename(f"net_{scalar_cube.name()}_horizontal_flux_to_region")
net_flux.attributes["region_str"] = str(region)
return net_flux