"""Generic calculus functions."""
from cf_units import Unit
from iris.analysis import Linear
from iris.analysis.calculus import (
_coord_cos,
_curl_differentiate,
_curl_regrid,
differentiate,
)
from iris.analysis.maths import add, divide, multiply
from iris.coord_systems import GeogCS, RotatedGeogCS
import numpy as np
from ..const import get_planet_radius
from ..exceptions import NotYetImplementedError
from ..meta import update_metadata
from ..model import um
__all__ = (
"d_dx",
"d_dy",
"d_dz",
"deriv",
"div_h",
"integrate",
)
[docs]
def d_dx(cube, model=um):
"""Calculate a derivative w.r.t. x-coordinate."""
return deriv(cube, model.x)
[docs]
def d_dy(cube, model=um):
"""Calculate a derivative w.r.t. y-coordinate."""
return deriv(cube, model.y)
[docs]
def d_dz(cube, model=um):
"""Calculate a derivative w.r.t. z-coordinate."""
return deriv(cube, model.z)
[docs]
def deriv(cube, coord):
"""
Calculate a derivative w.r.t. the given coordinate.
Uses `iris.analysis.calculus.differentiate` and then
interpolates the result to the grid points of the original cube.
Parameters
----------
cube: iris.cube.Cube
Input cube containing the given coordinate.
coord: str or iris.coords.Coord
Coordinate for differentiation.
Returns
-------
iris.cube.Cube
d(cube)/d(coord).
See Also
--------
aeolus.calc.calculus.d_dx, aeolus.calc.calculus.d_dy
"""
pnts = cube.coord(coord).points
diff = differentiate(cube, coord)
res = diff.interpolate([(coord, pnts)], Linear())
return res
[docs]
@update_metadata(name="horizontal_divergence")
def div_h(i_cube, j_cube, r_planet=None, model=um):
r"""
Calculate horizontal divergence.
Note: currently works only in spherical coordinates.
Parameters
----------
i_cube:
i-component (e.g. u-wind)
j_cube:
j-component (e.g. v-wind)
r_planet: float, optional
Radius of the planet (m). If not given, an attempt is made
to get it from the cube metadata.
model: aeolus.model.Model, optional
Model class with relevant variable names.
Returns
-------
iris.cube.Cube
Cube of horizontal divergence.
Notes
-----
Divergence in spherical coordinates is defined as
.. math::
\nabla\cdot \vec A = \frac{1}{r cos \theta} (
\frac{\partial A_\lambda}{\partial \lambda}
+ \frac{\partial}{\partial \theta}
(A_\theta cos \theta))
where \lambda is longitude, \theta is latitude.
"""
x_coord = i_cube.coord(model.x)
y_coord = i_cube.coord(model.y)
y_dim = i_cube.coord_dims(y_coord)[0]
horiz_cs = i_cube.coord_system("CoordSystem")
# Check for spherical coords
spherical_coords = isinstance(horiz_cs, (GeogCS, RotatedGeogCS))
if spherical_coords:
# Get the radius of the planet
if r_planet is None:
r = get_planet_radius(i_cube)
else:
r = r_planet
r_unit = Unit("m")
lon_coord = x_coord.copy()
lat_coord = y_coord.copy()
lon_coord.convert_units("radians")
lat_coord.convert_units("radians")
lat_cos_coord = _coord_cos(lat_coord)
# j-comp: \frac{\partial}{\partial \theta} (\vec A_\theta cos \theta))
temp = multiply(j_cube, lat_cos_coord, y_dim)
djcos_dtheta = _curl_differentiate(temp, lat_coord)
prototype_diff = djcos_dtheta
# i-comp: \frac{\partial \vec A_\lambda}{\partial \lambda}
d_i_cube_dlambda = _curl_differentiate(i_cube, lon_coord)
d_i_cube_dlambda = _curl_regrid(d_i_cube_dlambda, prototype_diff)
new_lat_coord = d_i_cube_dlambda.coord(model.y)
new_lat_cos_coord = _coord_cos(new_lat_coord)
lat_dim = d_i_cube_dlambda.coord_dims(new_lat_coord)[0]
# Sum and divide
div = divide(
add(d_i_cube_dlambda, djcos_dtheta),
r * new_lat_cos_coord,
dim=lat_dim,
)
div.units /= r_unit
div = div.regrid(i_cube, Linear())
else:
raise NotYetImplementedError(
"Non-spherical coordinates are not implemented yet."
)
return div
[docs]
def integrate(cube, coord):
"""
Integrate the cube along a 1D coordinate using the trapezoidal rule.
Note: `coord` must be one of the dimensional coordinates of the cube.
Parameters
----------
cube: iris.cube.Cube
Input cube containing the given coordinate.
coord: str or iris.coords.Coord
Coordinate for integration.
Returns
-------
iris.cube.Cube
Integrated cube.
"""
# TODO: allow non-dim coordinates
c = cube.coord(coord)
others = [
dc.name()
for dc in cube.dim_coords
if cube.coord_dims(dc) != cube.coord_dims(c)
]
dim = cube.coord_dims(c)[0]
data = np.trapz(cube.data, c.points, axis=dim)
res = next(cube.slices(others)).copy(data=data)
res.units = cube.units * c.units
res.remove_coord(c)
res.rename(f"integral_of_{cube.name()}_wrt_{c.name()}")
# ensure_bounds(cube, [c])
# delta = AuxCoord(c.bounds[:, 1] - c.bounds[:, 0], units=c.units)
# res = multiply(cube, delta, dim=dim).collapsed(c, iris.analysis.SUM)
return res