# -*- coding: utf-8 -*-
"""Integrated fluxes."""
import iris
import numpy as np
from ..const import get_planet_radius
from ..coord import nearest_coord_value, vertical_cross_section_area
from ..exceptions import _warn
from ..model import um
__all__ = ("horizontal_fluxes_through_region_boundaries", "net_horizontal_flux_to_region")
def horizontal_fluxes_through_region_boundaries(
scalar_cube, region, u, v, r_planet=None, vertical_constraint=None, warn_thresh=10, model=um
):
"""Calculate horizontal fluxes of `scalar_cube` through planes of a rectangular region."""
perpendicular_wind_cmpnts = {um.x: u, um.y: 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:
_warn(
f"Nearest value is {np.round(nearest - bound['value'], 2)} deg away"
f" from the given value of {this_coord}",
)
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 side 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, model=um
):
"""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,
model=model,
)
net_flux = (
total_h_fluxes.extract_cube(iris.Constraint(cube_func=lambda x: "through_west" in x.name()))
- total_h_fluxes.extract_cube(
iris.Constraint(cube_func=lambda x: "through_east" in x.name())
)
+ total_h_fluxes.extract_cube(
iris.Constraint(cube_func=lambda x: "through_south" in x.name())
)
- total_h_fluxes.extract_cube(
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