Improved xarray compatibility on function input and output

examples/meteogram_metpy.py

@@ -81,7 +81,7 @@ def plot_winds(self, ws, wd, wsmax, plot_range=None):
         ax7.set_ylim(0, 360)
         ax7.set_yticks(np.arange(45, 405, 90), ['NE', 'SE', 'SW', 'NW'])
         lns = ln1 + ln2 + ln3
-        labs = [l.get_label() for l in lns]
+        labs = [ln.get_label() for ln in lns]
         ax7.xaxis.set_major_formatter(mpl.dates.DateFormatter('%d/%H UTC'))
         ax7.legend(lns, labs, loc='upper center',
                    bbox_to_anchor=(0.5, 1.2), ncol=3, prop={'size': 12})
@@ -112,7 +112,7 @@ def plot_thermo(self, t, td, plot_range=None):
         ax_twin = self.ax2.twinx()
         ax_twin.set_ylim(plot_range[0], plot_range[1], plot_range[2])
         lns = ln4 + ln5
-        labs = [l.get_label() for l in lns]
+        labs = [ln.get_label() for ln in lns]
         ax_twin.xaxis.set_major_formatter(mpl.dates.DateFormatter('%d/%H UTC'))
 
         self.ax2.legend(lns, labs, loc='upper center',

src/metpy/__init__.py

@@ -32,6 +32,6 @@
 os.environ['PINT_ARRAY_PROTOCOL_FALLBACK'] = '0'
 
 from ._version import get_version  # noqa: E402
-from .xarray import *  # noqa: F401, F403
+from .xarray import *  # noqa: F401, F403, E402
 __version__ = get_version()
 del get_version

src/metpy/calc/basic.py

@@ -15,11 +15,10 @@
 import numpy as np
 from scipy.ndimage import gaussian_filter
 
-from .tools import wrap_output_like
 from .. import constants as mpconsts
 from ..package_tools import Exporter
 from ..units import check_units, masked_array, units
-from ..xarray import preprocess_xarray
+from ..xarray import preprocess_and_wrap
 
 exporter = Exporter(globals())
 
@@ -29,7 +28,7 @@
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='u')
 @check_units('[speed]', '[speed]')
 def wind_speed(u, v):
     r"""Compute the wind speed from u and v-components.
@@ -56,7 +55,7 @@ def wind_speed(u, v):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='u')
 @check_units('[speed]', '[speed]')
 def wind_direction(u, v, convention='from'):
     r"""Compute the wind direction from u and v-components.
@@ -108,7 +107,7 @@ def wind_direction(u, v, convention='from'):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like=('speed', 'speed'))
 @check_units('[speed]')
 def wind_components(speed, wind_direction):
     r"""Calculate the U, V wind vector components from the speed and direction.
@@ -147,7 +146,7 @@ def wind_components(speed, wind_direction):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='temperature')
 @check_units(temperature='[temperature]', speed='[speed]')
 def windchill(temperature, speed, face_level_winds=False, mask_undefined=True):
     r"""Calculate the Wind Chill Temperature Index (WCTI).
@@ -209,7 +208,7 @@ def windchill(temperature, speed, face_level_winds=False, mask_undefined=True):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='temperature')
 @check_units('[temperature]')
 def heat_index(temperature, relative_humidity, mask_undefined=True):
     r"""Calculate the Heat Index from the current temperature and relative humidity.
@@ -313,7 +312,7 @@ def heat_index(temperature, relative_humidity, mask_undefined=True):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='temperature')
 @check_units(temperature='[temperature]', speed='[speed]')
 def apparent_temperature(temperature, relative_humidity, speed, face_level_winds=False,
                          mask_undefined=True):
@@ -392,7 +391,7 @@ def apparent_temperature(temperature, relative_humidity, speed, face_level_winds
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='pressure')
 @check_units('[pressure]')
 def pressure_to_height_std(pressure):
     r"""Convert pressure data to height using the U.S. standard atmosphere [NOAA1976]_.
