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

# Source code for aeolus.calc.diag

    """Some commonly used diagnostics in atmospheric science."""
    from cf_units import Unit

    from iris.analysis.calculus import _coord_cos
    from iris.analysis.maths import add, apply_ufunc, multiply
    from iris.exceptions import ConstraintMismatchError as ConMisErr
    from iris.util import reverse

    import numpy as np

    from .calculus import d_dz, integrate
    from .meta import const_from_attrs, update_metadata
    from .stats import cumsum, spatial_mean, time_mean, zonal_mean
    from ..const import init_const
    from ..coord import coord_to_cube, ensure_bounds, regrid_3d
    from ..exceptions import ArgumentError, MissingCubeError
    from ..model import um
    from ..subset import _dim_constr

    __all__ = (
        "air_density",
        "air_potential_temperature",
        "air_temperature",
        "bond_albedo",
        "dry_lapse_rate",
        "flux",
        "geopotential_height",
        "ghe_norm",
        "heat_redist_eff",
        "horiz_wind_cmpnts",
        "meridional_mass_streamfunction",
        "precip_sum",
        "sfc_net_energy",
        "sfc_water_balance",
        "sigma_p",
        "toa_cloud_radiative_effect",
        "toa_eff_temp",
        "toa_net_energy",
        "water_path",
        "wind_speed",
        "zonal_mass_streamfunction",
    )

    def _precip_name_mapping(model=um):
        """Generate lists of variable names for `precip_sum()`."""
        return {
            "total": [model.ls_rain, model.ls_snow, model.cv_rain, model.cv_snow],
            "conv": [model.cv_rain, model.cv_snow],
            "stra": [model.ls_rain, model.ls_snow],
            "rain": [model.ls_rain, model.cv_rain],
            "snow": [model.ls_snow, model.cv_snow],
        }

    @update_metadata(name="cos_lat", units="1")
    def lat_cos(cube, model=um):
        """Convert the latitude coordinate to a cube and apply cosine to it."""
        lat_cube = coord_to_cube(cube, model.y, broadcast=False)
        lat_cos_cube = apply_ufunc(np.cos, apply_ufunc(np.deg2rad, lat_cube))
        return lat_cos_cube

    @update_metadata(units="K")
    def air_temperature(cubelist, const=None, model=um):
        """
        Get the real temperature from the given cube list.

        If not present, it is attempted to calculate it from other variables,
        Exner pressure or air pressure, and potential temperature.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes containing temperature.
        const: aeolus.const.const.ConstContainer, optional
            Must have `reference_surface_pressure` and `dry_air_gas_constant` as attributes.
            If not given, an attempt is made to retrieve it from cube attributes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of air temperature.
        """
        try:
            return cubelist.extract_cube(model.temp)
        except ConMisErr:
            try:
                thta = cubelist.extract_cube(model.thta)
            except ConMisErr:
                raise MissingCubeError(f"Unable to get air temperature from {cubelist}")

            if len(cubelist.extract(model.exner)) == 1:
                exner = cubelist.extract_cube(model.exner)
            elif len(cubelist.extract(model.pres)) == 1:
                if const is None:
                    const = thta.attributes["planet_conf"]
                pres = cubelist.extract_cube(model.pres)
                exner = (pres / const.reference_surface_pressure) ** (
                    const.dry_air_gas_constant / const.dry_air_spec_heat_press
                ).data
            else:
                raise MissingCubeError(f"Unable to get air temperature from {cubelist}")
            temp = thta * exner
            temp.rename(model.temp)
            return temp

    @update_metadata(units="K")
    def air_potential_temperature(cubelist, const=None, model=um):
        """
        Get the potential temperature from the given cube list.

        If not present, it is attempted to calculate it from other variables,
        Exner pressure or air pressure, and real temperature.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes containing temperature.
        const: aeolus.const.const.ConstContainer, optional
            Must have `reference_surface_pressure` and `dry_air_gas_constant` as attributes.
            If not given, an attempt is made to retrieve it from cube attributes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of air potential temperature.
        """
        try:
            return cubelist.extract_cube(model.thta)
        except ConMisErr:
            try:
                temp = cubelist.extract_cube(model.temp)
            except ConMisErr:
                raise MissingCubeError(f"Unable to get air potential temperature from {cubelist}")

            if len(cubelist.extract(model.exner)) == 1:
                exner = cubelist.extract_cube(model.exner)
            elif len(cubelist.extract(model.pres)) == 1:
                if const is None:
                    const = temp.attributes["planet_conf"]
                pres = cubelist.extract_cube(model.pres)
                exner = (pres / const.reference_surface_pressure) ** (
                    const.dry_air_gas_constant / const.dry_air_spec_heat_press
                ).data
            else:
                raise MissingCubeError(f"Unable to get air potential temperature from {cubelist}")
            thta = temp / exner
            thta.rename(model.thta)
            thta.convert_units("K")
            return thta

    @update_metadata(units="kg m-3")
    def air_density(cubelist, const=None, model=um):
        """
        Get air density from the given cube list.

        If not present, it is attempted to calculate it from pressure and temperature.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes containing temperature.
        const: aeolus.const.const.ConstContainer, optional
            Must have `dry_air_gas_constant` as an attribute.
            If not given, an attempt is made to retrieve it from cube attributes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of air density.
        """
        try:
            return cubelist.extract_cube(model.dens)
        except ConMisErr:
            try:
                temp = cubelist.extract_cube(model.temp)
                pres = cubelist.extract_cube(model.pres)
                if const is None:
                    const = pres.attributes["planet_conf"]
                rho = pres / (const.dry_air_gas_constant * temp)
                rho.rename(model.dens)
                return rho
            except ConMisErr:
                _msg = f"Unable to get variables from\n{cubelist}\nto calculate air density"
                raise MissingCubeError(_msg)

    @update_metadata(units="m2 s-2")
    def geopotential_height(cubelist, const=None, model=um):
        """
        Get geopotential height from the given cube list.

        If not present, the altitude coordinate is transformed into a cube.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes containing temperature.
        const: aeolus.const.const.ConstContainer, optional
            Must have `gravity` as an attribute.
            If not given, an attempt is made to retrieve it from cube attributes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of geopotential height.
        """
        try:
            return cubelist.extract_cube(model.ghgt)
        except ConMisErr:
            try:
                cube_w_height = cubelist.extract(_dim_constr(model.z, strict=False))[0]
                if const is None:
                    const = cube_w_height.attributes["planet_conf"]
                g_hgt = coord_to_cube(cube_w_height, model.z) * const.gravity
                g_hgt.attributes = {k: v for k, v in cube_w_height.attributes.items() if k != "STASH"}
                ensure_bounds(g_hgt, [model.z])
                g_hgt.rename(model.ghgt)
                return g_hgt
            except ConMisErr:
                _msg = f"No cubes in \n{cubelist}\nwith {model.z} as a coordinate."
                raise MissingCubeError(_msg)

    def flux(cubelist, quantity, axis, weight_by_density=True, model=um):
        r"""
        Calculate horizontal or vertical flux of some quantity.

        .. math::
            F_{x} = u (\rho q),
            F_{y} = v (\rho q),
            F_{z} = w (\rho q)

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes.
        quantity: str or iris.Constraint
            Quantity (present in the cube list).
        axis: str
            Axis of the flux component (x|y|z)
        weight_by_density: bool, optional
            Multiply by a cube of air density (must be present in the input cube list).
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of a flux component with the same dimensions as input cubes.
        """
        if axis.lower() == "x":
            u = cubelist.extract_cube(model.u)
        elif axis.lower() == "y":
            u = cubelist.extract_cube(model.v)
        elif axis.lower() == "z":
            u = cubelist.extract_cube(model.w)
        q = cubelist.extract_cube(quantity)
        fl = u * q
        if weight_by_density:
            fl *= cubelist.extract_cube(model.dens)
        fl.rename(f"flux_of_{quantity}_along_{axis.lower()}_axis")
        return fl

    @update_metadata(units="W m-2")
    def toa_cloud_radiative_effect(cubelist, kind, model=um):
        r"""
        Calculate domain-average TOA cloud radiative effect (CRE).

        .. math::
            CRE_{TOA} = F_{up,clear-sky} - F_{up,all-sky}

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes
        kind: str
            Shortwave ('sw'), longwave ('lw'), or 'total' CRE

        Returns
        -------
        iris.cube.Cube
            Cube of CRE.
        """
        name = f"toa_cloud_radiative_effect_{kind}"
        if kind == "sw":
            all_sky = model.toa_osr
            clr_sky = model.toa_osr_cs
        elif kind == "lw":
            all_sky = model.toa_olr
            clr_sky = model.toa_olr_cs
        elif kind == "total":
            sw = toa_cloud_radiative_effect(cubelist, "sw")
            lw = toa_cloud_radiative_effect(cubelist, "lw")
            cre = sw + lw
            cre.rename(name)
            return cre

        cube_clr = cubelist.extract_cube(clr_sky)
        cube_all = cubelist.extract_cube(all_sky)
        cre = cube_clr - cube_all
        cre.rename(name)
        return cre

    @update_metadata(name="toa_net_downward_energy_flux", units="W m-2")
    def toa_net_energy(cubelist, model=um):
        """
        Calculate domain-average TOA energy flux.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of total TOA downward energy flux.
        """
        varnames = [
            model.toa_isr,
            model.toa_osr,
            model.toa_olr,
        ]
        terms = cubelist.extract(varnames)
        if len(terms) != 3:
            raise MissingCubeError(
                f"{varnames} required for TOA energy balance are missing from cubelist:\n{cubelist}"
            )
        terms_ave = []
        for cube in terms:
            terms_ave.append(cube)
        toa_net = terms_ave[0] - terms_ave[1] - terms_ave[2]
        return toa_net

    @update_metadata(name="surface_net_downward_energy_flux", units="W m-2")
    def sfc_net_energy(cubelist, model=um):
        """
        Calculate domain-average surface energy flux.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes with net LW and SW radiation, sensible and latent surface fluxes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of total surface downward energy flux.
        """
        net_down_lw = cubelist.extract_cube(model.sfc_net_down_lw)
        net_down_sw = cubelist.extract_cube(model.sfc_net_down_sw)
        shf = cubelist.extract_cube(model.sfc_shf)
        lhf = cubelist.extract_cube(model.sfc_lhf)
        sfc_net = net_down_lw + net_down_sw - shf - lhf
        return sfc_net

    @const_from_attrs
    @update_metadata(name="surface_net_downward_water_flux", units="mm h-1")
    def sfc_water_balance(cubelist, const=None, model=um):
        """
        Calculate domain-average precipitation minus evaporation.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes.
        const: aeolus.const.const.ConstContainer, optional
            Must have a scalar cube of `condensible_density` as an attribute.
            If not given, attempt is made to retrieve it from cube attributes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of total surface downward water flux (P-E).
        """
        try:
            evap = cubelist.extract_cube(model.sfc_evap)
        except ConMisErr:
            try:
                lhf = cubelist.extract_cube(model.sfc_lhf)
                evap = lhf / const.condensible_heat_vaporization
                evap /= const.condensible_density
            except (KeyError, ConMisErr):
                raise MissingCubeError(f"Cannot retrieve evaporation from\n{cubelist}")
        try:
            precip = cubelist.extract_cube(model.ppn)
            precip /= const.condensible_density
        except ConMisErr:
            precip = precip_sum(cubelist, ptype="total", const=const, model=model)
        precip.convert_units("mm h-1")
        evap.convert_units("mm h-1")
        net = precip - evap
        return net

    @const_from_attrs
    def precip_sum(cubelist, ptype="total", const=None, model=um):
        """
        Calculate a sum of different types of precipitation [:math:`mm~day^{-1}`].

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes.
        ptype: str, optional
            Precipitation type (total|stra|conv|rain|snow).
        const: aeolus.const.const.ConstContainer, optional
            Must have a scalar cube of `condensible_density` as an attribute.
            If not given, attempt to retrieve it from cube attributes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Sum of cubes of precipitation with units converted to mm per day.
        """
        try:
            varnames = _precip_name_mapping(model=model)[ptype]
        except KeyError:
            raise ArgumentError(f"Unknown ptype={ptype}")
        if len(cubelist.extract(varnames)) == 0:
            raise MissingCubeError(
                f"{varnames} required for ptype={ptype} are missing from cubelist:\n{cubelist}"
            )
        precip = 0.0
        for varname in varnames:
            try:
                cube = cubelist.extract_cube(varname)
                precip += cube
            except ConMisErr:
                pass
        precip /= const.condensible_density
        precip.convert_units("mm day^-1")
        precip.rename(f"{ptype}_precip_rate")
        return precip

    @update_metadata(name="heat_redistribution_efficiency", units="1")
    def heat_redist_eff(cubelist, region_a, region_b, model=um):
        r"""
        Heat redistribution efficiency (Leconte et al. 2013).

        .. math::
            \eta=\frac{OLR_{TOA,night}}{OLR_{TOA,day}}

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes.
        region_a: aeolus.region.Region
            First region (usually nightside).
        region_b: aeolus.region.Region
            Second region (usually dayside).
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of eta parameter with collapsed spatial dimensions.
        """
        toa_olr = cubelist.extract_cube(model.toa_olr)
        toa_olr_a = spatial_mean(toa_olr.extract(region_a.constraint))
        toa_olr_b = spatial_mean(toa_olr.extract(region_b.constraint))
        eta = toa_olr_a / toa_olr_b
        return eta

    @update_metadata(name="toa_effective_temperature", units="K")
    def toa_eff_temp(cubelist, model=um):
        r"""
        Calculate effective temperature from TOA OLR.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of :math:`T_{eff}`.
        """
        toa_olr = cubelist.extract_cube(model.toa_olr)
        sbc = init_const().stefan_boltzmann
        t_eff = (toa_olr / sbc) ** 0.25
        return t_eff

    @update_metadata(name="normalised_greenhouse_effect_parameter", units="1")
    def ghe_norm(cubelist, model=um):
        r"""
        Normalised greenhouse effect parameter.

        .. math::
            GHE = 1 - \left(\frac{T_{eff}}{T_{sfc}}\right)^4

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of greenhouse effect parameter.
        """
        t_sfc = cubelist.extract_cube(model.t_sfc)
        t_eff = toa_eff_temp(cubelist, model=model)
        out = (t_eff / t_sfc) ** 4
        out = out.copy(data=1 - out.data)
        return out

    @const_from_attrs
    @update_metadata(name="bond_albedo", units="1")
    def bond_albedo(cubelist, const=None, model=um):
        r"""
        Bold albedo.

        .. math::
            4 \frac{OSR_{TOA}}{S_{0}}

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes.
        const: aeolus.const.const.ConstContainer, optional
            Must have a scalar cube of `solar_constant` as an attribute.
            If not given, attempt to retrieve it from cube attributes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of bond albedo.
        """
        toa_osr = cubelist.extract_cube(model.toa_osr)
        sc = const.solar_constant
        b_alb = 4 * toa_osr / sc
        return b_alb

    @update_metadata(units="kg m-2")
    def water_path(cubelist, kind="water_vapour", model=um):
        r"""
        Water vapour or condensate path, i.e. a vertical integral of a water phase.

        .. math::
            WP = \int_{z_{sfc}}^{z_{top}} \rho q dz

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes containing appropriate mixing ratio and air density.
        kind: str, optional
            Short name of the water phase to be integrated.
            Options are water_vapour (default) | liquid_water | ice_water | cloud_water
            `cloud_water` is the sum of liquid and ice phases.
        model: aeolus.model.Model, optional
            Model class with relevant coordinate names.
            `model.z` is used as a vertical coordinate for integration.

        Returns
        -------
        iris.cube.Cube
            Cube of water path with collapsed vertical dimension.
        """
        if kind == "water_vapour":
            q = cubelist.extract_cube(model.sh)
        elif kind == "liquid_water":
            q = cubelist.extract_cube(model.cld_liq_mf)
        elif kind == "ice_water":
            q = cubelist.extract_cube(model.cld_ice_mf)
        elif kind == "cloud_water":
            q = cubelist.extract_cube(model.cld_liq_mf).copy()
            q += cubelist.extract_cube(model.cld_ice_mf)
        rho = cubelist.extract_cube(model.dens)
        wp = integrate(q * rho, model.z)
        wp.rename(f"{kind}_path")
        return wp

    def dry_lapse_rate(cubelist, model=um):
        r"""
        Dry lapse rate, or the change of air temperature with altitude.

        .. math::
            \gamma = \partial T_{air} / \partial z

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes containing temperature.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube of dT/dz.
        """
        temp = cubelist.extract_cube(model.temp)
        return d_dz(temp, model=model)

    def horiz_wind_cmpnts(cubelist, model=um):
        """
        Extract u and v wind components from a cube list and interpolate v on u's grid if necessary.

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            List of cubes with horizontal wind components
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        u, v: iris.cube.Cube
            Cubes of wind components.
        """
        # use relaxed extracting
        u = cubelist.extract(model.u)[0]
        v = cubelist.extract(model.v)[0]
        # interpolate v on u's grid if coordinates are different
        v = regrid_3d(v, u, model=model)
        return u, v

    @const_from_attrs
    @update_metadata(name="local_superrotation_index", units="1")
    def superrotation_index(cubelist, const=None, model=um):
        r"""
        Local superrotation index.

        .. math::
            s = \frac{m}{\Omega a^2} - 1,

            m = a cos\phi(\Omega a cos\phi + u)

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            List of cubes containing a cube of zonal velocity (u).
        const: aeolus.const.const.ConstContainer, optional
            Constainer with the relevant planetary constants.
            If not given, attempt is made to retrieve it from cube attributes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        s_idx: iris.cube.Cube
            Cubes of superrotation index.

        References
        ----------
        Read (1986), https://doi.org/10.1002/qj.49711247114
        """
        # Zonal velocity
        u = cubelist.extract_cube(model.u)
        # Radius of the planet
        r = const.radius
        r.convert_units("m")
        # Rotation rate
        omega = const.planet_rotation_rate

        lat_coord = u.coord(model.y).copy()
        lat_dim = u.coord_dims(model.y)[0]
        lat_coord.convert_units("radians")
        lat_cos_coord = _coord_cos(lat_coord)

        # Calculate intermediate terms
        omega_r_coslat = omega.data * r.data * lat_cos_coord
        omega_r_coslat.units = Unit("m s-1")
        r_coslat = r.data * lat_cos_coord
        r_coslat.units = Unit("m")
        inner_sum = add(u, omega_r_coslat, dim=lat_dim)

        # Calculate axial component of specific absolute angular momentum
        ang_mom = multiply(inner_sum, r_coslat, dim=lat_dim)

        # Final index
        s_idx = ang_mom / omega / (r ** 2)
        s_idx.convert_units("1")
        s_idx = s_idx.copy(data=s_idx.data - 1)
        return s_idx

    @update_metadata(name="wind_speed", units="m s-1")
    def wind_speed(*components):
        r"""
        Calculate the wind speed (magnitude of the wind vector).

        Parameters
        ----------
        args: iris.cube.Cube
            Cubes of u, v, w wind components.

        .. math::
            \sqrt{u^2 + v^2 + w^2}
        """
        out = sum(cube ** 2 for cube in components) ** 0.5
        return out

    @const_from_attrs
    @update_metadata(name="atmosphere_hybrid_sigma_pressure_coordinate", units="1")
    def sigma_p(cubelist, const=None, model=um):
        r"""
        Calculate sigma (normalised pressure coordinate) from a cube of pressure.

        .. math::
            \sigma = p / p_{sfc}

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input list of cubes.
        const: aeolus.const.const.ConstContainer, optional
            Must have a cube of reference (surface) pressure as an attribute.
            If not given, attempt to retrieve it from cube attributes.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.
        """
        pres_cube = time_mean(spatial_mean(cubelist.extract_cube(model.pres)))
        pres_cube.convert_units("Pa")
        return pres_cube / const.reference_surface_pressure

    @const_from_attrs
    @update_metadata(name="zonal_mass_streamfunction", units="kg s^-1")
    def zonal_mass_streamfunction(cubelist, const=None, model=um):
        r"""
        Calculate mean zonal mass streamfunction.

        .. math::
            \Psi_Z = 2\pi a \int_{z_{sfc}}^{z_{top}}\overline{\rho}^* \overline{u}^* dz

        References
        ----------
        Haqq-Misra & Kopparapu (2015), eq. 5;
        Hartmann (1994), Global Physical Climatology, eq. 6.21

        Examples
        --------
        >>> lat_band_constr = iris.Constraint(
            latitude=lambda x: -30 <= x.point <= 30, longitude=lambda x: True
        )
        >>> mzsf = zonal_mass_streamfunction(cubes.extract(lat_band_constr))
        """
        streamf_const = 2 * np.pi * const.radius

        u = cubelist.extract_cube(model.u)
        u_tm = time_mean(u, model=model)
        u = -1 * (u_tm - zonal_mean(u_tm, model=model))

        if u.coord(axis="z").units.is_convertible("m"):
            rho = cubelist.extract_cube(model.dens).copy()
            rho = time_mean(rho, model=model)
            rho.coord(model.z).bounds = None
            u.coord(model.z).bounds = None
            integrand = u * rho
            # integrand = reverse(integrand, model.z)
            # print(integrand.coord(um.z))
            res = cumsum(integrand, "z", axis_weights=True, model=model)
            # res = cumsum(integrand, "z", axis_weights=True, model=model)
            # res = reverse(res, model.z)

        elif u.coord(axis="z").units.is_convertible("Pa"):
            integrand = u
            res = cumsum(integrand, "z", axis_weights=True, model=model)
            res /= const.gravity

        res *= streamf_const
        return res

    @const_from_attrs
    @update_metadata(name="meridional_mass_streamfunction", units="kg s^-1")
    def meridional_mass_streamfunction(cubelist, const=None, model=um):
        r"""
        Calculate the mean meridional mass streamfunction.

        * In height coordinates

        .. math::
            \Psi_M = - 2\pi cos\phi a \int_{z_{sfc}}^{z_{top}}\overline{\rho v} dz

        * In pressure coordinates

        .. math::
            \Psi_M = 2\pi cos\phi a \int_{0}^{p_{sfc}}\overline{\rho v} dp / g

        Parameters
        ----------
        cubelist: iris.cube.CubeList
            Input cubelist.
        const: aeolus.const.const.ConstContainer, optional
            If not given, constants are attempted to be retrieved from
            attributes of a cube in the cube list.
        model: aeolus.model.Model, optional
            Model class with relevant variable names.

        Returns
        -------
        iris.cube.Cube
            Cube with collapsed spatial dimensions.

        References
        ----------
        Haqq-Misra & Kopparapu (2015), eq. 4;
        Vallis (2017)

        Examples
        --------
        >>> from aeolus.calc import meridional_mass_streamfunction, time_mean
        >>> from aeolus.const import init_const
        >>> from aeolus.model import um
        >>> earth_constants = init_const("earth")
        >>> cubes = iris.cube.CubeList([time_mean(cube) for cube in input_cubes])
        >>> mmsf = meridional_mass_streamfunction(cubes, const=earth_constants, model=um)
        """
        v = cubelist.extract_cube(model.v)
        v = zonal_mean(v, model=model)
        if v.coord(model.z).units.is_convertible("m"):
            rho = zonal_mean(cubelist.extract_cube(model.dens), model=model)
            rho.coord(model.z).bounds = None
            v.coord(model.z).bounds = None
            # Reverse the coordinate to start from the model top (where p=0 or z=z_top)
            # TODO: check if the coordinate is ascending or descending
            integrand = reverse(v * rho, model.z)
            res = -1 * cumsum(integrand, "z", axis_weights=True, model=model)
            # Reverse the result back
            res = reverse(res, model.z)
        elif v.coord(model.z).units.is_convertible("Pa"):
            res = cumsum(v, "z", axis_weights=True, model=model)
            res /= const.gravity
        # Calculate the constant: 2 pi cos\phi a
        streamf_const = 2 * np.pi * const.radius * lat_cos(res, model=model)
        res *= streamf_const
        return res