aeolus.calc.flux_h — aeolus 0.4.9+9.g5a71939 documentation

# Source code for aeolus.calc.flux_h

    # -*- 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

    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