"""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
[docs]@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
[docs]@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
[docs]@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)
[docs]@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)
[docs]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
[docs]@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
[docs]@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
[docs]@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
[docs]@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
[docs]@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
[docs]@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
[docs]@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
[docs]@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
[docs]@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
[docs]@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
[docs]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)
[docs]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
[docs]@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
[docs]@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
[docs]@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
[docs]@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