From 5d94b4ef895eca7f791badcd9b06b1adf243c7aa Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Wed, 23 Sep 2020 11:12:28 -0500 Subject: [PATCH] Misc updates based on review feedback --- src/metpy/calc/thermo.py | 13 +++++++------ src/metpy/calc/tools.py | 28 ++++++++++++++-------------- src/metpy/xarray.py | 6 +++--- tests/test_xarray.py | 25 +++++++++++++++++++++++++ 4 files changed, 49 insertions(+), 23 deletions(-) diff --git a/src/metpy/calc/thermo.py b/src/metpy/calc/thermo.py index 903323c191b..9bbb1c9c768 100644 --- a/src/metpy/calc/thermo.py +++ b/src/metpy/calc/thermo.py @@ -258,7 +258,10 @@ def dry_lapse(pressure, temperature, reference_pressure=None, vertical_dim=0): @exporter.export -@preprocess_and_wrap() +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'reference_pressure') +) @check_units('[pressure]', '[temperature]', '[pressure]') def moist_lapse(pressure, temperature, reference_pressure=None): r"""Calculate the temperature at a level assuming liquid saturation processes. @@ -280,8 +283,7 @@ def moist_lapse(pressure, temperature, reference_pressure=None): Returns ------- `pint.Quantity` - The temperature corresponding to the starting temperature and - pressure levels. + The resulting parcel temperature at levels given by `pressure` See Also -------- @@ -298,9 +300,8 @@ def moist_lapse(pressure, temperature, reference_pressure=None): This equation comes from [Bakhshaii2013]_. - Only functions on 1D profiles (not higher-dimension vertical cross sections or grids). - Since this function returns scalar values when given a profile, this will return Pint - Quantities even when given xarray DataArray profiles. + Only reliably functions on 1D profiles (not higher-dimension vertical cross sections or + grids). """ def dt(t, p): diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index d3996c71eb4..3a1f8c66e03 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -1175,7 +1175,7 @@ def second_derivative(f, axis=None, x=None, delta=None): @exporter.export -def gradient(f, coordinates=None, deltas=None, axes=None): +def gradient(f, axes=None, coordinates=None, deltas=None): """Calculate the gradient of a grid of values. Works for both regularly-spaced data, and grids with varying spacing. @@ -1191,13 +1191,6 @@ def gradient(f, coordinates=None, deltas=None, axes=None): ---------- f : array-like Array of values of which to calculate the derivative - coordinates : array-like, optional - Sequence of arrays containing the coordinate values corresponding to the - grid points in `f` in axis order. - deltas : array-like, optional - Sequence of arrays or scalars that specify the spacing between the grid points in `f` - in axis order. There should be one item less than the size of `f` along the applicable - axis. axes : sequence, optional Sequence of strings (if `f` is a `xarray.DataArray` and implicit conversion to `pint.Quantity` is not used) or integers that specify the array axes along which to @@ -1206,6 +1199,13 @@ def gradient(f, coordinates=None, deltas=None, axes=None): `coordinates` or `deltas` given. In general, each axis can be an axis number (integer), dimension coordinate name (string) or a standard axis type (string). The current standard axis types are 'time', 'vertical', 'y', and 'x'. + coordinates : array-like, optional + Sequence of arrays containing the coordinate values corresponding to the + grid points in `f` in axis order. + deltas : array-like, optional + Sequence of arrays or scalars that specify the spacing between the grid points in `f` + in axis order. There should be one item less than the size of `f` along the applicable + axis. Returns ------- @@ -1228,7 +1228,7 @@ def gradient(f, coordinates=None, deltas=None, axes=None): @exporter.export -def laplacian(f, coordinates=None, deltas=None, axes=None): +def laplacian(f, axes=None, coordinates=None, deltas=None): """Calculate the laplacian of a grid of values. Works for both regularly-spaced data, and grids with varying spacing. @@ -1244,11 +1244,6 @@ def laplacian(f, coordinates=None, deltas=None, axes=None): ---------- f : array-like Array of values of which to calculate the derivative - coordinates : array-like, optional - The coordinate values corresponding to the grid points in `f` - deltas : array-like, optional - Spacing between the grid points in `f`. There should be one item less than the size - of `f` along the applicable axis. axes : sequence, optional Sequence of strings (if `f` is a `xarray.DataArray` and implicit conversion to `pint.Quantity` is not used) or integers that specify the array axes along which to @@ -1257,6 +1252,11 @@ def laplacian(f, coordinates=None, deltas=None, axes=None): `coordinates` or `deltas` given. In general, each axis can be an axis number (integer), dimension coordinate name (string) or a standard axis type (string). The current standard axis types are 'time', 'vertical', 'y', and 'x'. + coordinates : array-like, optional + The coordinate values corresponding to the grid points in `f` + deltas : array-like, optional + Spacing between the grid points in `f`. There should be one item less than the size + of `f` along the applicable axis. Returns ------- diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index c528984d1da..9a082ad965c 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -1071,7 +1071,7 @@ def preprocess_and_wrap(broadcast=None, wrap_like=None, match_unit=False, to_mag with default of None. wrap_like : str or array-like or tuple of str or tuple of array-like or None Wrap the calculation output following a particular input argument (if str) or data - data object (if array-like). If tuple, will assume output is in the form of a tuple, + object (if array-like). If tuple, will assume output is in the form of a tuple, and wrap iteratively according to the str or array-like contained within. If None, will not wrap output. match_unit : bool @@ -1182,11 +1182,11 @@ def _wrap_output_like_not_matching_units(result, match): else: # Determine if need to upcast to Quantity if ( - ( + not isinstance(result, units.Quantity) + and ( isinstance(match, units.Quantity) or (output_xarray and isinstance(match.data, units.Quantity)) ) - and not isinstance(result, units.Quantity) ): result = units.Quantity(result) return ( diff --git a/tests/test_xarray.py b/tests/test_xarray.py index dbbfa206269..5a38ce98f91 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -1256,3 +1256,28 @@ def func(a, b): return a * b np.testing.assert_array_equal(func(data, data2), np.array([0, 0, 1])) + + +def test_preprocess_and_wrap_with_variable(): + """Test preprocess and wrapping decorator when using an xarray Variable.""" + data1 = xr.DataArray([1, 0, 1], dims=('x',), attrs={'units': 'meter'}) + data2 = xr.Variable(data=[0, 1, 1], dims=('x',), attrs={'units': 'meter'}) + + @preprocess_and_wrap(wrap_like='a') + def func(a, b): + return a * b + + # Note, expected units are meter, not meter**2, since attributes are stripped from + # Variables + expected_12 = xr.DataArray([0, 0, 1] * units.m, dims=('x',)) + expected_21 = [0, 0, 1] * units.m + + with pytest.warns(UserWarning, match='Argument b given as xarray Variable'): + result_12 = func(data1, data2) + with pytest.warns(UserWarning, match='Argument a given as xarray Variable'): + result_21 = func(data2, data1) + + assert isinstance(result_12, xr.DataArray) + xr.testing.assert_identical(func(data1, data2), expected_12) + assert isinstance(result_21, units.Quantity) + assert_array_equal(func(data2, data1), expected_21)