@@ -420,7 +419,7 @@ def pressure_to_height_std(pressure):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='height')
 @check_units('[length]')
 def height_to_geopotential(height):
     r"""Compute geopotential for a given height above sea level.
@@ -477,7 +476,7 @@ def height_to_geopotential(height):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='geopotential')
 def geopotential_to_height(geopotential):
     r"""Compute height above sea level from a given geopotential.
 
@@ -537,7 +536,7 @@ def geopotential_to_height(geopotential):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='height')
 @check_units('[length]')
 def height_to_pressure_std(height):
     r"""Convert height data to pressures using the U.S. standard atmosphere [NOAA1976]_.
@@ -564,7 +563,7 @@ def height_to_pressure_std(height):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='latitude')
 def coriolis_parameter(latitude):
     r"""Calculate the coriolis parameter at each point.
 
@@ -586,7 +585,7 @@ def coriolis_parameter(latitude):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='pressure')
 @check_units('[pressure]', '[length]')
 def add_height_to_pressure(pressure, height):
     r"""Calculate the pressure at a certain height above another pressure level.
@@ -615,7 +614,7 @@ def add_height_to_pressure(pressure, height):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='height')
 @check_units('[length]', '[pressure]')
 def add_pressure_to_height(height, pressure):
     r"""Calculate the height at a certain pressure above another height.
@@ -644,7 +643,7 @@ def add_pressure_to_height(height, pressure):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='sigma')
 @check_units('[dimensionless]', '[pressure]', '[pressure]')
 def sigma_to_pressure(sigma, pressure_sfc, pressure_top):
     r"""Calculate pressure from sigma values.
@@ -687,7 +686,7 @@ def sigma_to_pressure(sigma, pressure_sfc, pressure_top):
 
 
 @exporter.export
-@wrap_output_like(argument='scalar_grid', match_unit=True)
+@preprocess_and_wrap(wrap_like='scalar_grid', match_unit=True, to_magnitude=True)
 def smooth_gaussian(scalar_grid, n):
     """Filter with normal distribution of weights.
 
@@ -779,7 +778,7 @@ def smooth_gaussian(scalar_grid, n):
 
 
 @exporter.export
-@wrap_output_like(argument='scalar_grid', match_unit=True)
+@preprocess_and_wrap(wrap_like='scalar_grid', match_unit=True, to_magnitude=True)
 def smooth_window(scalar_grid, window, passes=1, normalize_weights=True):
     """Filter with an arbitrary window smoother.
 
@@ -1007,7 +1006,7 @@ def smooth_n_point(scalar_grid, n=5, passes=1):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='altimeter_value')
 @check_units('[pressure]', '[length]')
 def altimeter_to_station_pressure(altimeter_value, height):
     r"""Convert the altimeter measurement to station pressure.
@@ -1089,7 +1088,7 @@ def altimeter_to_station_pressure(altimeter_value, height):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='altimeter_value')
 @check_units('[pressure]', '[length]', '[temperature]')
 def altimeter_to_sea_level_pressure(altimeter_value, height, temperature):
     r"""Convert the altimeter setting to sea-level pressure.

src/metpy/calc/indices.py

@@ -9,13 +9,13 @@
 from .. import constants as mpconsts
 from ..package_tools import Exporter
 from ..units import check_units, concatenate, units
-from ..xarray import preprocess_xarray
+from ..xarray import preprocess_and_wrap
 
 exporter = Exporter(globals())
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='dewpoint')
 @check_units('[pressure]', '[temperature]', bottom='[pressure]', top='[pressure]')
 def precipitable_water(pressure, dewpoint, *, bottom=None, top=None):
     r"""Calculate precipitable water through the depth of a sounding.
@@ -48,6 +48,10 @@ def precipitable_water(pressure, dewpoint, *, bottom=None, top=None):
     >>> dewpoint = np.array([20, 15, 10]) * units.degC
     >>> pw = precipitable_water(pressure, dewpoint)
 
