diff --git a/docs/_templates/overrides/metpy.calc.rst b/docs/_templates/overrides/metpy.calc.rst index e5e7a4f787f..a3be725bc3e 100644 --- a/docs/_templates/overrides/metpy.calc.rst +++ b/docs/_templates/overrides/metpy.calc.rst @@ -164,6 +164,7 @@ Mathematical Functions tangential_component unit_vectors_from_cross_section vector_derivative + cumulative_integrate Apparent Temperature diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 091545d8856..323fb91eb9a 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -17,6 +17,7 @@ import numpy.ma as ma from pyproj import CRS, Geod, Proj +from scipy.integrate import cumulative_trapezoid from scipy.spatial import cKDTree import xarray as xr @@ -1942,3 +1943,44 @@ def _remove_nans(*variables): for v in variables: ret.append(v[~mask]) return ret + + +@exporter.export +@xarray_derivative_wrap +def cumulative_integrate(field, axis=None, x=None, delta=None): + """Return cumulative integral of field along axis. + + Uses trapezoidal rule for integration. + + Parameters + ---------- + field : array-like + Array of values for which to calculate the integral + axis : int or str + The axis along which to integrate. If `field` is an + `np.ndarray` or `pint.Quantity`, must be an integer. Defaults + to zero. + x : array-like, optional + The coordinate values along which to integrate + delta : array-like, optional + Spacing between grid points in `field`. + + Examples + -------- + >>> cumulative_integrate(units.Quantity(np.arange(5), 'm'), delta=units('1 m')) + + """ + n, axis, delta = _process_deriv_args(field, axis, x, delta) + take = make_take(n, axis) + right = np.cumsum(delta, axis=axis) + left = np.zeros_like(right[take(slice(1))]) + x = concatenate([left, right], axis=axis) + try: + result = cumulative_trapezoid(field.magnitude, x=x.magnitude, axis=axis, initial=0) + return units.Quantity(result, field.units * delta.units) + except AttributeError: + # NumPy arrays without units + raise TypeError( + 'cumulative_integrate called with unitless arguments\n' + 'Either add units or use scipy.integrate.cumulative_trapezoid' + ) from None diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 6f92f03d385..df640d86e2c 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -12,7 +12,8 @@ import pytest import xarray as xr -from metpy.calc import (angle_to_direction, find_bounding_indices, find_intersections, +from metpy.calc import (angle_to_direction, cumulative_integrate, + find_bounding_indices, find_intersections, first_derivative, geospatial_gradient, get_layer, get_layer_heights, gradient, laplacian, lat_lon_grid_deltas, nearest_intersection_idx, parse_angle, pressure_to_height_std, reduce_point_density, @@ -1557,3 +1558,53 @@ def test_vector_derivative_return_subset(return_only, length): u, v, longitude=lons, latitude=lats, crs=crs, return_only=return_only) assert len(ddx) == length + + +def test_cumulative_integrate_numpy(): + """Test that cumulative_integrate works with numpy arrays.""" + field = np.arange(5) + with pytest.raises( + TypeError, match='cumulative_integrate called with unitless arguments' + ): + cumulative_integrate(field, delta=1) + + +def test_cumulative_integrate_pint(): + """Test that cumulative_integrate works with pint Quantities.""" + field = np.arange(6) * units('kg/m^3') + delta = np.array([1, 2, 3, 2, 1]) * units('cm') + integral = cumulative_integrate(field, delta=delta) + assert integral.magnitude == pytest.approx( + np.array([0, 0.5, 3.5, 11, 18, 22.5]) + ) + assert units.Quantity(1, integral.units).to('dag/m^2').magnitude == 1 + + +def test_cumulative_integrate_xarray(): + """Test that cumulative_integrate works with XArray DataArrays.""" + field = xr.DataArray( + np.arange(10) / 100, + {'x': (('x',), np.arange(100, 1001, 100), {'units': 'hPa'})}, + attrs={'units': 'g/kg'} + ) + integral = cumulative_integrate(field, axis='x') + assert integral.metpy.magnitude == pytest.approx( + np.array([0, 0.5, 2, 4.5, 8, 12.5, 18, 24.5, 32, 40.5]) + ) + assert units.Quantity(1, integral.metpy.units).to('hg/(m s^2)').magnitude == 1 + + +def test_cumulative_integrate_xr_2d(): + """Test that cumulative_integrate works with 2D DataArrays.""" + arr = np.arange(5) + data_xr = xr.DataArray( + np.ones((5, 5)), + {'y': ('y', arr, {'units': 'm'}), 'x': ('x', arr, {'units': 'm'})}, + ('y', 'x'), + 'height', + {'units': 'm'} + ) + integral = cumulative_integrate(data_xr, axis='x') + assert integral.dims == data_xr.dims + assert integral.coords.keys() == data_xr.coords.keys() + assert units.Quantity(1, integral.metpy.units).to('m^2').magnitude == 1