+    Notes
+    -----
+    Only functions on 1D profiles (not higher-dimension vertical cross sections or grids).
+
     """
     # Sort pressure and dewpoint to be in decreasing pressure order (increasing height)
     sort_inds = np.argsort(pressure)[::-1]
@@ -74,7 +78,7 @@ def precipitable_water(pressure, dewpoint, *, bottom=None, top=None):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like=('pressure', 'pressure'))
 @check_units('[pressure]')
 def mean_pressure_weighted(pressure, *args, height=None, bottom=None, depth=None):
     r"""Calculate pressure-weighted mean of an arbitrary variable through a layer.
@@ -105,6 +109,10 @@ def mean_pressure_weighted(pressure, *args, height=None, bottom=None, depth=None
     `pint.Quantity`
         v_mean: v-component of layer mean wind.
 
+    Notes
+    -----
+    Only functions on 1D profiles (not higher-dimension vertical cross sections or grids).
+
     """
     ret = []  # Returned variable means in layer
     layer_arg = get_layer(pressure, *args, height=height,
@@ -123,7 +131,7 @@ def mean_pressure_weighted(pressure, *args, height=None, bottom=None, depth=None
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap()
 @check_units('[pressure]', '[speed]', '[speed]', '[length]')
 def bunkers_storm_motion(pressure, u, v, height):
     r"""Calculate the Bunkers right-mover and left-mover storm motions and sfc-6km mean flow.
@@ -150,6 +158,12 @@ def bunkers_storm_motion(pressure, u, v, height):
     wind_mean: `pint.Quantity`
         U and v component of sfc-6km mean flow
 
+    Notes
+    -----
+    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.
+
     """
     # mean wind from sfc-6km
     wind_mean = concatenate(mean_pressure_weighted(pressure, u, v, height=height,
@@ -183,7 +197,7 @@ def bunkers_storm_motion(pressure, u, v, height):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap()
 @check_units('[pressure]', '[speed]', '[speed]')
 def bulk_shear(pressure, u, v, height=None, bottom=None, depth=None):
     r"""Calculate bulk shear through a layer.
@@ -215,6 +229,12 @@ def bulk_shear(pressure, u, v, height=None, bottom=None, depth=None):
     v_shr: `pint.Quantity`
         v-component of layer bulk shear
 
+    Notes
+    -----
+    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.
+
     """
     _, u_layer, v_layer = get_layer(pressure, u, v, height=height,
                                     bottom=bottom, depth=depth)
@@ -226,7 +246,7 @@ def bulk_shear(pressure, u, v, height=None, bottom=None, depth=None):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='mucape')
 @check_units('[energy] / [mass]', '[speed] * [speed]', '[speed]')
 def supercell_composite(mucape, effective_storm_helicity, effective_shear):
     r"""Calculate the supercell composite parameter.
@@ -268,7 +288,7 @@ def supercell_composite(mucape, effective_storm_helicity, effective_shear):
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap(wrap_like='sbcape')
 @check_units('[energy] / [mass]', '[length]', '[speed] * [speed]', '[speed]')
 def significant_tornado(sbcape, surface_based_lcl_height, storm_helicity_1km, shear_6km):
     r"""Calculate the significant tornado parameter (fixed layer).
@@ -322,7 +342,7 @@ def significant_tornado(sbcape, surface_based_lcl_height, storm_helicity_1km, sh
 
 
 @exporter.export
-@preprocess_xarray
+@preprocess_and_wrap()
 @check_units('[pressure]', '[speed]', '[speed]', '[length]', '[speed]', '[speed]')
 def critical_angle(pressure, u, v, height, u_storm, v_storm):
     r"""Calculate the critical angle.
@@ -354,6 +374,12 @@ def critical_angle(pressure, u, v, height, u_storm, v_storm):
     `pint.Quantity`
         critical angle in degrees
 
+    Notes
+    -----
+    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.
+
     """
     # Convert everything to m/s
     u = u.to('m/s')