diff --git a/examples/meteogram_metpy.py b/examples/meteogram_metpy.py index 844209136f7..4c306fa9069 100644 --- a/examples/meteogram_metpy.py +++ b/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', diff --git a/src/metpy/__init__.py b/src/metpy/__init__.py index 9b01fb469a7..ee1199205ec 100644 --- a/src/metpy/__init__.py +++ b/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 diff --git a/src/metpy/calc/basic.py b/src/metpy/calc/basic.py index 5d24e816377..06312a7de03 100644 --- a/src/metpy/calc/basic.py +++ b/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. diff --git a/src/metpy/calc/indices.py b/src/metpy/calc/indices.py index f297e25ceb1..af1f1a8abd8 100644 --- a/src/metpy/calc/indices.py +++ b/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') diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 5b7e876e4e1..2c882bfd4e8 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -2,15 +2,20 @@ # Distributed under the terms of the BSD 3-Clause License. # SPDX-License-Identifier: BSD-3-Clause """Contains calculation of kinematic parameters (e.g. divergence or vorticity).""" +import functools +from inspect import signature +import warnings + import numpy as np +import xarray as xr from . import coriolis_parameter -from .tools import first_derivative, get_layer_heights, gradient +from .tools import first_derivative, get_layer_heights, gradient, grid_deltas_from_dataarray from .. import constants as mpconsts from ..cbook import iterable 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()) @@ -19,184 +24,279 @@ def _stack(arrs): return concatenate([a[np.newaxis] if iterable(a) else a for a in arrs], axis=0) +def add_grid_arguments_from_xarray(func): + """Fill in optional arguments like dx/dy from DataArray arguments.""" + @functools.wraps(func) + def wrapper(*args, **kwargs): + bound_args = signature(func).bind(*args, **kwargs) + bound_args.apply_defaults() + + # Search for DataArray with valid latitude and longitude coordinates to find grid + # deltas and any other needed parameter + dataarray_arguments = [ + value for value in bound_args.arguments.values() + if isinstance(value, xr.DataArray) + ] + grid_prototype = None + for da in dataarray_arguments: + if hasattr(da.metpy, 'latitude') and hasattr(da.metpy, 'longitude'): + grid_prototype = da + break + + # Fill in x_dim/y_dim + if ( + grid_prototype is not None + and 'x_dim' in bound_args.arguments + and 'y_dim' in bound_args.arguments + ): + try: + bound_args.arguments['x_dim'] = grid_prototype.metpy.find_axis_number('x') + bound_args.arguments['y_dim'] = grid_prototype.metpy.find_axis_number('y') + except AttributeError: + # If axis number not found, fall back to default but warn. + warnings.warn('Horizontal dimension numbers not found. Defaulting to ' + '(..., Y, X) order.') + + # Fill in dx/dy + if ( + 'dx' in bound_args.arguments and bound_args.arguments['dx'] is None + and 'dy' in bound_args.arguments and bound_args.arguments['dy'] is None + ): + if grid_prototype is not None: + bound_args.arguments['dx'], bound_args.arguments['dy'] = ( + grid_deltas_from_dataarray(grid_prototype, kind='actual') + ) + else: + raise ValueError('Must provide dx/dy arguments or input DataArray with ' + 'latitude/longitude coordinates.') + + # Fill in latitude + if 'latitude' in bound_args.arguments and bound_args.arguments['latitude'] is None: + if grid_prototype is not None: + bound_args.arguments['latitude'] = ( + grid_prototype.metpy.latitude.metpy.unit_array + ) + else: + raise ValueError('Must provide latitude argument or input DataArray with ' + 'latitude/longitude coordinates.') + + # Fill in Coriolis parameter + if 'f' in bound_args.arguments and bound_args.arguments['f'] is None: + if grid_prototype is not None: + bound_args.arguments['f'] = coriolis_parameter( + grid_prototype.metpy.latitude.metpy.unit_array + ) + else: + raise ValueError('Must provide f argument or input DataArray with' + 'latitude/longitude coordinates.') + + return func(*bound_args.args, **bound_args.kwargs) + return wrapper + + @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like='u') @check_units('[speed]', '[speed]', '[length]', '[length]') -def vorticity(u, v, dx, dy): +def vorticity(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): r"""Calculate the vertical vorticity of the horizontal wind. Parameters ---------- - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the wind - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` vertical vorticity See Also -------- divergence - Notes - ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - """ - dudy = first_derivative(u, delta=dy, axis=-2) - dvdx = first_derivative(v, delta=dx, axis=-1) + dudy = first_derivative(u, delta=dy, axis=y_dim) + dvdx = first_derivative(v, delta=dx, axis=x_dim) return dvdx - dudy @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like='u') @check_units(dx='[length]', dy='[length]') -def divergence(u, v, dx, dy): +def divergence(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): r"""Calculate the horizontal divergence of a vector. Parameters ---------- - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the vector - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the vector - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` The horizontal divergence See Also -------- vorticity - Notes - ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - """ - dudx = first_derivative(u, delta=dx, axis=-1) - dvdy = first_derivative(v, delta=dy, axis=-2) + dudx = first_derivative(u, delta=dx, axis=x_dim) + dvdy = first_derivative(v, delta=dy, axis=y_dim) return dudx + dvdy @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like='u') @check_units('[speed]', '[speed]', '[length]', '[length]') -def shearing_deformation(u, v, dx, dy): +def shearing_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): r"""Calculate the shearing deformation of the horizontal wind. Parameters ---------- - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the wind - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` Shearing Deformation See Also -------- stretching_deformation, total_deformation - Notes - ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - """ - dudy = first_derivative(u, delta=dy, axis=-2) - dvdx = first_derivative(v, delta=dx, axis=-1) + dudy = first_derivative(u, delta=dy, axis=y_dim) + dvdx = first_derivative(v, delta=dx, axis=x_dim) return dvdx + dudy @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like='u') @check_units('[speed]', '[speed]', '[length]', '[length]') -def stretching_deformation(u, v, dx, dy): +def stretching_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): r"""Calculate the stretching deformation of the horizontal wind. Parameters ---------- - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the wind - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` Stretching Deformation See Also -------- shearing_deformation, total_deformation - Notes - ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - """ - dudx = first_derivative(u, delta=dx, axis=-1) - dvdy = first_derivative(v, delta=dy, axis=-2) + dudx = first_derivative(u, delta=dx, axis=x_dim) + dvdy = first_derivative(v, delta=dy, axis=y_dim) return dudx - dvdy @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like='u') @check_units('[speed]', '[speed]', '[length]', '[length]') -def total_deformation(u, v, dx, dy): +def total_deformation(u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): r"""Calculate the horizontal total deformation of the horizontal wind. Parameters ---------- - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the wind - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` Total Deformation See Also @@ -209,13 +309,12 @@ def total_deformation(u, v, dx, dy): of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. """ - dudy, dudx = gradient(u, deltas=(dy, dx), axes=(-2, -1)) - dvdy, dvdx = gradient(v, deltas=(dy, dx), axes=(-2, -1)) + dudy, dudx = gradient(u, deltas=(dy, dx), axes=(y_dim, x_dim)) + dvdy, dvdx = gradient(v, deltas=(dy, dx), axes=(y_dim, x_dim)) return np.sqrt((dvdx + dudy)**2 + (dudx - dvdy)**2) @exporter.export -@preprocess_xarray def advection(scalar, wind, deltas): r"""Calculate the advection of a scalar field by the wind. @@ -244,6 +343,9 @@ def advection(scalar, wind, deltas): An N-dimensional array containing the advection at all grid points. """ + # TODO: This requires custom handling with xarray, since it does not fit the API paradigm + # of the rest of kinematics + # This allows passing in a list of wind components or an array. wind = _stack(wind) @@ -266,9 +368,13 @@ def advection(scalar, wind, deltas): @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap( + wrap_like='potential_temperature', + broadcast=('potential_temperature', 'u', 'v') +) @check_units('[temperature]', '[speed]', '[speed]', '[length]', '[length]') -def frontogenesis(potential_temperature, u, v, dx, dy): +def frontogenesis(potential_temperature, u, v, dx=None, dy=None, x_dim=-1, y_dim=-2): r"""Calculate the 2D kinematic frontogenesis of a temperature field. The implementation is a form of the Petterssen Frontogenesis and uses the formula @@ -284,47 +390,52 @@ def frontogenesis(potential_temperature, u, v, dx, dy): Parameters ---------- - potential_temperature : (M, N) `pint.Quantity` + potential_temperature : (..., M, N) `xarray.DataArray` or `pint.Quantity` Potential temperature - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the wind - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` 2D Frontogenesis in [temperature units]/m/s Notes ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - Conversion factor to go from [temperature units]/m/s to [temperature units/100km/3h] :math:`1.08e4*1.e5` """ # Get gradients of potential temperature in both x and y - ddy_thta = first_derivative(potential_temperature, delta=dy, axis=-2) - ddx_thta = first_derivative(potential_temperature, delta=dx, axis=-1) + ddy_thta = first_derivative(potential_temperature, delta=dy, axis=y_dim) + ddx_thta = first_derivative(potential_temperature, delta=dx, axis=x_dim) # Compute the magnitude of the potential temperature gradient mag_thta = np.sqrt(ddx_thta**2 + ddy_thta**2) # Get the shearing, stretching, and total deformation of the wind field - shrd = shearing_deformation(u, v, dx, dy) - strd = stretching_deformation(u, v, dx, dy) - tdef = total_deformation(u, v, dx, dy) + shrd = shearing_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim) + strd = stretching_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim) + tdef = total_deformation(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim) # Get the divergence of the wind field - div = divergence(u, v, dx, dy) + div = divergence(u, v, dx, dy, x_dim=x_dim, y_dim=y_dim) # Compute the angle (beta) between the wind field and the gradient of potential temperature psi = 0.5 * np.arctan2(shrd, strd) @@ -334,92 +445,97 @@ def frontogenesis(potential_temperature, u, v, dx, dy): @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like=('height', 'height'), broadcast=('height', 'f')) @check_units(f='[frequency]', dx='[length]', dy='[length]') -def geostrophic_wind(height, f, dx, dy): +def geostrophic_wind(height, f=None, dx=None, dy=None, x_dim=-1, y_dim=-2): r"""Calculate the geostrophic wind given from the height or geopotential. Parameters ---------- - height : (M, N) `pint.Quantity` - The height field, with either leading dimensions of (x, y) or trailing dimensions - of (y, x), depending on the value of ``dim_order``. - f : array_like - The coriolis parameter. This can be a scalar to be applied - everywhere or an array of values. - dx : `pint.Quantity` + height : (..., M, N) `xarray.DataArray` or `pint.Quantity` + The height or geopotential field. + f : `pint.Quantity`, optional + The Coriolis parameter. This can be a scalar to be applied everywhere or an array of + values. Optional if `xarray.DataArray` with latitude/longitude coordinates used as + input for height. + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `height` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `height` along the applicable axis. + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - A 2-item tuple of arrays, `pint.Quantity` + A 2-item tuple of arrays A tuple of the u-component and v-component of the geostrophic wind. - Notes - ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - """ if height.dimensionality['[length]'] == 2.0: norm_factor = 1. / f else: norm_factor = mpconsts.g / f - dhdy = first_derivative(height, delta=dy, axis=-2) - dhdx = first_derivative(height, delta=dx, axis=-1) + dhdy = first_derivative(height, delta=dy, axis=y_dim) + dhdx = first_derivative(height, delta=dx, axis=x_dim) return -norm_factor * dhdy, norm_factor * dhdx @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like=('height', 'height'), broadcast=('height', 'u', 'v', 'f')) @check_units(f='[frequency]', u='[speed]', v='[speed]', dx='[length]', dy='[length]') -def ageostrophic_wind(height, u, v, f, dx, dy): +def ageostrophic_wind(height, u, v, f=None, dx=None, dy=None, x_dim=-1, y_dim=-2): r"""Calculate the ageostrophic wind given from the height or geopotential. Parameters ---------- height : (M, N) ndarray The height or geopotential field. - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` The u wind field. - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` The u wind field. - f : array_like - The coriolis parameter. This can be a scalar to be applied - everywhere or an array of values. - dx : `pint.Quantity` + f : `pint.Quantity`, optional + The Coriolis parameter. This can be a scalar to be applied everywhere or an array of + values. Optional if `xarray.DataArray` with latitude/longitude coordinates used as + input for height. + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `height` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `height` along the applicable axis. + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- A 2-item tuple of arrays A tuple of the u-component and v-component of the ageostrophic wind. - Notes - ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - - This function contains an updated input variable order from the same function in the - kinematics module. This version will be fully implemented in 1.0 and moved from the - `future` module back to the `kinematics` module. - """ u_geostrophic, v_geostrophic = geostrophic_wind(height, f, dx, dy) return u - u_geostrophic, v - v_geostrophic @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='height', broadcast=('height', 'temperature')) @check_units('[length]', '[temperature]') def montgomery_streamfunction(height, temperature): r"""Compute the Montgomery Streamfunction on isentropic surfaces. @@ -461,7 +577,7 @@ def montgomery_streamfunction(height, temperature): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[length]', '[speed]', '[speed]', '[length]', bottom='[length]', storm_u='[speed]', storm_v='[speed]') def storm_relative_helicity(height, u, v, depth, *, bottom=0 * units.m, @@ -504,6 +620,12 @@ def storm_relative_helicity(height, u, v, depth, *, bottom=0 * units.m, `pint.Quantity` total storm-relative helicity + 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, v = get_layer_heights(height, depth, u, v, with_agl=True, bottom=bottom) @@ -528,36 +650,42 @@ def storm_relative_helicity(height, u, v, depth, *, bottom=0 * units.m, @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like='u') @check_units('[speed]', '[speed]', '[length]', '[length]') -def absolute_vorticity(u, v, dx, dy, latitude): +def absolute_vorticity(u, v, dx=None, dy=None, latitude=None, x_dim=-1, y_dim=-2): """Calculate the absolute vorticity of the horizontal wind. Parameters ---------- - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the wind - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - latitude : (M, N) ndarray - latitude of the wind data in radians or with appropriate unit information attached + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + latitude : `pint.Quantity`, optional + Latitude of the wind data. Optional if `xarray.DataArray` with latitude/longitude + coordinates used as input. Note that an argument without units is treated as + dimensionless, which translates to radians. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` absolute vorticity - Notes - ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - """ f = coriolis_parameter(latitude) relative_vorticity = vorticity(u, v, dx, dy) @@ -565,10 +693,25 @@ def absolute_vorticity(u, v, dx, dy, latitude): @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap( + wrap_like='potential_temperature', + broadcast=('potential_temperature', 'pressure', 'u', 'v') +) @check_units('[temperature]', '[pressure]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') -def potential_vorticity_baroclinic(potential_temperature, pressure, u, v, dx, dy, latitude): +def potential_vorticity_baroclinic( + potential_temperature, + pressure, + u, + v, + dx=None, + dy=None, + latitude=None, + x_dim=-1, + y_dim=-2, + vertical_dim=-3 +): r"""Calculate the baroclinic potential vorticity. .. math:: PV = -g \left(\frac{\partial u}{\partial p}\frac{\partial \theta}{\partial y} @@ -579,34 +722,43 @@ def potential_vorticity_baroclinic(potential_temperature, pressure, u, v, dx, dy Parameters ---------- - potential_temperature : (P, M, N) `pint.Quantity` + potential_temperature : (..., P, M, N) `xarray.DataArray` or `pint.Quantity` potential temperature - pressure : (P, M, N) `pint.Quantity` + pressure : (..., P, M, N) `xarray.DataArray` or `pint.Quantity` vertical pressures - u : (P, M, N) `pint.Quantity` + u : (..., P, M, N) `xarray.DataArray` or `pint.Quantity` x component of the wind - v : (P, M, N) `pint.Quantity` + v : (..., P, M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - latitude : (M, N) ndarray - latitude of the wind data in radians or with appropriate unit information attached + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + latitude : `pint.Quantity`, optional + Latitude of the wind data. Optional if `xarray.DataArray` with latitude/longitude + coordinates used as input. Note that an argument without units is treated as + dimensionless, which translates to radians. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Z, Y, X] order). + Automatically parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Z, Y, X] order). + Automatically parsed from input if using `xarray.DataArray`. + vertical_dim : int, optional + Axis number of vertical dimension. Defaults to -3 (implying [..., Z, Y, X] order). + Automatically parsed from input if using `xarray.DataArray`. Returns ------- - (P, M, N) `pint.Quantity` + (..., P, M, N) `xarray.DataArray` or `pint.Quantity` baroclinic potential vorticity Notes ----- - This function will only work with data that is in (P, Y, X) format. If your data - is in a different order you will need to re-order your data in order to get correct - results from this function. - The same function can be used for isobaric and isentropic PV analysis. Provide winds for vorticity calculations on the desired isobaric or isentropic surface. At least three layers of pressure/potential temperature are required in order to calculate the vertical @@ -616,26 +768,32 @@ def potential_vorticity_baroclinic(potential_temperature, pressure, u, v, dx, dy This function expects pressure/isentropic level to increase with increasing array element (e.g., from higher in the atmosphere to closer to the surface. If the pressure array is - one-dimensional p[:, None, None] can be used to make it appear multi-dimensional.) + one-dimensional, and not given as `xarray.DataArray`, p[:, None, None] can be used to make + it appear multi-dimensional.) """ - if ((np.shape(potential_temperature)[-3] < 3) or (np.shape(pressure)[-3] < 3) - or (np.shape(potential_temperature)[-3] != (np.shape(pressure)[-3]))): - raise ValueError('Length of potential temperature along the pressure axis ' - '{} must be at least 3.'.format(-3)) - - avor = absolute_vorticity(u, v, dx, dy, latitude) - dthtadp = first_derivative(potential_temperature, x=pressure, axis=-3) - - if ((np.shape(potential_temperature)[-2] == 1) - and (np.shape(potential_temperature)[-1] == 1)): - dthtady = 0 * units.K / units.m # axis=-2 only has one dimension - dthtadx = 0 * units.K / units.m # axis=-1 only has one dimension + if ( + np.shape(potential_temperature)[vertical_dim] < 3 + or np.shape(pressure)[vertical_dim] < 3 + or np.shape(potential_temperature)[vertical_dim] != np.shape(pressure)[vertical_dim] + ): + raise ValueError('Length of potential temperature along the vertical axis ' + '{} must be at least 3.'.format(vertical_dim)) + + avor = absolute_vorticity(u, v, dx, dy, latitude, x_dim=x_dim, y_dim=y_dim) + dthtadp = first_derivative(potential_temperature, x=pressure, axis=vertical_dim) + + if ( + (np.shape(potential_temperature)[y_dim] == 1) + and (np.shape(potential_temperature)[x_dim] == 1) + ): + dthtady = 0 * units.K / units.m # axis=y_dim only has one dimension + dthtadx = 0 * units.K / units.m # axis=x_dim only has one dimension else: - dthtady = first_derivative(potential_temperature, delta=dy, axis=-2) - dthtadx = first_derivative(potential_temperature, delta=dx, axis=-1) - dudp = first_derivative(u, x=pressure, axis=-3) - dvdp = first_derivative(v, x=pressure, axis=-3) + dthtady = first_derivative(potential_temperature, delta=dy, axis=y_dim) + dthtadx = first_derivative(potential_temperature, delta=dx, axis=x_dim) + dudp = first_derivative(u, x=pressure, axis=vertical_dim) + dvdp = first_derivative(v, x=pressure, axis=vertical_dim) return (-mpconsts.g * (dudp * dthtady - dvdp * dthtadx + avor * dthtadp)).to(units.kelvin * units.meter**2 @@ -643,9 +801,19 @@ def potential_vorticity_baroclinic(potential_temperature, pressure, u, v, dx, dy @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like='height', broadcast=('height', 'u', 'v')) @check_units('[length]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') -def potential_vorticity_barotropic(height, u, v, dx, dy, latitude): +def potential_vorticity_barotropic( + height, + u, + v, + dx=None, + dy=None, + latitude=None, + x_dim=-1, + y_dim=-2 +): r"""Calculate the barotropic (Rossby) potential vorticity. .. math:: PV = \frac{f + \zeta}{H} @@ -654,41 +822,57 @@ def potential_vorticity_barotropic(height, u, v, dx, dy, latitude): Parameters ---------- - height : (M, N) `pint.Quantity` + height : (..., M, N) `xarray.DataArray` or `pint.Quantity` atmospheric height - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the wind - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - latitude : (M, N) ndarray - latitude of the wind data in radians or with appropriate unit information attached + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + latitude : `pint.Quantity`, optional + Latitude of the wind data. Optional if `xarray.DataArray` with latitude/longitude + coordinates used as input. Note that an argument without units is treated as + dimensionless, which translates to radians. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` barotropic potential vorticity - Notes - ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - """ - avor = absolute_vorticity(u, v, dx, dy, latitude) + avor = absolute_vorticity(u, v, dx, dy, latitude, x_dim=x_dim, y_dim=y_dim) return (avor / height).to('meter**-1 * second**-1') @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap(wrap_like=('u', 'u')) @check_units('[speed]', '[speed]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') -def inertial_advective_wind(u, v, u_geostrophic, v_geostrophic, dx, dy, latitude): +def inertial_advective_wind( + u, + v, + u_geostrophic, + v_geostrophic, + dx=None, + dy=None, + latitude=None, + x_dim=-1, + y_dim=-2 +): r"""Calculate the inertial advective wind. .. math:: \frac{\hat k}{f} \times (\vec V \cdot \nabla)\hat V_g @@ -706,28 +890,38 @@ def inertial_advective_wind(u, v, u_geostrophic, v_geostrophic, dx, dy, latitude Parameters ---------- - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the advecting wind - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the advecting wind - u_geostrophic : (M, N) `pint.Quantity` + u_geostrophic : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the geostrophic (advected) wind - v_geostrophic : (M, N) `pint.Quantity` + v_geostrophic : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the geostrophic (advected) wind - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - latitude : (M, N) ndarray - latitude of the wind data in radians or with appropriate unit information attached + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + latitude : `pint.Quantity`, optional + Latitude of the wind data. Optional if `xarray.DataArray` with latitude/longitude + coordinates used as input. Note that an argument without units is treated as + dimensionless, which translates to radians. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of inertial advective wind - (M, N) `pint.Quantity` + (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of inertial advective wind Notes @@ -736,14 +930,11 @@ def inertial_advective_wind(u, v, u_geostrophic, v_geostrophic, dx, dy, latitude wind to both be the geostrophic wind. To do so, pass the x and y components of the geostrophic with for u and u_geostrophic/v and v_geostrophic. - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - """ f = coriolis_parameter(latitude) - dugdy, dugdx = gradient(u_geostrophic, deltas=(dy, dx), axes=(-2, -1)) - dvgdy, dvgdx = gradient(v_geostrophic, deltas=(dy, dx), axes=(-2, -1)) + dugdy, dugdx = gradient(u_geostrophic, deltas=(dy, dx), axes=(y_dim, x_dim)) + dvgdy, dvgdx = gradient(v_geostrophic, deltas=(dy, dx), axes=(y_dim, x_dim)) u_component = -(u * dvgdx + v * dvgdy) / f v_component = (u * dugdx + v * dugdy) / f @@ -752,9 +943,23 @@ def inertial_advective_wind(u, v, u_geostrophic, v_geostrophic, dx, dy, latitude @exporter.export -@preprocess_xarray +@add_grid_arguments_from_xarray +@preprocess_and_wrap( + wrap_like=('u', 'u'), + broadcast=('u', 'v', 'temperature', 'pressure', 'static_stability') +) @check_units('[speed]', '[speed]', '[temperature]', '[pressure]', '[length]', '[length]') -def q_vector(u, v, temperature, pressure, dx, dy, static_stability=1): +def q_vector( + u, + v, + temperature, + pressure, + dx=None, + dy=None, + static_stability=1, + x_dim=-1, + y_dim=-2 +): r"""Calculate Q-vector at a given pressure level using the u, v winds and temperature. .. math:: \vec{Q} = (Q_1, Q_2) @@ -774,42 +979,45 @@ def q_vector(u, v, temperature, pressure, dx, dy, static_stability=1): Parameters ---------- - u : (M, N) `pint.Quantity` + u : (..., M, N) `xarray.DataArray` or `pint.Quantity` x component of the wind (geostrophic in QG-theory) - v : (M, N) `pint.Quantity` + v : (..., M, N) `xarray.DataArray` or `pint.Quantity` y component of the wind (geostrophic in QG-theory) - temperature : (M, N) `pint.Quantity` + temperature : (..., M, N) `xarray.DataArray` or `pint.Quantity` Array of temperature at pressure level pressure : `pint.Quantity` Pressure at level - dx : `pint.Quantity` + dx : `pint.Quantity`, optional The grid spacing(s) in the x-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. - dy : `pint.Quantity` + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. + dy : `pint.Quantity`, optional The grid spacing(s) in the y-direction. If an array, there should be one item less than - the size of `u` along the applicable axis. + the size of `u` along the applicable axis. Optional if `xarray.DataArray` with + latitude/longitude coordinates used as input. static_stability : `pint.Quantity`, optional The static stability at the pressure level. Defaults to 1 if not given to calculate the Q-vector without factoring in static stability. + x_dim : int, optional + Axis number of x dimension. Defaults to -1 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. + y_dim : int, optional + Axis number of y dimension. Defaults to -2 (implying [..., Y, X] order). Automatically + parsed from input if using `xarray.DataArray`. Returns ------- - tuple of (M, N) `pint.Quantity` + tuple of (..., M, N) `xarray.DataArray` or `pint.Quantity` The components of the Q-vector in the u- and v-directions respectively See Also -------- static_stability - Notes - ----- - If inputs have more than two dimensions, they are assumed to have either leading dimensions - of (x, y) or trailing dimensions of (y, x), depending on the value of ``dim_order``. - """ - dudy, dudx = gradient(u, deltas=(dy, dx), axes=(-2, -1)) - dvdy, dvdx = gradient(v, deltas=(dy, dx), axes=(-2, -1)) - dtempdy, dtempdx = gradient(temperature, deltas=(dy, dx), axes=(-2, -1)) + dudy, dudx = gradient(u, deltas=(dy, dx), axes=(y_dim, x_dim)) + dvdy, dvdx = gradient(v, deltas=(dy, dx), axes=(y_dim, x_dim)) + dtempdy, dtempdx = gradient(temperature, deltas=(dy, dx), axes=(y_dim, x_dim)) q1 = -mpconsts.Rd / (pressure * static_stability) * (dudx * dtempdx + dvdx * dtempdy) q2 = -mpconsts.Rd / (pressure * static_stability) * (dudy * dtempdx + dvdy * dtempdy) diff --git a/src/metpy/calc/thermo.py b/src/metpy/calc/thermo.py index f7007ee1209..c588f56fbe4 100644 --- a/src/metpy/calc/thermo.py +++ b/src/metpy/calc/thermo.py @@ -7,6 +7,7 @@ import numpy as np import scipy.integrate as si import scipy.optimize as so +import xarray as xr from .tools import (_greater_or_close, _less_or_close, _remove_nans, find_bounding_indices, find_intersections, first_derivative, get_layer) @@ -15,7 +16,7 @@ from ..interpolate.one_dimension import interpolate_1d 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()) @@ -23,13 +24,13 @@ @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='temperature', broadcast=('temperature', 'dewpoint')) @check_units('[temperature]', '[temperature]') def relative_humidity_from_dewpoint(temperature, dewpoint): r"""Calculate the relative humidity. - Uses temperature and dewpoint in celsius to calculate relative - humidity using the ratio of vapor pressure to saturation vapor pressures. + Uses temperature and dewpoint to calculate relative humidity as the ratio of vapor + pressure to saturation vapor pressures. Parameters ---------- @@ -54,7 +55,7 @@ def relative_humidity_from_dewpoint(temperature, dewpoint): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='pressure') @check_units('[pressure]', '[pressure]') def exner_function(pressure, reference_pressure=mpconsts.P0): r"""Calculate the Exner function. @@ -89,7 +90,7 @@ def exner_function(pressure, reference_pressure=mpconsts.P0): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='temperature', broadcast=('pressure', 'temperature')) @check_units('[pressure]', '[temperature]') def potential_temperature(pressure, temperature): r"""Calculate the potential temperature. @@ -131,7 +132,10 @@ def potential_temperature(pressure, temperature): @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='potential_temperature', + broadcast=('pressure', 'potential_temperature') +) @check_units('[pressure]', '[temperature]') def temperature_from_potential_temperature(pressure, potential_temperature): r"""Calculate the temperature from a given potential temperature. @@ -176,9 +180,12 @@ def temperature_from_potential_temperature(pressure, potential_temperature): @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'reference_pressure') +) @check_units('[pressure]', '[temperature]', '[pressure]') -def dry_lapse(pressure, temperature, reference_pressure=None): +def dry_lapse(pressure, temperature, reference_pressure=None, vertical_dim=0): r"""Calculate the temperature at a level assuming only dry processes. This function lifts a parcel starting at `temperature`, conserving @@ -205,6 +212,11 @@ def dry_lapse(pressure, temperature, reference_pressure=None): parcel_profile : Calculate complete parcel profile potential_temperature + Notes + ----- + Only reliably functions on 1D profiles (not higher-dimension vertical cross sections or + grids) unless reference_pressure is specified. + """ if reference_pressure is None: reference_pressure = pressure[0] @@ -212,7 +224,7 @@ def dry_lapse(pressure, temperature, reference_pressure=None): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[pressure]') def moist_lapse(pressure, temperature, reference_pressure=None): r"""Calculate the temperature at a level assuming liquid saturation processes. @@ -252,6 +264,10 @@ 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. + """ def dt(t, p): t = units.Quantity(t, temperature.units) @@ -300,7 +316,7 @@ def dt(t, p): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') def lcl(pressure, temperature, dewpoint, max_iters=50, eps=1e-5): r"""Calculate the lifted condensation level (LCL) using from the starting point. @@ -347,6 +363,10 @@ def lcl(pressure, temperature, dewpoint, max_iters=50, eps=1e-5): The function is guaranteed to finish by virtue of the `max_iters` counter. + 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. + """ def _lcl_iter(p, p0, w, t): td = globals()['dewpoint'](vapor_pressure(units.Quantity(p, pressure.units), w)) @@ -364,7 +384,7 @@ def _lcl_iter(p, p0, w, t): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]', '[temperature]') def lfc(pressure, temperature, dewpoint, parcel_temperature_profile=None, dewpoint_start=None, which='top'): @@ -408,6 +428,12 @@ def lfc(pressure, temperature, dewpoint, parcel_temperature_profile=None, dewpoi -------- parcel_profile + 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. + """ pressure, temperature, dewpoint = _remove_nans(pressure, temperature, dewpoint) # Default to surface parcel if no profile or starting pressure level is given @@ -539,7 +565,7 @@ def _most_cape_option(intersect_type, p_list, t_list, pressure, temperature, dew @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]', '[temperature]') def el(pressure, temperature, dewpoint, parcel_temperature_profile=None, which='top'): r"""Calculate the equilibrium level. @@ -577,6 +603,12 @@ def el(pressure, temperature, dewpoint, parcel_temperature_profile=None, which=' -------- parcel_profile + 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. + """ pressure, temperature, dewpoint = _remove_nans(pressure, temperature, dewpoint) # Default to surface parcel if no profile or starting pressure level is given @@ -604,7 +636,7 @@ def el(pressure, temperature, dewpoint, parcel_temperature_profile=None, which=' @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='pressure') @check_units('[pressure]', '[temperature]', '[temperature]') def parcel_profile(pressure, temperature, dewpoint): r"""Calculate the profile a parcel takes through the atmosphere. @@ -632,13 +664,17 @@ def parcel_profile(pressure, temperature, dewpoint): -------- lcl, moist_lapse, dry_lapse + Notes + ----- + Only functions on 1D profiles (not higher-dimension vertical cross sections or grids). + """ _, _, _, t_l, _, t_u = _parcel_profile_helper(pressure, temperature, dewpoint) return concatenate((t_l, t_u)) @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') def parcel_profile_with_lcl(pressure, temperature, dewpoint): r"""Calculate the profile a parcel takes through the atmosphere. @@ -674,7 +710,13 @@ def parcel_profile_with_lcl(pressure, temperature, dewpoint): See Also -------- - lcl, moist_lapse, dry_lapse, parcel_profile + lcl, moist_lapse, dry_lapse, parcel_profile, parcel_profile_with_lcl_as_dataset + + Notes + ----- + Only functions on 1D profiles (not higher-dimension vertical cross sections or grids). + Also, will only return Pint Quantities, even when given xarray DataArray profiles. To + obtain a xarray Dataset instead, use `parcel_profile_with_lcl_as_dataset` instead. """ p_l, p_lcl, p_u, t_l, t_lcl, t_u = _parcel_profile_helper(pressure, temperature[0], @@ -686,6 +728,76 @@ def parcel_profile_with_lcl(pressure, temperature, dewpoint): return new_press, new_temp, new_dewp, prof_temp +@exporter.export +def parcel_profile_with_lcl_as_dataset(pressure, temperature, dewpoint): + r"""Calculate the profile a parcel takes through the atmosphere, returning a Dataset. + + The parcel starts at `temperature`, and `dewpoint`, lifted up + dry adiabatically to the LCL, and then moist adiabatically from there. + `pressure` specifies the pressure levels for the profile. This function returns + a profile that includes the LCL. + + Parameters + ---------- + pressure : `pint.Quantity` + The atmospheric pressure level(s) of interest. This array must be from + high to low pressure. + temperature : `pint.Quantity` + The atmospheric temperature at the levels in `pressure`. The first entry should be at + the same level as the first `pressure` data point. + dewpoint : `pint.Quantity` + The atmospheric dewpoint at the levels in `pressure`. The first entry should be at + the same level as the first `pressure` data point. + + Returns + ------- + profile : `xarray.Dataset` + The interpolated profile with three data variables: ambient_temperature, + ambient_dew_point, and profile_temperature, all of which are on an isobaric + coordinate. + + See Also + -------- + lcl, moist_lapse, dry_lapse, parcel_profile, parcel_profile_with_lcl + + Notes + ----- + Only functions on 1D profiles (not higher-dimension vertical cross sections or grids). + + """ + p, ambient_temperature, ambient_dew_point, profile_temperature = parcel_profile_with_lcl( + pressure, + temperature, + dewpoint + ) + return xr.Dataset( + { + 'ambient_temperature': ( + ('isobaric',), + ambient_temperature, + {'standard_name': 'air_temperature'} + ), + 'ambient_dew_point': ( + ('isobaric',), + ambient_dew_point, + {'standard_name': 'dew_point_temperature'} + ), + 'parcel_temperature': ( + ('isobaric',), + profile_temperature, + {'long_name': 'air_temperature_of_lifted_parcel'} + ) + }, + coords={ + 'isobaric': ( + 'isobaric', + p.m, + {'units': str(p.units), 'standard_name': 'air_pressure'} + ) + } + ) + + def _parcel_profile_helper(pressure, temperature, dewpoint): """Help calculate parcel profiles. @@ -728,7 +840,7 @@ def _insert_lcl_level(pressure, temperature, lcl_pressure): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='mixing_ratio', broadcast=('pressure', 'mixing_ratio')) @check_units('[pressure]', '[dimensionless]') def vapor_pressure(pressure, mixing_ratio): r"""Calculate water vapor (partial) pressure. @@ -765,7 +877,7 @@ def vapor_pressure(pressure, mixing_ratio): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='temperature') @check_units('[temperature]') def saturation_vapor_pressure(temperature): r"""Calculate the saturation water vapor (partial) pressure. @@ -801,7 +913,7 @@ def saturation_vapor_pressure(temperature): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='temperature', broadcast=('temperature', 'relative_humidity')) @check_units('[temperature]', '[dimensionless]') def dewpoint_from_relative_humidity(temperature, relative_humidity): r"""Calculate the ambient dewpoint given air temperature and relative humidity. @@ -829,7 +941,7 @@ def dewpoint_from_relative_humidity(temperature, relative_humidity): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='vapor_pressure') @check_units('[pressure]') def dewpoint(vapor_pressure): r"""Calculate the ambient dewpoint given the vapor pressure. @@ -862,7 +974,7 @@ def dewpoint(vapor_pressure): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='partial_press', broadcast=('partial_press', 'total_press')) @check_units('[pressure]', '[pressure]', '[dimensionless]') def mixing_ratio(partial_press, total_press, molecular_weight_ratio=mpconsts.epsilon): r"""Calculate the mixing ratio of a gas. @@ -904,7 +1016,7 @@ def mixing_ratio(partial_press, total_press, molecular_weight_ratio=mpconsts.eps @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='temperature', broadcast=('total_press', 'temperature')) @check_units('[pressure]', '[temperature]') def saturation_mixing_ratio(total_press, temperature): r"""Calculate the saturation mixing ratio of water vapor. @@ -929,7 +1041,10 @@ def saturation_mixing_ratio(total_press, temperature): @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'dewpoint') +) @check_units('[pressure]', '[temperature]', '[temperature]') def equivalent_potential_temperature(pressure, temperature, dewpoint): r"""Calculate equivalent potential temperature. @@ -982,11 +1097,11 @@ def equivalent_potential_temperature(pressure, temperature, dewpoint): th_l = t * (1000 / (p - e)) ** mpconsts.kappa * (t / t_l) ** (0.28 * r) th_e = th_l * np.exp((3036. / t_l - 1.78) * r * (1 + 0.448 * r)) - return th_e * units.kelvin + return units.Quantity(th_e, units.kelvin) @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='temperature', broadcast=('pressure', 'temperature')) @check_units('[pressure]', '[temperature]') def saturation_equivalent_potential_temperature(pressure, temperature): r"""Calculate saturation equivalent potential temperature. @@ -1049,11 +1164,11 @@ def saturation_equivalent_potential_temperature(pressure, temperature): th_l = t * (1000 / (p - e)) ** mpconsts.kappa th_es = th_l * np.exp((3036. / t - 1.78) * r * (1 + 0.448 * r)) - return th_es * units.kelvin + return units.Quantity(th_es, units.kelvin) @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='temperature', broadcast=('temperature', 'mixing_ratio')) @check_units('[temperature]', '[dimensionless]', '[dimensionless]') def virtual_temperature(temperature, mixing_ratio, molecular_weight_ratio=mpconsts.epsilon): r"""Calculate virtual temperature. @@ -1087,7 +1202,10 @@ def virtual_temperature(temperature, mixing_ratio, molecular_weight_ratio=mpcons @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'mixing_ratio') +) @check_units('[pressure]', '[temperature]', '[dimensionless]', '[dimensionless]') def virtual_potential_temperature(pressure, temperature, mixing_ratio, molecular_weight_ratio=mpconsts.epsilon): @@ -1124,7 +1242,10 @@ def virtual_potential_temperature(pressure, temperature, mixing_ratio, @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'mixing_ratio') +) @check_units('[pressure]', '[temperature]', '[dimensionless]', '[dimensionless]') def density(pressure, temperature, mixing_ratio, molecular_weight_ratio=mpconsts.epsilon): r"""Calculate density. @@ -1160,9 +1281,12 @@ def density(pressure, temperature, mixing_ratio, molecular_weight_ratio=mpconsts @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='dry_bulb_temperature', + broadcast=('pressure', 'dry_bulb_temperature', 'wet_bulb_temperature') +) @check_units('[pressure]', '[temperature]', '[temperature]') -def relative_humidity_wet_psychrometric(pressure, dry_bulb_temperature, web_bulb_temperature, +def relative_humidity_wet_psychrometric(pressure, dry_bulb_temperature, wet_bulb_temperature, **kwargs): r"""Calculate the relative humidity with wet bulb and dry bulb temperatures. @@ -1175,7 +1299,7 @@ def relative_humidity_wet_psychrometric(pressure, dry_bulb_temperature, web_bulb Total atmospheric pressure dry_bulb_temperature: `pint.Quantity` Dry bulb temperature - web_bulb_temperature: `pint.Quantity` + wet_bulb_temperature: `pint.Quantity` Wet bulb temperature Returns @@ -1197,12 +1321,15 @@ def relative_humidity_wet_psychrometric(pressure, dry_bulb_temperature, web_bulb """ return (psychrometric_vapor_pressure_wet(pressure, dry_bulb_temperature, - web_bulb_temperature, **kwargs) + wet_bulb_temperature, **kwargs) / saturation_vapor_pressure(dry_bulb_temperature)) @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='dry_bulb_temperature', + broadcast=('pressure', 'dry_bulb_temperature', 'wet_bulb_temperature') +) @check_units('[pressure]', '[temperature]', '[temperature]') def psychrometric_vapor_pressure_wet(pressure, dry_bulb_temperature, wet_bulb_temperature, psychrometer_coefficient=6.21e-4 / units.kelvin): @@ -1252,7 +1379,10 @@ def psychrometric_vapor_pressure_wet(pressure, dry_bulb_temperature, wet_bulb_te @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'relative_humidity') +) @check_units('[pressure]', '[temperature]', '[dimensionless]') def mixing_ratio_from_relative_humidity(pressure, temperature, relative_humidity): r"""Calculate the mixing ratio from relative humidity, temperature, and pressure. @@ -1292,7 +1422,10 @@ def mixing_ratio_from_relative_humidity(pressure, temperature, relative_humidity @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'mixing_ratio') +) @check_units('[pressure]', '[temperature]', '[dimensionless]') def relative_humidity_from_mixing_ratio(pressure, temperature, mixing_ratio): r"""Calculate the relative humidity from mixing ratio, temperature, and pressure. @@ -1330,7 +1463,7 @@ def relative_humidity_from_mixing_ratio(pressure, temperature, mixing_ratio): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='specific_humidity') @check_units('[dimensionless]') def mixing_ratio_from_specific_humidity(specific_humidity): r"""Calculate the mixing ratio from specific humidity. @@ -1367,7 +1500,7 @@ def mixing_ratio_from_specific_humidity(specific_humidity): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='mixing_ratio') @check_units('[dimensionless]') def specific_humidity_from_mixing_ratio(mixing_ratio): r"""Calculate the specific humidity from the mixing ratio. @@ -1404,7 +1537,10 @@ def specific_humidity_from_mixing_ratio(mixing_ratio): @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'specific_humidity') +) @check_units('[pressure]', '[temperature]', '[dimensionless]') def relative_humidity_from_specific_humidity(pressure, temperature, specific_humidity): r"""Calculate the relative humidity from specific humidity, temperature, and pressure. @@ -1443,7 +1579,7 @@ def relative_humidity_from_specific_humidity(pressure, temperature, specific_hum @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]', '[temperature]') def cape_cin(pressure, temperature, dewpoint, parcel_profile, which_lfc='bottom', which_el='top'): @@ -1499,6 +1635,10 @@ def cape_cin(pressure, temperature, dewpoint, parcel_profile, which_lfc='bottom' * :math:`T_{env}` Environment temperature * :math:`p` Atmospheric pressure + 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. + See Also -------- lfc, el @@ -1593,7 +1733,7 @@ def _find_append_zero_crossings(x, y): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') def most_unstable_parcel(pressure, temperature, dewpoint, height=None, bottom=None, depth=300 * units.hPa): @@ -1631,6 +1771,12 @@ def most_unstable_parcel(pressure, temperature, dewpoint, height=None, -------- get_layer + 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. + """ p_layer, t_layer, td_layer = get_layer(pressure, temperature, dewpoint, bottom=bottom, depth=depth, height=height, interpolate=False) @@ -1640,7 +1786,7 @@ def most_unstable_parcel(pressure, temperature, dewpoint, height=None, @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[temperature]', '[pressure]', '[temperature]') def isentropic_interpolation(levels, pressure, temperature, *args, axis=0, temperature_out=False, max_iters=50, eps=1e-6, @@ -1685,9 +1831,12 @@ def isentropic_interpolation(levels, pressure, temperature, *args, axis=0, [Ziv1994]_. Any additional arguments are assumed to vary linearly with temperature and will be linearly interpolated to the new isentropic levels. + Will only return Pint Quantities, even when given xarray DataArray profiles. To + obtain a xarray Dataset instead, use `isentropic_interpolation_as_dataset` instead. + See Also -------- - potential_temperature + potential_temperature, isentropic_interpolation_as_dataset """ # iteration function to be used later @@ -1787,7 +1936,119 @@ def _isen_iter(iter_log_p, isentlevs_nd, ka, a, b, pok): @exporter.export -@preprocess_xarray +def isentropic_interpolation_as_dataset( + levels, + temperature, + *args, + max_iters=50, + eps=1e-6, + bottom_up_search=True +): + r"""Interpolate xarray data in isobaric coords to isentropic coords, returning a Dataset. + + Parameters + ---------- + levels : `pint.Quantity` + One-dimensional array of desired potential temperature surfaces + temperature : `xarray.DataArray` + Array of temperature + args : `xarray.DataArray`, optional + Any other given variables will be interpolated to each isentropic level. Must have + names in order to have a well-formed output Dataset. + max_iters : int, optional + The maximum number of iterations to use in calculation, defaults to 50. + eps : float, optional + The desired absolute error in the calculated value, defaults to 1e-6. + bottom_up_search : bool, optional + Controls whether to search for levels bottom-up, or top-down. Defaults to + True, which is bottom-up search. + + Returns + ------- + xarray.Dataset + Dataset with pressure, temperature, and each additional argument, all on the specified + isentropic coordinates. + + Notes + ----- + Input variable arrays must have the same number of vertical levels as the pressure levels + array. Pressure is calculated on isentropic surfaces by assuming that temperature varies + linearly with the natural log of pressure. Linear interpolation is then used in the + vertical to find the pressure at each isentropic level. Interpolation method from + [Ziv1994]_. Any additional arguments are assumed to vary linearly with temperature and will + be linearly interpolated to the new isentropic levels. + + This formulation relies upon xarray functionality. If using Pint Quantities, use + `isentropic_interpolation` instead. + + See Also + -------- + potential_temperature, isentropic_interpolation + + """ + # Ensure matching coordinates by broadcasting + all_args = xr.broadcast(temperature, *args) + + # Obtain result as list of Quantities + ret = isentropic_interpolation( + levels, + all_args[0].metpy.vertical, + all_args[0].metpy.unit_array, + *(arg.metpy.unit_array for arg in all_args[1:]), + axis=all_args[0].metpy.find_axis_number('vertical'), + temperature_out=True, + max_iters=max_iters, + eps=eps, + bottom_up_search=bottom_up_search + ) + + # Reconstruct coordinates and dims (add isentropic levels, remove isobaric levels) + vertical_dim = all_args[0].metpy.find_axis_name('vertical') + new_coords = { + 'isentropic_level': xr.DataArray( + levels.m, + dims=('isentropic_level',), + coords=levels.m, + name='isentropic_level', + attrs={ + 'units': str(levels.units), + 'positive': 'up' + } + ), + **{ + key: value + for key, value in all_args[0].coords + if key != vertical_dim + } + } + new_dims = [ + dim if dim != vertical_dim else 'isentropic_level' for dim in all_args[0].dims + ] + + # Build final dataset from interpolated Quantities and original DataArrays + return xr.Dataset( + { + 'pressure': ( + new_dims, + ret[0], + {'standard_name': 'air_pressure'} + ), + 'temperature': ( + new_dims, + ret[-1], + {'standard_name': 'air_temperature'} + ), + **{ + all_args[i].name: (new_dims, ret[i], all_args[i].attrs) + for i in range(1, len(all_args)) + } + }, + coords=new_coords + ) + + +@exporter.export +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') def surface_based_cape_cin(pressure, temperature, dewpoint): r"""Calculate surface-based CAPE and CIN. @@ -1819,6 +2080,12 @@ def surface_based_cape_cin(pressure, temperature, dewpoint): -------- cape_cin, parcel_profile + 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. + """ pressure, temperature, dewpoint = _remove_nans(pressure, temperature, dewpoint) p, t, td, profile = parcel_profile_with_lcl(pressure, temperature, dewpoint) @@ -1826,7 +2093,7 @@ def surface_based_cape_cin(pressure, temperature, dewpoint): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') def most_unstable_cape_cin(pressure, temperature, dewpoint, **kwargs): r"""Calculate most unstable CAPE/CIN. @@ -1859,6 +2126,12 @@ def most_unstable_cape_cin(pressure, temperature, dewpoint, **kwargs): -------- cape_cin, most_unstable_parcel, parcel_profile + 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. + """ pressure, temperature, dewpoint = _remove_nans(pressure, temperature, dewpoint) _, _, _, parcel_idx = most_unstable_parcel(pressure, temperature, dewpoint, **kwargs) @@ -1869,7 +2142,7 @@ def most_unstable_cape_cin(pressure, temperature, dewpoint, **kwargs): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') def mixed_layer_cape_cin(pressure, temperature, dewpoint, **kwargs): r"""Calculate mixed-layer CAPE and CIN. @@ -1902,6 +2175,13 @@ def mixed_layer_cape_cin(pressure, temperature, dewpoint, **kwargs): See Also -------- cape_cin, mixed_parcel, parcel_profile + + 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. + """ depth = kwargs.get('depth', 100 * units.hPa) parcel_pressure, parcel_temp, parcel_dewpoint = mixed_parcel(pressure, temperature, @@ -1920,7 +2200,7 @@ def mixed_layer_cape_cin(pressure, temperature, dewpoint, **kwargs): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') def mixed_parcel(pressure, temperature, dewpoint, parcel_start_pressure=None, height=None, bottom=None, depth=100 * units.hPa, interpolate=True): @@ -1959,6 +2239,12 @@ def mixed_parcel(pressure, temperature, dewpoint, parcel_start_pressure=None, `pint.Quantity` The dewpoint of the mixed parcel + 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. + """ # If a parcel starting pressure is not provided, use the surface if not parcel_start_pressure: @@ -1988,7 +2274,7 @@ def mixed_parcel(pressure, temperature, dewpoint, parcel_start_pressure=None, @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]') def mixed_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa, interpolate=True): @@ -2019,6 +2305,12 @@ def mixed_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa `pint.Quantity` The mixed value of the data variable. + 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. + """ layer = get_layer(pressure, *args, height=height, bottom=bottom, depth=depth, interpolate=interpolate) @@ -2034,7 +2326,7 @@ def mixed_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='temperature', broadcast=('height', 'temperature')) @check_units('[length]', '[temperature]') def dry_static_energy(height, temperature): r"""Calculate the dry static energy of parcels. @@ -2066,7 +2358,10 @@ def dry_static_energy(height, temperature): @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('height', 'temperature', 'specific_humidity') +) @check_units('[length]', '[temperature]', '[dimensionless]') def moist_static_energy(height, temperature, specific_humidity): r"""Calculate the moist static energy of parcels. @@ -2102,7 +2397,7 @@ def moist_static_energy(height, temperature, specific_humidity): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]') def thickness_hydrostatic(pressure, temperature, mixing_ratio=None, molecular_weight_ratio=mpconsts.epsilon, bottom=None, depth=None): @@ -2147,6 +2442,12 @@ def thickness_hydrostatic(pressure, temperature, mixing_ratio=None, -------- thickness_hydrostatic_from_relative_humidity, pressure_to_height_std, virtual_temperature + 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. + """ # Get the data for the layer, conditional upon bottom/depth being specified and mixing # ratio being given @@ -2172,7 +2473,7 @@ def thickness_hydrostatic(pressure, temperature, mixing_ratio=None, @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]') def thickness_hydrostatic_from_relative_humidity(pressure, temperature, relative_humidity, bottom=None, depth=None): @@ -2217,6 +2518,12 @@ def thickness_hydrostatic_from_relative_humidity(pressure, temperature, relative thickness_hydrostatic, pressure_to_height_std, virtual_temperature, mixing_ratio_from_relative_humidity + 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. + """ mixing = mixing_ratio_from_relative_humidity(pressure, temperature, relative_humidity) @@ -2225,7 +2532,6 @@ def thickness_hydrostatic_from_relative_humidity(pressure, temperature, relative @exporter.export -@preprocess_xarray @check_units('[length]', '[temperature]') def brunt_vaisala_frequency_squared(height, potential_temperature, axis=0): r"""Calculate the square of the Brunt-Vaisala frequency. @@ -2239,33 +2545,51 @@ def brunt_vaisala_frequency_squared(height, potential_temperature, axis=0): Parameters ---------- - height : `pint.Quantity` - One-dimensional profile of atmospheric height - potential_temperature : `pint.Quantity` + height : `xarray.DataArray` or `pint.Quantity` + Atmospheric (geopotential) height + potential_temperature : `xarray.DataArray` or `pint.Quantity` Atmospheric potential temperature axis : int, optional - The axis corresponding to vertical in the potential temperature array, defaults to 0. + The axis corresponding to vertical in the potential temperature array, defaults to 0, + unless `height` and `potential_temperature` given as `xarray.DataArray`, in which case + it is automatically determined from the coordinate metadata. Returns ------- - `pint.Quantity` - The square of the Brunt-Vaisala frequency. + `pint.Quantity` or `xarray.DataArray` + The square of the Brunt-Vaisala frequency. Given as `pint.Quantity`, unless both + `height` and `potential_temperature` arguments are given as `xarray.DataArray`, in + which case will be `xarray.DataArray`. See Also -------- brunt_vaisala_frequency, brunt_vaisala_period, potential_temperature """ - # Ensure validity of temperature units - potential_temperature = potential_temperature.to('K') + if isinstance(height, xr.DataArray) and isinstance(potential_temperature, xr.DataArray): + # Prepare arguments when DataArrays + potential_temperature = potential_temperature.metpy.convert_units('K') + height, potential_temperature = xr.broadcast( + height.metpy.quantify(), + potential_temperature.metpy.quantify() + ) + dtheta_dz = first_derivative( + potential_temperature, + x=height, + axis=potential_temperature.metpy.find_axis_number('vertical') + ) + else: + if isinstance(potential_temperature, xr.DataArray): + potential_temperature = potential_temperature.metpy.unit_array.to('K') + else: + potential_temperature = potential_temperature.to('K') + dtheta_dz = first_derivative(potential_temperature, x=height, axis=axis) # Calculate and return the square of Brunt-Vaisala frequency - return mpconsts.g / potential_temperature * first_derivative(potential_temperature, - x=height, axis=axis) + return mpconsts.g / potential_temperature * dtheta_dz @exporter.export -@preprocess_xarray @check_units('[length]', '[temperature]') def brunt_vaisala_frequency(height, potential_temperature, axis=0): r"""Calculate the Brunt-Vaisala frequency. @@ -2281,17 +2605,21 @@ def brunt_vaisala_frequency(height, potential_temperature, axis=0): Parameters ---------- - height : `pint.Quantity` - One-dimensional profile of atmospheric height - potential_temperature : `pint.Quantity` + height : `xarray.DataArray` or `pint.Quantity` + Atmospheric (geopotential) height + potential_temperature : `xarray.DataArray` or `pint.Quantity` Atmospheric potential temperature axis : int, optional - The axis corresponding to vertical in the potential temperature array, defaults to 0. + The axis corresponding to vertical in the potential temperature array, defaults to 0, + unless `height` and `potential_temperature` given as `xarray.DataArray`, in which case + it is automatically determined from the coordinate metadata. Returns ------- - `pint.Quantity` - Brunt-Vaisala frequency. + `pint.Quantity` or `xarray.DataArray` + Brunt-Vaisala frequency. Given as `pint.Quantity`, unless both + `height` and `potential_temperature` arguments are given as `xarray.DataArray`, in + which case will be `xarray.DataArray`. See Also -------- @@ -2300,13 +2628,15 @@ def brunt_vaisala_frequency(height, potential_temperature, axis=0): """ bv_freq_squared = brunt_vaisala_frequency_squared(height, potential_temperature, axis=axis) - bv_freq_squared[bv_freq_squared.magnitude < 0] = np.nan + if isinstance(bv_freq_squared, xr.DataArray): + bv_freq_squared.data[bv_freq_squared.data.magnitude < 0] = np.nan + else: + bv_freq_squared[bv_freq_squared.magnitude < 0] = np.nan return np.sqrt(bv_freq_squared) @exporter.export -@preprocess_xarray @check_units('[length]', '[temperature]') def brunt_vaisala_period(height, potential_temperature, axis=0): r"""Calculate the Brunt-Vaisala period. @@ -2320,17 +2650,22 @@ def brunt_vaisala_period(height, potential_temperature, axis=0): Parameters ---------- - height : `pint.Quantity` - One-dimensional profile of atmospheric height - potential_temperature : pint.Quantity` + height : `xarray.DataArray` or `pint.Quantity` + Atmospheric (geopotential) height + potential_temperature : `xarray.DataArray` or `pint.Quantity` Atmospheric potential temperature axis : int, optional - The axis corresponding to vertical in the potential temperature array, defaults to 0. + The axis corresponding to vertical in the potential temperature array, defaults to 0, + unless `height` and `potential_temperature` given as `xarray.DataArray`, in which case + it is automatically determined from the coordinate metadata. Returns ------- - `pint.Quantity` - Brunt-Vaisala period. + `pint.Quantity` or `xarray.DataArray` + Brunt-Vaisala period. Given as `pint.Quantity`, unless both + `height` and `potential_temperature` arguments are given as `xarray.DataArray`, in + which case will be `xarray.DataArray`. + See Also -------- @@ -2339,13 +2674,19 @@ def brunt_vaisala_period(height, potential_temperature, axis=0): """ bv_freq_squared = brunt_vaisala_frequency_squared(height, potential_temperature, axis=axis) - bv_freq_squared[bv_freq_squared.magnitude <= 0] = np.nan + if isinstance(bv_freq_squared, xr.DataArray): + bv_freq_squared.data[bv_freq_squared.data.magnitude <= 0] = np.nan + else: + bv_freq_squared[bv_freq_squared.magnitude <= 0] = np.nan return 2 * np.pi / np.sqrt(bv_freq_squared) @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'dewpoint') +) @check_units('[pressure]', '[temperature]', '[temperature]') def wet_bulb_temperature(pressure, temperature, dewpoint): """Calculate the wet-bulb temperature using Normand's rule. @@ -2372,6 +2713,11 @@ def wet_bulb_temperature(pressure, temperature, dewpoint): -------- lcl, moist_lapse + Notes + ----- + Since this function iteratively applies a parcel calculation, it should be used with + caution on large arrays. + """ if not hasattr(pressure, 'shape'): pressure = np.atleast_1d(pressure) @@ -2398,7 +2744,7 @@ def wet_bulb_temperature(pressure, temperature, dewpoint): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='temperature', broadcast=('pressure', 'temperature')) @check_units('[pressure]', '[temperature]') def static_stability(pressure, temperature, axis=0): r"""Calculate the static stability within a vertical profile. @@ -2430,7 +2776,10 @@ def static_stability(pressure, temperature, axis=0): @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='temperature', + broadcast=('pressure', 'temperature', 'specific_humdiity') +) @check_units('[pressure]', '[temperature]', '[dimensionless]') def dewpoint_from_specific_humidity(pressure, temperature, specific_humidity): r"""Calculate the dewpoint from specific humidity, temperature, and pressure. @@ -2460,7 +2809,7 @@ def dewpoint_from_specific_humidity(pressure, temperature, specific_humidity): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='w', broadcast=('w', 'pressure', 'temperature')) @check_units('[length]/[time]', '[pressure]', '[temperature]') def vertical_velocity_pressure(w, pressure, temperature, mixing_ratio=0): r"""Calculate omega from w assuming hydrostatic conditions. @@ -2503,7 +2852,10 @@ def vertical_velocity_pressure(w, pressure, temperature, mixing_ratio=0): @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='omega', + broadcast=('omega', 'pressure', 'temperature', 'mixing_ratio') +) @check_units('[pressure]/[time]', '[pressure]', '[temperature]') def vertical_velocity(omega, pressure, temperature, mixing_ratio=0): r"""Calculate w from omega assuming hydrostatic conditions. @@ -2549,7 +2901,7 @@ def vertical_velocity(omega, pressure, temperature, mixing_ratio=0): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='dewpoint', broadcast=('dewpoint', 'pressure')) @check_units('[pressure]', '[temperature]') def specific_humidity_from_dewpoint(pressure, dewpoint): r"""Calculate the specific humidity from the dewpoint temperature and pressure. @@ -2577,7 +2929,7 @@ def specific_humidity_from_dewpoint(pressure, dewpoint): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]', '[temperature]', '[temperature]') def lifted_index(pressure, temperature, parcel_profile): """Calculate Lifted Index from the pressure temperature and parcel profile. @@ -2620,7 +2972,10 @@ def lifted_index(pressure, temperature, parcel_profile): @exporter.export -@preprocess_xarray +@preprocess_and_wrap( + wrap_like='potential_temperature', + broadcast=('height', 'potential_temperature', 'u', 'v') +) @check_units('[length]', '[temperature]', '[speed]', '[speed]') def gradient_richardson_number(height, potential_temperature, u, v, axis=0): r"""Calculate the gradient (or flux) Richardson number. diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index f6e450a02f0..53d3abb64e4 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -3,7 +3,6 @@ # SPDX-License-Identifier: BSD-3-Clause """Contains a collection of generally useful calculation tools.""" import functools -from inspect import signature from operator import itemgetter import warnings @@ -17,7 +16,7 @@ from ..interpolate import interpolate_1d, log_interpolate_1d from ..package_tools import Exporter from ..units import check_units, concatenate, units -from ..xarray import check_axis, preprocess_xarray +from ..xarray import check_axis, preprocess_and_wrap exporter = Exporter(globals()) @@ -39,21 +38,21 @@ @exporter.export -@preprocess_xarray def resample_nn_1d(a, centers): """Return one-dimensional nearest-neighbor indexes based on user-specified centers. Parameters ---------- a : array-like - 1-dimensional array of numeric values from which to - extract indexes of nearest-neighbors + 1-dimensional array of numeric values from which to extract indexes of + nearest-neighbors centers : array-like 1-dimensional array of numeric values representing a subset of values to approximate Returns ------- - An array of indexes representing values closest to given array values + A list of indexes (in type given by `array.argmin()`) representing values closest to + given array values. """ ix = [] @@ -65,7 +64,6 @@ def resample_nn_1d(a, centers): @exporter.export -@preprocess_xarray def nearest_intersection_idx(a, b): """Determine the index of the point just before two lines with common x values. @@ -93,7 +91,7 @@ def nearest_intersection_idx(a, b): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @units.wraps(('=A', '=B'), ('=A', '=B', '=B', None, None)) def find_intersections(x, a, b, direction='all', log_x=False): """Calculate the best estimate of intersection. @@ -121,6 +119,11 @@ def find_intersections(x, a, b, direction='all', log_x=False): A tuple (x, y) of array-like with the x and y coordinates of the intersections of the lines. + Notes + ----- + This function implicity converts `xarray.DataArray` to `pint.Quantity`, with the results + given as `pint.Quantity`. + """ # Change x to logarithmic if log_x=True if log_x is True: @@ -233,7 +236,7 @@ def _delete_masked_points(*arrs): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() def reduce_point_density(points, radius, priority=None): r"""Return a mask to reduce the density of points in irregularly-spaced data. @@ -417,7 +420,7 @@ def _get_bound_pressure_height(pressure, bound, height=None, interpolate=True): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[length]') def get_layer_heights(height, depth, *args, bottom=None, interpolate=True, with_agl=False): """Return an atmospheric layer from upper air data with the requested bottom and depth. @@ -447,6 +450,11 @@ def get_layer_heights(height, depth, *args, bottom=None, interpolate=True, with_ `pint.Quantity, pint.Quantity` The height and data variables of the layer + Notes + ----- + Only functions on 1D profiles (not higher-dimension vertical cross sections or grids). + Also, this will return Pint Quantities even when given xarray DataArray profiles. + """ # Make sure pressure and datavars are the same length for datavar in args: @@ -507,7 +515,7 @@ def get_layer_heights(height, depth, *args, bottom=None, interpolate=True, with_ @exporter.export -@preprocess_xarray +@preprocess_and_wrap() @check_units('[pressure]') def get_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa, interpolate=True): @@ -543,6 +551,11 @@ def get_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa, `pint.Quantity, pint.Quantity` The pressure and data variables of the layer + Notes + ----- + Only functions on 1D profiles (not higher-dimension vertical cross sections or grids). + Also, this will return Pint Quantities even when given xarray DataArray profiles. + """ # If we get the depth kwarg, but it's None, set it to the default as well if depth is None: @@ -609,7 +622,7 @@ def get_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa, @exporter.export -@preprocess_xarray +@preprocess_and_wrap() def find_bounding_indices(arr, values, axis, from_below=True): """Find the indices surrounding the values within arr along axis. @@ -750,7 +763,7 @@ def take(indexer): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() def lat_lon_grid_deltas(longitude, latitude, y_dim=-2, x_dim=-1, **kwargs): r"""Calculate the actual delta between grid points that are in latitude/longitude format. @@ -781,6 +794,9 @@ def lat_lon_grid_deltas(longitude, latitude, y_dim=-2, x_dim=-1, **kwargs): Assumes [..., Y, X] dimension order for input and output, unless keyword arguments `y_dim` and `x_dim` are otherwise specified. + This function will only return `pint.Quantity` arrays (not `xarray.DataArray` or another + array-like type). It will also "densify" your data if using Dask or lazy-loading. + """ from pyproj import Geod @@ -794,8 +810,8 @@ def lat_lon_grid_deltas(longitude, latitude, y_dim=-2, x_dim=-1, **kwargs): # pyproj requires ndarrays, not Quantities try: - longitude = longitude.m_as('degrees') - latitude = latitude.m_as('degrees') + longitude = np.asarray(longitude.m_as('degrees')) + latitude = np.asarray(latitude.m_as('degrees')) except AttributeError: longitude = np.asarray(longitude) latitude = np.asarray(latitude) @@ -905,7 +921,7 @@ def xarray_derivative_wrap(func): def wrapper(f, **kwargs): if 'x' in kwargs or 'delta' in kwargs: # Use the usual DataArray to pint.Quantity preprocessing wrapper - return preprocess_xarray(func)(f, **kwargs) + return preprocess_and_wrap()(func)(f, **kwargs) elif isinstance(f, xr.DataArray): # Get axis argument, defaulting to first dimension axis = f.metpy.find_axis_name(kwargs.get('axis', 0)) @@ -941,7 +957,7 @@ def wrapper(f, **kwargs): @exporter.export @xarray_derivative_wrap -def first_derivative(f, **kwargs): +def first_derivative(f, axis=None, x=None, delta=None): """Calculate the first derivative of a grid of values. Works for both regularly-spaced data and grids with varying spacing. @@ -984,7 +1000,7 @@ def first_derivative(f, **kwargs): second_derivative """ - n, axis, delta = _process_deriv_args(f, kwargs) + n, axis, delta = _process_deriv_args(f, axis, x, delta) take = make_take(n, axis) # First handle centered case @@ -1031,7 +1047,7 @@ def first_derivative(f, **kwargs): @exporter.export @xarray_derivative_wrap -def second_derivative(f, **kwargs): +def second_derivative(f, axis=None, x=None, delta=None): """Calculate the second derivative of a grid of values. Works for both regularly-spaced data and grids with varying spacing. @@ -1074,7 +1090,7 @@ def second_derivative(f, **kwargs): first_derivative """ - n, axis, delta = _process_deriv_args(f, kwargs) + n, axis, delta = _process_deriv_args(f, axis, x, delta) take = make_take(n, axis) # First handle centered case @@ -1117,7 +1133,7 @@ def second_derivative(f, **kwargs): @exporter.export -def gradient(f, **kwargs): +def gradient(f, coordinates=None, deltas=None, axes=None): """Calculate the gradient of a grid of values. Works for both regularly-spaced data, and grids with varying spacing. @@ -1164,13 +1180,13 @@ def gradient(f, **kwargs): `deltas` (as applicable) should match the number of dimensions of `f`. """ - pos_kwarg, positions, axes = _process_gradient_args(f, kwargs) + pos_kwarg, positions, axes = _process_gradient_args(f, axes, coordinates, deltas) return tuple(first_derivative(f, axis=axis, **{pos_kwarg: positions[ind]}) for ind, axis in enumerate(axes)) @exporter.export -def laplacian(f, **kwargs): +def laplacian(f, coordinates=None, deltas=None, axes=None): """Calculate the laplacian of a grid of values. Works for both regularly-spaced data, and grids with varying spacing. @@ -1215,7 +1231,7 @@ def laplacian(f, **kwargs): `deltas` (as applicable) should match the number of dimensions of `f`. """ - pos_kwarg, positions, axes = _process_gradient_args(f, kwargs) + pos_kwarg, positions, axes = _process_gradient_args(f, axes, coordinates, deltas) derivs = [second_derivative(f, axis=axis, **{pos_kwarg: positions[ind]}) for ind, axis in enumerate(axes)] laplac = sum(derivs) @@ -1237,26 +1253,27 @@ def _broadcast_to_axis(arr, axis, ndim): return arr -def _process_gradient_args(f, kwargs): +def _process_gradient_args(f, axes, coordinates, deltas): """Handle common processing of arguments for gradient and gradient-like functions.""" - axes = kwargs.get('axes', range(f.ndim)) + axes_given = axes is not None + axes = axes if axes_given else range(f.ndim) def _check_length(positions): - if 'axes' in kwargs and len(positions) < len(axes): + if axes_given and len(positions) < len(axes): raise ValueError('Length of "coordinates" or "deltas" cannot be less than that ' 'of "axes".') - elif 'axes' not in kwargs and len(positions) != len(axes): + elif not axes_given and len(positions) != len(axes): raise ValueError('Length of "coordinates" or "deltas" must match the number of ' 'dimensions of "f" when "axes" is not given.') - if 'deltas' in kwargs: - if 'coordinates' in kwargs or 'x' in kwargs: + if deltas is not None: + if coordinates is not None: raise ValueError('Cannot specify both "coordinates" and "deltas".') - _check_length(kwargs['deltas']) - return 'delta', kwargs['deltas'], axes - elif 'coordinates' in kwargs: - _check_length(kwargs['coordinates']) - return 'x', kwargs['coordinates'], axes + _check_length(deltas) + return 'delta', deltas, axes + elif coordinates is not None: + _check_length(coordinates) + return 'x', coordinates, axes elif isinstance(f, xr.DataArray): return 'pass', axes, axes # only the axis argument matters else: @@ -1264,19 +1281,19 @@ def _check_length(positions): 'when "f" is not a DataArray.') -def _process_deriv_args(f, kwargs): +def _process_deriv_args(f, axis, x, delta): """Handle common processing of arguments for derivative functions.""" n = f.ndim - axis = normalize_axis_index(kwargs.get('axis', 0), n) + axis = normalize_axis_index(axis if axis is not None else 0, n) if f.shape[axis] < 3: raise ValueError('f must have at least 3 point along the desired axis.') - if 'delta' in kwargs: - if 'x' in kwargs: + if delta is not None: + if x is not None: raise ValueError('Cannot specify both "x" and "delta".') - delta = np.atleast_1d(kwargs['delta']) + delta = np.atleast_1d(delta) if delta.size == 1: diff_size = list(f.shape) diff_size[axis] -= 1 @@ -1286,8 +1303,8 @@ def _process_deriv_args(f, kwargs): delta = delta * delta_units else: delta = _broadcast_to_axis(delta, axis, n) - elif 'x' in kwargs: - x = _broadcast_to_axis(kwargs['x'], axis, n) + elif x is not None: + x = _broadcast_to_axis(x, axis, n) delta = np.diff(x, axis=axis) else: raise ValueError('Must specify either "x" or "delta" for value positions.') @@ -1296,7 +1313,7 @@ def _process_deriv_args(f, kwargs): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='input_dir') def parse_angle(input_dir): """Calculate the meteorological angle from directional text. @@ -1352,7 +1369,7 @@ def _abbrieviate_direction(ext_dir_str): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() def angle_to_direction(input_angle, full=False, level=3): """Convert the meteorological angle to directional text. @@ -1475,146 +1492,3 @@ def _remove_nans(*variables): for v in variables: ret.append(v[~mask]) return ret - - -def wrap_output_like(**wrap_kwargs): - """Wrap the output from a function to be like some other data object type. - - Wraps given data to match the units/coordinates/object type of another array. Currently - supports: - - - As input (output from wrapped function): - - * ``pint.Quantity`` - * ``xarray.DataArray`` - * any type wrappable by ``pint.Quantity`` - - - As matched output (final returned value): - - * ``pint.Quantity`` - * ``xarray.DataArray`` wrapping a ``pint.Quantity`` - - (if matched output is not one of these types, we instead treat the match as if it was a - dimenionless Quantity.) - - This wrapping/conversion follows the following rules: - - - If match_unit is False, for output of Quantity or DataArary respectively, - - * ndarray becomes dimensionless Quantity or unitless DataArray with matching coords - * Quantity is unchanged or becomes DataArray with input units and output coords - * DataArray is converted to Quantity by accessor or is unchanged - - - If match_unit is True, for output of Quantity or DataArary respectively, with a given - unit, - - * ndarray becomes Quantity or DataArray (with matching coords) with output unit - * Quantity is converted to output unit, then returned or converted to DataArray with - matching coords - * DataArray is has units converted via the accessor, then converted to Quantity via - the accessor or returned - - The output to match can be specified two ways: - - - Using the `argument` keyword argument, the output is taken from argument of that name - from the wrapped function's signature - - Using the `other` keyword argument, the output is given directly - - Parameters - ---------- - argument : str - specify the name of a single argument from the function signature from which - to take the other data object - other : `numpy.ndarray` or `pint.Quantity` or `xarray.DataArray` - specify the other data object directly - match_unit : bool - if True and other data object has units, convert output to those units - (defaults to False) - - Notes - ----- - This can be extended in the future to support: - - - ``units.wraps``-like behavior - - Python scalars vs. NumPy scalars (Issue #1209) - - dask (and other duck array) compatibility - - dimensionality reduction (particularly with xarray) - - See Also - -------- - preprocess_xarray - - """ - def decorator(func): - @functools.wraps(func) - def wrapper(*args, **kwargs): - # Determine other - if 'other' in wrap_kwargs: - other = wrap_kwargs['other'] - elif 'argument' in wrap_kwargs: - other = signature(func).bind(*args, **kwargs).arguments[ - wrap_kwargs['argument']] - else: - raise ValueError('Must specify keyword "other" or "argument".') - - # Get result from wrapped function - result = func(*args, **kwargs) - - # Proceed with wrapping rules - if wrap_kwargs.get('match_unit', False): - return _wrap_output_like_matching_units(result, other) - else: - return _wrap_output_like_not_matching_units(result, other) - - return wrapper - return decorator - - -def _wrap_output_like_matching_units(result, match): - """Convert result to be like match with matching units for output wrapper.""" - if isinstance(match, xr.DataArray): - output_xarray = True - match_units = match.metpy.units - elif isinstance(match, units.Quantity): - output_xarray = False - match_units = match.units - else: - output_xarray = False - match_units = '' - - if isinstance(result, xr.DataArray): - result = result.metpy.convert_units(match_units) - return result.metpy.quantify() if output_xarray else result.metpy.unit_array - else: - result = result.m_as(match_units) if isinstance(result, units.Quantity) else result - if output_xarray: - return xr.DataArray( - units.Quantity(result, match_units), - dims=match.dims, - coords=match.coords - ) - else: - return units.Quantity(result, match_units) - - -def _wrap_output_like_not_matching_units(result, match): - """Convert result to be like match without matching units for output wrapper.""" - output_xarray = isinstance(match, xr.DataArray) - if isinstance(result, xr.DataArray): - return result.metpy.quantify() if output_xarray else result.metpy.unit_array - else: - if isinstance(result, units.Quantity): - result_magnitude = result.magnitude - result_units = str(result.units) - else: - result_magnitude = result - result_units = '' - - if output_xarray: - return xr.DataArray( - units.Quantity(result_magnitude, result_units), - dims=match.dims, - coords=match.coords - ) - else: - return units.Quantity(result_magnitude, result_units) diff --git a/src/metpy/calc/turbulence.py b/src/metpy/calc/turbulence.py index fa55bb46f51..461ba0eeba1 100644 --- a/src/metpy/calc/turbulence.py +++ b/src/metpy/calc/turbulence.py @@ -7,13 +7,13 @@ from .tools import make_take from ..package_tools import Exporter -from ..xarray import preprocess_xarray +from ..xarray import preprocess_and_wrap exporter = Exporter(globals()) @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='ts') def get_perturbation(ts, axis=-1): r"""Compute the perturbation from the mean of a time series. @@ -46,7 +46,7 @@ def get_perturbation(ts, axis=-1): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='u') def tke(u, v, w, perturbation=False, axis=-1): r"""Compute turbulence kinetic energy. @@ -112,7 +112,7 @@ def tke(u, v, w, perturbation=False, axis=-1): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='vel') def kinematic_flux(vel, b, perturbation=False, axis=-1): r"""Compute the kinematic flux from two time series. @@ -181,7 +181,7 @@ def kinematic_flux(vel, b, perturbation=False, axis=-1): @exporter.export -@preprocess_xarray +@preprocess_and_wrap(wrap_like='u') def friction_velocity(u, w, v=None, perturbation=False, axis=-1): r"""Compute the friction velocity from the time series of velocity components. diff --git a/src/metpy/interpolate/one_dimension.py b/src/metpy/interpolate/one_dimension.py index 0e9e85fceff..6a5d83fe319 100644 --- a/src/metpy/interpolate/one_dimension.py +++ b/src/metpy/interpolate/one_dimension.py @@ -8,13 +8,13 @@ from ..cbook import broadcast_indices from ..package_tools import Exporter -from ..xarray import preprocess_xarray +from ..xarray import preprocess_and_wrap exporter = Exporter(globals()) @exporter.export -@preprocess_xarray +@preprocess_and_wrap() def interpolate_nans_1d(x, y, kind='linear'): """Interpolate NaN values in y. @@ -49,7 +49,7 @@ def interpolate_nans_1d(x, y, kind='linear'): @exporter.export -@preprocess_xarray +@preprocess_and_wrap() def interpolate_1d(x, xp, *args, axis=0, fill_value=np.nan, return_list_always=False): r"""Interpolates data with any shape over a specified axis. @@ -174,7 +174,7 @@ def interpolate_1d(x, xp, *args, axis=0, fill_value=np.nan, return_list_always=F @exporter.export -@preprocess_xarray +@preprocess_and_wrap() def log_interpolate_1d(x, xp, *args, axis=0, fill_value=np.nan): r"""Interpolates data with logarithmic x-scale over a specified axis. diff --git a/src/metpy/plots/_mpl.py b/src/metpy/plots/_mpl.py index a65b246bcc7..a7848c92022 100644 --- a/src/metpy/plots/_mpl.py +++ b/src/metpy/plots/_mpl.py @@ -100,6 +100,12 @@ def scattertext(self, x, y, texts, loc=(0, 0), **kw): # Add it to the axes and update range self.add_artist(text_obj) + + # Matplotlib at least up to 3.2.2 does not properly clip text with paths, so + # work-around by setting to the bounding box of the Axes + # TODO: Remove when fixed in our minimum supported version of matplotlib + text_obj.clipbox = self.bbox + self.update_datalim(text_obj.get_datalim(self.transData)) self.autoscale_view() return text_obj diff --git a/src/metpy/plots/cartopy_utils.py b/src/metpy/plots/cartopy_utils.py index 0140946f990..dd7c5891927 100644 --- a/src/metpy/plots/cartopy_utils.py +++ b/src/metpy/plots/cartopy_utils.py @@ -3,28 +3,39 @@ # SPDX-License-Identifier: BSD-3-Clause """Cartopy specific mapping utilities.""" +import cartopy.crs as ccrs import cartopy.feature as cfeature from ..cbook import get_test_data -class MetPyMapFeature(cfeature.NaturalEarthFeature): - """A simple interface to US County shapefiles.""" +class MetPyMapFeature(cfeature.Feature): + """A simple interface to MetPy-included shapefiles.""" def __init__(self, name, scale, **kwargs): - """Create USCountiesFeature instance.""" - super().__init__('', name, scale, **kwargs) + """Create MetPyMapFeature instance.""" + super().__init__(ccrs.PlateCarree(), **kwargs) + self.name = name + + if isinstance(scale, str): + scale = cfeature.Scaler(scale) + self.scaler = scale def geometries(self): """Return an iterator of (shapely) geometries for this feature.""" import cartopy.io.shapereader as shapereader # Ensure that the associated files are in the cache - fname = '{}_{}'.format(self.name, self.scale) + fname = '{}_{}'.format(self.name, self.scaler.scale) for extension in ['.dbf', '.shx']: get_test_data(fname + extension) path = get_test_data(fname + '.shp', as_file_obj=False) return iter(tuple(shapereader.Reader(path).geometries())) + def intersecting_geometries(self, extent): + """Return geometries that intersect the extent.""" + self.scaler.scale_from_extent(extent) + return super().intersecting_geometries(extent) + def with_scale(self, new_scale): """ Return a copy of the feature with a new scale. diff --git a/src/metpy/plots/wx_symbols.py b/src/metpy/plots/wx_symbols.py index 8808c6182ff..6a0ac0a6817 100644 --- a/src/metpy/plots/wx_symbols.py +++ b/src/metpy/plots/wx_symbols.py @@ -106,13 +106,13 @@ def __init__(self, num, font_start, font_jumps=None, char_jumps=None): font_point += 1 @staticmethod - def _safe_pop(l): + def _safe_pop(li): """Safely pop from a list. Returns None if list empty. """ - return l.pop(0) if l else None + return li.pop(0) if li else None def __call__(self, code): """Return the Unicode code point corresponding to `code`.""" diff --git a/src/metpy/units.py b/src/metpy/units.py index dde51b5f176..61c9bb5e314 100644 --- a/src/metpy/units.py +++ b/src/metpy/units.py @@ -180,9 +180,9 @@ def _check_argument_units(args, defaults, dimensionality): # Argument did not have units specified in decorator continue - if arg in defaults: + if arg in defaults and (defaults[arg] is not None or val is None): check = val == defaults[arg] - if check: + if isinstance(check, bool) and check: continue # See if the value passed in is appropriate diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index f01b5832578..6dc50b82984 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -16,6 +16,7 @@ See Also: :doc:`xarray with MetPy Tutorial `. """ import functools +from inspect import signature import logging import re import warnings @@ -140,18 +141,37 @@ def magnitude(self): @property def unit_array(self): - """Return the data values of this DataArray as a `pint.Quantity`.""" + """Return the data values of this DataArray as a `pint.Quantity`. + + Notes + ----- + If not already existing as a `pint.Quantity` or Dask array, the data of this DataArray + will be loaded into memory by this operation. + """ if isinstance(self._data_array.data, units.Quantity): return self._data_array.data else: - return units.Quantity(self._data_array.values, self.units) + return units.Quantity(self._data_array.data, self.units) def convert_units(self, units): - """Return new DataArray with values converted to different units.""" + """Return new DataArray with values converted to different units. + + Notes + ----- + Any cached/lazy-loaded data (except that in a Dask array) will be loaded into memory + by this operation. Do not utilize on moderate- to large-sized remote datasets before + subsetting! + """ return self.quantify().copy(data=self.unit_array.to(units)) def convert_coordinate_units(self, coord, units): - """Return new DataArray with coordinate converted to different units.""" + """Return new DataArray with coordinate converted to different units. + + Notes + ----- + Any cached/lazy-loaded coordinate data (except that in a Dask array) will be loaded + into memory by this operation. + """ new_coord_var = self._data_array[coord].copy( data=self._data_array[coord].metpy.unit_array.m_as(units) ) @@ -159,7 +179,14 @@ def convert_coordinate_units(self, coord, units): return self._data_array.assign_coords(coords={coord: new_coord_var}) def quantify(self): - """Return a DataArray with the data converted to a `pint.Quantity`.""" + """Return a DataArray with the data converted to a `pint.Quantity`. + + Notes + ----- + Any cached/lazy-loaded data (except that in a Dask array) will be loaded into memory + by this operation. Do not utilize on moderate- to large-sized remote datasets before + subsetting! + """ if ( not isinstance(self._data_array.data, units.Quantity) and np.issubdtype(self._data_array.data.dtype, np.number) @@ -203,14 +230,17 @@ def cartopy_globe(self): def _fixup_coordinate_map(self, coord_map): """Ensure sure we have coordinate variables in map, not coordinate names.""" + new_coord_map = {} for axis in coord_map: if coord_map[axis] is not None and not isinstance(coord_map[axis], xr.DataArray): - coord_map[axis] = self._data_array[coord_map[axis]] + new_coord_map[axis] = self._data_array[coord_map[axis]] + else: + new_coord_map[axis] = coord_map[axis] - return coord_map + return new_coord_map def assign_coordinates(self, coordinates): - """Assign the given coordinates to the given MetPy axis types. + """Return new DataArray with given coordinates assigned to the given MetPy axis types. Parameters ---------- @@ -221,18 +251,32 @@ def assign_coordinates(self, coordinates): which will trigger reparsing of all coordinates on next access. """ + coord_updates = {} if coordinates: # Assign the _metpy_axis attributes according to supplied mapping coordinates = self._fixup_coordinate_map(coordinates) for axis in coordinates: if coordinates[axis] is not None: - _assign_axis(coordinates[axis].attrs, axis) + coord_updates[coordinates[axis].name] = ( + coordinates[axis].assign_attrs( + _assign_axis(coordinates[axis].attrs.copy(), axis) + ) + ) else: # Clear _metpy_axis attribute on all coordinates - for coord_var in self._data_array.coords.values(): - coord_var.attrs.pop('_metpy_axis', None) + for coord_name, coord_var in self._data_array.coords.items(): + coord_updates[coord_name] = coord_var.copy(deep=False) - return self._data_array # allow method chaining + # Some coordinates remained linked in old form under other coordinates. We + # need to remove from these. + sub_coords = coord_updates[coord_name].coords + for sub_coord in sub_coords: + coord_updates[coord_name].coords[sub_coord].attrs.pop('_metpy_axis', None) + + # Now we can remove the _metpy_axis attr from the coordinate itself + coord_updates[coord_name].attrs.pop('_metpy_axis', None) + + return self._data_array.assign_coords(coord_updates) def _generate_coordinate_map(self): """Generate a coordinate map via CF conventions and other methods.""" @@ -291,6 +335,11 @@ def _metpy_axis_search(self, metpy_axis): return coord_var # Opportunistically parse all coordinates, and assign if not already assigned + # Note: since this is generally called by way of the coordinate properties, to cache + # the coordinate parsing results in coord_map on the coordinates means modifying the + # DataArray in-place (an exception to the usual behavior of MetPy's accessor). This is + # considered safe because it only effects the "_metpy_axis" attribute on the + # coordinates, and nothing else. coord_map = self._generate_coordinate_map() for axis, coord_var in coord_map.items(): if (coord_var is not None @@ -625,7 +674,7 @@ def parse_cf(self, varname=None, coordinates=None): # Assign coordinates if the coordinates argument is given if coordinates is not None: - var.metpy.assign_coordinates(coordinates) + var = var.metpy.assign_coordinates(coordinates) # Attempt to build the crs coordinate crs = None @@ -658,7 +707,7 @@ def _has_coord(coord_type): var = self._rebuild_coords(var, crs) if crs is not None: var = var.assign_coords(coords={'crs': crs}) - return var.metpy.quantify() + return var def _rebuild_coords(self, var, crs): """Clean up the units on the coordinate variables.""" @@ -814,7 +863,7 @@ def assign_y_x(self, force=False, tolerance=None): return self._dataset.assign_coords(**{y.name: y, x.name: x}) def update_attribute(self, attribute, mapping): - """Update attribute of all Dataset variables. + """Return new Dataset with specified attribute updated on all Dataset variables. Parameters ---------- @@ -829,24 +878,41 @@ def update_attribute(self, attribute, mapping): Returns ------- `xarray.Dataset` - Dataset with attribute updated (modified in place, and returned to allow method - chaining) + New Dataset with attribute updated """ # Make mapping uniform - if callable(mapping): - mapping_func = mapping - else: - def mapping_func(varname, **kwargs): - return mapping.get(varname, None) + if not callable(mapping): + old_mapping = mapping + + def mapping(varname, **kwargs): + return old_mapping.get(varname, None) - # Apply across all variables - for varname in list(self._dataset.data_vars) + list(self._dataset.coords): - value = mapping_func(varname, **self._dataset[varname].attrs) - if value is not None: - self._dataset[varname].attrs[attribute] = value + # Define mapping function for Dataset.map + def mapping_func(da): + new_value = mapping(da.name, **da.attrs) + if new_value is None: + return da + else: + return da.assign_attrs(**{attribute: new_value}) + + # Apply across all variables and coordinates + return ( + self._dataset + .map(mapping_func, keep_attrs=True) + .assign_coords({ + coord_name: mapping_func(coord_var) + for coord_name, coord_var in self._dataset.coords.items() + }) + ) + + def quantify(self): + """Return new dataset with all numeric variables quantified and cached data loaded.""" + return self._dataset.map(lambda da: da.metpy.quantify(), keep_attrs=True) - return self._dataset + def dequantify(self): + """Return new dataset with variables cast to magnitude and units on attribute.""" + return self._dataset.map(lambda da: da.metpy.dequantify(), keep_attrs=True) def _assign_axis(attributes, axis): @@ -986,18 +1052,130 @@ def _build_y_x(da, tolerance): 'correpsond to your CRS coordinate.') -def preprocess_xarray(func): - """Decorate a function to convert all DataArray arguments to pint.Quantities. +def preprocess_and_wrap(broadcast=None, wrap_like=None, match_unit=False, to_magnitude=False): + """Return decorator to wrap array calculations for type flexibility. + + Assuming you have a calculation that works internally with `pint.Quantity` or + `numpy.ndarray`, this will wrap the function to be able to handle `xarray.DataArray` and + `pint.Quantity` as well (assuming appropriate match to one of the input arguments). - This uses the metpy xarray accessors to do the actual conversion. + Parameters + ---------- + broadcast : iterable of str or None + Iterable of string labels for arguments to broadcast against each other using xarray, + assuming they are supplied as `xarray.DataArray`. No automatic broadcasting will occur + 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, + and wrap iteratively according to the str or array-like contained within. If None, + will not wrap output. + match_unit : bool + If true, force the unit of the final output to be that of wrapping object (as + determined by wrap_like), no matter the original calculation output. Defaults to + False. + to_magnitude : bool + If true, downcast xarray and Pint arguments to their magnitude. If false, downcast + xarray arguments to Quantity, and do not change other array-like arguments. """ - @functools.wraps(func) - def wrapper(*args, **kwargs): - args = tuple(a.metpy.unit_array if isinstance(a, xr.DataArray) else a for a in args) - kwargs = {name: (v.metpy.unit_array if isinstance(v, xr.DataArray) else v) - for name, v in kwargs.items()} - return func(*args, **kwargs) - return wrapper + def decorator(func): + @functools.wraps(func) + def wrapper(*args, **kwargs): + bound_args = signature(func).bind(*args, **kwargs) + + # Auto-broadcast select xarray arguments, and update bound_args + if broadcast is not None: + arg_names_to_broadcast = tuple( + arg_name for arg_name in broadcast + if arg_name in bound_args.arguments + and isinstance(bound_args.arguments[arg_name], xr.DataArray) + ) + broadcasted_args = xr.broadcast( + *(bound_args.arguments[arg_name] for arg_name in arg_names_to_broadcast) + ) + for i, arg_name in enumerate(arg_names_to_broadcast): + bound_args.arguments[arg_name] = broadcasted_args[i] + + # Obtain proper match if referencing an input + match = list(wrap_like) if isinstance(wrap_like, tuple) else wrap_like + if isinstance(wrap_like, str): + match = bound_args.arguments[wrap_like] + elif isinstance(wrap_like, tuple): + for i in range(len(wrap_like)): + if isinstance(wrap_like[i], str): + match[i] = bound_args.arguments[wrap_like[i]] + + # Cast all DataArrays to Pint Quantities + for arg_name in bound_args.arguments: + if isinstance(bound_args.arguments[arg_name], xr.DataArray): + bound_args.arguments[arg_name] = ( + bound_args.arguments[arg_name].metpy.unit_array + ) + + # Optionally cast all Quantities to their magnitudes + if to_magnitude: + for arg_name in bound_args.arguments: + if isinstance(bound_args.arguments[arg_name], units.Quantity): + bound_args.arguments[arg_name] = bound_args.arguments[arg_name].m + + # Evaluate inner calculation + result = func(*bound_args.args, **bound_args.kwargs) + + # Wrap output based on match and match_unit + if match is None: + return result + else: + if match_unit: + wrapping = _wrap_output_like_matching_units + else: + wrapping = _wrap_output_like_not_matching_units + + if isinstance(match, list): + return tuple(wrapping(*args) for args in zip(result, match)) + else: + return wrapping(result, match) + return wrapper + return decorator + + +def _wrap_output_like_matching_units(result, match): + """Convert result to be like match with matching units for output wrapper.""" + output_xarray = isinstance(match, xr.DataArray) + match_units = str(match.metpy.units if output_xarray else getattr(match, 'units', '')) + + if isinstance(result, xr.DataArray): + result = result.metpy.convert_units(match_units) + return result if output_xarray else result.metpy.unit_array + else: + result = ( + result.to(match_units) if isinstance(result, units.Quantity) + else units.Quantity(result, match_units) + ) + return ( + xr.DataArray(result, coords=match.coords, dims=match.dims) if output_xarray + else result + ) + + +def _wrap_output_like_not_matching_units(result, match): + """Convert result to be like match without matching units for output wrapper.""" + output_xarray = isinstance(match, xr.DataArray) + if isinstance(result, xr.DataArray): + return result if output_xarray else result.metpy.unit_array + else: + # Determine if need to upcast to Quantity + if ( + ( + isinstance(match, units.Quantity) + or (output_xarray and isinstance(match.data, units.Quantity)) + ) + and not isinstance(result, units.Quantity) + ): + result = units.Quantity(result) + return ( + xr.DataArray(result, coords=match.coords, dims=match.dims) if output_xarray + else result + ) def check_matching_coordinates(func): diff --git a/tests/calc/test_calc_tools.py b/tests/calc/test_calc_tools.py index 319078434e8..056f5ceef1d 100644 --- a/tests/calc/test_calc_tools.py +++ b/tests/calc/test_calc_tools.py @@ -19,10 +19,9 @@ reduce_point_density, resample_nn_1d, second_derivative) from metpy.calc.tools import (_delete_masked_points, _get_bound_pressure_height, _greater_or_close, _less_or_close, _next_non_masked_element, - _remove_nans, BASE_DEGREE_MULTIPLIER, DIR_STRS, UND, - wrap_output_like) + _remove_nans, BASE_DEGREE_MULTIPLIER, DIR_STRS, UND) from metpy.testing import (assert_almost_equal, assert_array_almost_equal, assert_array_equal) -from metpy.units import DimensionalityError, units +from metpy.units import units FULL_CIRCLE_DEGREES = np.arange(0, 360, BASE_DEGREE_MULTIPLIER.m) * units.degree @@ -1289,213 +1288,3 @@ def test_remove_nans(): y_expected = np.array([0, 1, 3, 4]) assert_array_almost_equal(x_expected, x_test, 0) assert_almost_equal(y_expected, y_test, 0) - - -@pytest.mark.parametrize('test, other, match_unit, expected', [ - (np.arange(4), np.arange(4), False, np.arange(4) * units('dimensionless')), - (np.arange(4), np.arange(4), True, np.arange(4) * units('dimensionless')), - (np.arange(4), [0] * units.m, False, np.arange(4) * units('dimensionless')), - (np.arange(4), [0] * units.m, True, np.arange(4) * units.m), - ( - np.arange(4), - xr.DataArray( - np.zeros(4) * units.meter, - dims=('x',), - coords={'x': np.linspace(0, 1, 4)}, - attrs={'description': 'Just some zeros'} - ), - False, - xr.DataArray( - np.arange(4) * units.dimensionless, - dims=('x',), - coords={'x': np.linspace(0, 1, 4)} - ) - ), - ( - np.arange(4), - xr.DataArray( - np.zeros(4) * units.meter, - dims=('x',), - coords={'x': np.linspace(0, 1, 4)}, - attrs={'description': 'Just some zeros'} - ), - True, - xr.DataArray( - np.arange(4) * units.meter, - dims=('x',), - coords={'x': np.linspace(0, 1, 4)} - ) - ), - ([2, 4, 8] * units.kg, [0] * units.m, False, [2, 4, 8] * units.kg), - ([2, 4, 8] * units.kg, [0] * units.g, True, [2000, 4000, 8000] * units.g), - ( - [2, 4, 8] * units.kg, - xr.DataArray( - np.zeros(3) * units.meter, - dims=('x',), - coords={'x': np.linspace(0, 1, 3)} - ), - False, - xr.DataArray( - [2, 4, 8] * units.kilogram, - dims=('x',), - coords={'x': np.linspace(0, 1, 3)} - ) - ), - ( - [2, 4, 8] * units.kg, - xr.DataArray( - np.zeros(3) * units.gram, - dims=('x',), - coords={'x': np.linspace(0, 1, 3)} - ), - True, - xr.DataArray( - [2000, 4000, 8000] * units.gram, - dims=('x',), - coords={'x': np.linspace(0, 1, 3)} - ) - ), - ( - xr.DataArray( - np.linspace(0, 1, 5) * units.meter, - attrs={'description': 'A range of values'} - ), - np.arange(4, dtype=np.float64), - False, - units.Quantity(np.linspace(0, 1, 5), 'meter') - ), - ( - xr.DataArray( - np.linspace(0, 1, 5) * units.meter, - attrs={'description': 'A range of values'} - ), - [0] * units.kg, - False, - np.linspace(0, 1, 5) * units.m - ), - ( - xr.DataArray( - np.linspace(0, 1, 5) * units.meter, - attrs={'description': 'A range of values'} - ), - [0] * units.cm, - True, - np.linspace(0, 100, 5) * units.cm - ), - ( - xr.DataArray( - np.linspace(0, 1, 5) * units.meter, - attrs={'description': 'A range of values'} - ), - xr.DataArray( - np.zeros(3) * units.kilogram, - dims=('x',), - coords={'x': np.linspace(0, 1, 3)}, - attrs={'description': 'Alternative data'} - ), - False, - xr.DataArray( - np.linspace(0, 1, 5) * units.meter, - attrs={'description': 'A range of values'} - ) - ), - ( - xr.DataArray( - np.linspace(0, 1, 5) * units.meter, - attrs={'description': 'A range of values'} - ), - xr.DataArray( - np.zeros(3) * units.centimeter, - dims=('x',), - coords={'x': np.linspace(0, 1, 3)}, - attrs={'description': 'Alternative data'} - ), - True, - xr.DataArray( - np.linspace(0, 100, 5) * units.centimeter, - attrs={'description': 'A range of values'} - ) - ), -]) -def test_wrap_output_like_with_other_kwarg(test, other, match_unit, expected): - """Test the wrap output like decorator when using the output kwarg.""" - @wrap_output_like(other=other, match_unit=match_unit) - def almost_identity(arg): - return arg - - result = almost_identity(test) - - if hasattr(expected, 'units'): - assert expected.units == result.units - if isinstance(expected, xr.DataArray): - xr.testing.assert_identical(result, expected) - else: - assert_array_equal(result, expected) - - -@pytest.mark.parametrize('test, other', [ - ([2, 4, 8] * units.kg, [0] * units.m), - ( - [2, 4, 8] * units.kg, - xr.DataArray( - np.zeros(3) * units.meter, - dims=('x',), - coords={'x': np.linspace(0, 1, 3)} - ) - ), - ( - xr.DataArray( - np.linspace(0, 1, 5) * units.meter - ), - [0] * units.kg - ), - ( - xr.DataArray( - np.linspace(0, 1, 5) * units.meter - ), - xr.DataArray( - np.zeros(3) * units.kg, - dims=('x',), - coords={'x': np.linspace(0, 1, 3)} - ) - ), - ( - xr.DataArray( - np.linspace(0, 1, 5) * units.meter, - attrs={'description': 'A range of values'} - ), - np.arange(4, dtype=np.float64) - ) -]) -def test_wrap_output_like_with_other_kwarg_raising_dimensionality_error(test, other): - """Test the wrap output like decorator when when a dimensionality error is raised.""" - @wrap_output_like(other=other, match_unit=True) - def almost_identity(arg): - return arg - - with pytest.raises(DimensionalityError): - almost_identity(test) - - -def test_wrap_output_like_with_argument_kwarg(): - """Test the wrap output like decorator with signature recognition.""" - @wrap_output_like(argument='a') - def double(a): - return units.Quantity(2) * a.metpy.unit_array - - test = xr.DataArray([1, 3, 5, 7] * units.m) - expected = xr.DataArray([2, 6, 10, 14] * units.m) - - xr.testing.assert_identical(double(test), expected) - - -def test_wrap_output_like_without_control_kwarg(): - """Test that the wrap output like decorator fails when not provided a control param.""" - @wrap_output_like() - def func(arg): - """Do nothing.""" - - with pytest.raises(ValueError) as exc: - func(0) - assert 'Must specify keyword' in str(exc) diff --git a/tests/calc/test_kinematics.py b/tests/calc/test_kinematics.py index 2b46e5957d4..6db5f0cf05c 100644 --- a/tests/calc/test_kinematics.py +++ b/tests/calc/test_kinematics.py @@ -1436,8 +1436,8 @@ def test_geostrophic_wind_4d(data_4d): assert_array_almost_equal(v_g, v_g_truth, 6) -def test_intertial_advective_wind_4d(data_4d): - """Test intertial_advective_wind on a 4D (time, pressure, y, x) grid.""" +def test_inertial_advective_wind_4d(data_4d): + """Test inertial_advective_wind on a 4D (time, pressure, y, x) grid.""" f = coriolis_parameter(data_4d.latitude) u_g, v_g = geostrophic_wind(data_4d.height, f, data_4d.dx, data_4d.dy) u_i, v_i = inertial_advective_wind(u_g, v_g, u_g, v_g, data_4d.dx, data_4d.dy, diff --git a/tests/calc/test_thermo.py b/tests/calc/test_thermo.py index 69ec2511256..3b33cbd1cca 100644 --- a/tests/calc/test_thermo.py +++ b/tests/calc/test_thermo.py @@ -476,6 +476,24 @@ def test_equivalent_potential_temperature(): assert_almost_equal(ept, 311.18586467284007 * units.kelvin, 3) +def test_equivalent_potential_temperature_masked(): + """Test equivalent potential temperature calculation with masked arrays.""" + p = 1000 * units.mbar + t = units.Quantity(np.ma.array([293., 294., 295.]), units.kelvin) + td = units.Quantity( + np.ma.array([280., 281., 282.], mask=[False, True, False]), + units.kelvin + ) + ept = equivalent_potential_temperature(p, t, td) + expected = units.Quantity( + np.ma.array([311.18586, 313.51781, 315.93971], mask=[False, True, False]), + units.kelvin + ) + assert isinstance(ept, units.Quantity) + assert isinstance(ept.m, np.ma.MaskedArray) + assert_array_almost_equal(ept, expected, 3) + + def test_saturation_equivalent_potential_temperature(): """Test saturation equivalent potential temperature calculation.""" p = 700 * units.mbar @@ -486,6 +504,20 @@ def test_saturation_equivalent_potential_temperature(): assert_almost_equal(s_ept, 299.096584 * units.kelvin, 3) +def test_saturation_equivalent_potential_temperature_masked(): + """Test saturation equivalent potential temperature calculation with masked arrays.""" + p = 1000 * units.mbar + t = units.Quantity(np.ma.array([293., 294., 295.]), units.kelvin) + s_ept = saturation_equivalent_potential_temperature(p, t) + expected = units.Quantity( + np.ma.array([335.02750, 338.95813, 343.08740]), + units.kelvin + ) + assert isinstance(s_ept, units.Quantity) + assert isinstance(s_ept.m, np.ma.MaskedArray) + assert_array_almost_equal(s_ept, expected, 3) + + def test_virtual_temperature(): """Test virtual temperature calculation.""" t = 288. * units.kelvin diff --git a/tests/plots/baseline/test_arrow_projection.png b/tests/plots/baseline/test_arrow_projection.png index e99617251cd..e338acc350c 100644 Binary files a/tests/plots/baseline/test_arrow_projection.png and b/tests/plots/baseline/test_arrow_projection.png differ diff --git a/tests/plots/baseline/test_barb_projection.png b/tests/plots/baseline/test_barb_projection.png index 0b18e53cd81..765921a3266 100644 Binary files a/tests/plots/baseline/test_barb_projection.png and b/tests/plots/baseline/test_barb_projection.png differ diff --git a/tests/plots/baseline/test_colorfill.png b/tests/plots/baseline/test_colorfill.png index 21b19663067..f4b382a8c41 100644 Binary files a/tests/plots/baseline/test_colorfill.png and b/tests/plots/baseline/test_colorfill.png differ diff --git a/tests/plots/baseline/test_colorfill_horiz_colorbar.png b/tests/plots/baseline/test_colorfill_horiz_colorbar.png index 8d3630a6a11..7958a2cc28f 100644 Binary files a/tests/plots/baseline/test_colorfill_horiz_colorbar.png and b/tests/plots/baseline/test_colorfill_horiz_colorbar.png differ diff --git a/tests/plots/baseline/test_colorfill_no_colorbar.png b/tests/plots/baseline/test_colorfill_no_colorbar.png index b70baf7de5d..26da7345d10 100644 Binary files a/tests/plots/baseline/test_colorfill_no_colorbar.png and b/tests/plots/baseline/test_colorfill_no_colorbar.png differ diff --git a/tests/plots/baseline/test_declarative_contour_options.png b/tests/plots/baseline/test_declarative_contour_options.png index db82cc68264..485899f86f7 100644 Binary files a/tests/plots/baseline/test_declarative_contour_options.png and b/tests/plots/baseline/test_declarative_contour_options.png differ diff --git a/tests/plots/test_declarative.py b/tests/plots/test_declarative.py index 04639f1d50b..b2d6a702161 100644 --- a/tests/plots/test_declarative.py +++ b/tests/plots/test_declarative.py @@ -542,7 +542,7 @@ def test_declarative_sfc_obs_changes(): return pc.figure -@pytest.mark.mpl_image_compare(remove_text=True, tolerance=0) +@pytest.mark.mpl_image_compare(remove_text=True, tolerance=0.00586) def test_declarative_colored_barbs(): """Test making a surface plot with a colored barb (gh-1274).""" data = pd.read_csv(get_test_data('SFC_obs.csv', as_file_obj=False), @@ -575,7 +575,8 @@ def test_declarative_colored_barbs(): @pytest.mark.mpl_image_compare(remove_text=True, - tolerance={'3.1': 9.771, '2.1': 9.771}.get(MPL_VERSION, 0.)) + tolerance={'3.1': 9.771, + '2.1': 9.771}.get(MPL_VERSION, 0.00651)) def test_declarative_sfc_obs_full(): """Test making a full surface observation plot.""" data = pd.read_csv(get_test_data('SFC_obs.csv', as_file_obj=False), diff --git a/tests/plots/test_station_plot.py b/tests/plots/test_station_plot.py index 72c6883ca55..712b328bb37 100644 --- a/tests/plots/test_station_plot.py +++ b/tests/plots/test_station_plot.py @@ -287,7 +287,7 @@ def test_barb_projection(wind_plot): # Plot and check barbs (they should align with grid lines) fig = plt.figure() ax = fig.add_subplot(1, 1, 1, projection=ccrs.LambertConformal()) - ax.gridlines(xlocs=[-135, -120, -105, -90, -75, -60, -45]) + ax.gridlines(xlocs=[-120, -105, -90, -75, -60], ylocs=np.arange(24, 55, 6)) sp = StationPlot(ax, x, y, transform=ccrs.PlateCarree()) sp.plot_barb(u, v) @@ -302,7 +302,7 @@ def test_arrow_projection(wind_plot): # Plot and check barbs (they should align with grid lines) fig = plt.figure() ax = fig.add_subplot(1, 1, 1, projection=ccrs.LambertConformal()) - ax.gridlines(xlocs=[-135, -120, -105, -90, -75, -60, -45]) + ax.gridlines(xlocs=[-120, -105, -90, -75, -60], ylocs=np.arange(24, 55, 6)) sp = StationPlot(ax, x, y, transform=ccrs.PlateCarree()) sp.plot_arrow(u, v) sp.plot_arrow(u, v) # plot_arrow used twice to hit removal if statement diff --git a/tests/test_xarray.py b/tests/test_xarray.py index 60cde770898..d6f6bd54672 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -14,8 +14,8 @@ from metpy.plots.mapping import CFProjection from metpy.testing import (assert_almost_equal, assert_array_almost_equal, assert_array_equal, get_test_data) -from metpy.units import units -from metpy.xarray import check_axis, check_matching_coordinates, preprocess_xarray +from metpy.units import DimensionalityError, units +from metpy.xarray import check_axis, check_matching_coordinates, preprocess_and_wrap # Seed RandomState for deterministic tests @@ -151,13 +151,37 @@ def test_quantify(test_ds_generic): np.testing.assert_array_almost_equal(result.data, units.Quantity(original)) -def test_dequantify(test_var): +def test_dequantify(): """Test dequantify method for converting data away from Quantity.""" - original = test_var.data - result = test_var.metpy.dequantify() + original = xr.DataArray(units.Quantity([280, 290, 300], 'K')) + result = original.metpy.dequantify() assert isinstance(result.data, np.ndarray) assert result.attrs['units'] == 'kelvin' - np.testing.assert_array_almost_equal(result.data, original.magnitude) + np.testing.assert_array_almost_equal(result.data, original.data.magnitude) + + +def test_dataset_quantify(test_ds_generic): + """Test quantify method for converting data to Quantity on Datasets.""" + result = test_ds_generic.metpy.quantify() + assert isinstance(result['test'].data, units.Quantity) + assert result['test'].data.units == units.dimensionless + assert 'units' not in result['test'].attrs + np.testing.assert_array_almost_equal( + result['test'].data, + units.Quantity(test_ds_generic['test'].data) + ) + + +def test_dataset_dequantify(): + """Test dequantify method for converting data away from Quantity on Datasets.""" + original = xr.Dataset({ + 'test': ('x', units.Quantity([280, 290, 300], 'K')), + 'x': np.arange(3) + }) + result = original.metpy.dequantify() + assert isinstance(result['test'].data, np.ndarray) + assert result['test'].attrs['units'] == 'kelvin' + np.testing.assert_array_almost_equal(result['test'].data, original['test'].data.magnitude) def test_radian_projection_coords(): @@ -208,12 +232,12 @@ def test_missing_grid_mapping_var(caplog): assert 'Could not find' in caplog.text -def test_preprocess_xarray(): - """Test xarray preprocessing decorator.""" +def test_preprocess_and_wrap_only_preprocessing(): + """Test xarray preprocessing and wrapping decorator for only preprocessing.""" data = xr.DataArray(np.ones(3), attrs={'units': 'km'}) data2 = xr.DataArray(np.ones(3), attrs={'units': 'm'}) - @preprocess_xarray + @preprocess_and_wrap() def func(a, b): return a.to('m') + b @@ -278,7 +302,7 @@ def test_assign_coordinates_not_overwrite(test_ds_generic): """Test that assign_coordinates does not overwrite past axis attributes.""" data = test_ds_generic.copy() data['c'].attrs['axis'] = 'X' - data['test'].metpy.assign_coordinates({'y': data['c']}) + data['test'] = data['test'].metpy.assign_coordinates({'y': data['c']}) assert data['c'].identical(data['test'].metpy.y) assert data['c'].attrs['axis'] == 'X' @@ -598,9 +622,12 @@ def test_data_array_sel_dict_with_units(test_var): def test_data_array_sel_kwargs_with_units(test_var): """Test .sel on the metpy accessor with kwargs and axis type.""" truth = test_var.loc[:, 500.][..., 122] - selection = test_var.metpy.sel(vertical=5e4 * units.Pa, x=-16.569 * units.km, - tolerance=1., method='nearest') - selection.metpy.assign_coordinates(None) # truth was not parsed for coordinates + selection = ( + test_var.metpy + .sel(vertical=5e4 * units.Pa, x=-16.569 * units.km, tolerance=1., method='nearest') + .metpy + .assign_coordinates(None) + ) assert truth.identical(selection) @@ -959,13 +986,19 @@ def test_update_attribute_dictionary(test_ds_generic): 'test': 'Filler data', 'c': 'The third coordinate' } - test_ds_generic.metpy.update_attribute('description', descriptions) - assert 'description' not in test_ds_generic['a'].attrs - assert 'description' not in test_ds_generic['b'].attrs - assert test_ds_generic['c'].attrs['description'] == 'The third coordinate' - assert 'description' not in test_ds_generic['d'].attrs - assert 'description' not in test_ds_generic['e'].attrs - assert test_ds_generic['test'].attrs['description'] == 'Filler data' + result = test_ds_generic.metpy.update_attribute('description', descriptions) + + # Test attribute updates + assert 'description' not in result['a'].attrs + assert 'description' not in result['b'].attrs + assert result['c'].attrs['description'] == 'The third coordinate' + assert 'description' not in result['d'].attrs + assert 'description' not in result['e'].attrs + assert result['test'].attrs['description'] == 'Filler data' + + # Test for no side effects + assert 'description' not in test_ds_generic['c'].attrs + assert 'description' not in test_ds_generic['test'].attrs def test_update_attribute_callable(test_ds_generic): @@ -973,10 +1006,245 @@ def test_update_attribute_callable(test_ds_generic): def even_ascii(varname, **kwargs): if ord(varname[0]) % 2 == 0: return 'yes' + result = test_ds_generic.metpy.update_attribute('even', even_ascii) + + # Test attribute updates + assert 'even' not in result['a'].attrs + assert result['b'].attrs['even'] == 'yes' + assert 'even' not in result['c'].attrs + assert result['d'].attrs['even'] == 'yes' + assert 'even' not in result['e'].attrs + assert result['test'].attrs['even'] == 'yes' + + # Test for no side effects + assert 'even' not in test_ds_generic['b'].attrs + assert 'even' not in test_ds_generic['d'].attrs + assert 'even' not in test_ds_generic['test'].attrs test_ds_generic.metpy.update_attribute('even', even_ascii) - assert 'even' not in test_ds_generic['a'].attrs - assert test_ds_generic['b'].attrs['even'] == 'yes' - assert 'even' not in test_ds_generic['c'].attrs - assert test_ds_generic['d'].attrs['even'] == 'yes' - assert 'even' not in test_ds_generic['e'].attrs - assert test_ds_generic['test'].attrs['even'] == 'yes' + + +@pytest.mark.parametrize('test, other, match_unit, expected', [ + (np.arange(4), np.arange(4), False, np.arange(4)), + (np.arange(4), np.arange(4), True, np.arange(4) * units('dimensionless')), + (np.arange(4), [0] * units.m, False, np.arange(4) * units('dimensionless')), + (np.arange(4), [0] * units.m, True, np.arange(4) * units.m), + ( + np.arange(4), + xr.DataArray( + np.zeros(4) * units.meter, + dims=('x',), + coords={'x': np.linspace(0, 1, 4)}, + attrs={'description': 'Just some zeros'} + ), + False, + xr.DataArray( + np.arange(4) * units.dimensionless, + dims=('x',), + coords={'x': np.linspace(0, 1, 4)} + ) + ), + ( + np.arange(4), + xr.DataArray( + np.zeros(4) * units.meter, + dims=('x',), + coords={'x': np.linspace(0, 1, 4)}, + attrs={'description': 'Just some zeros'} + ), + True, + xr.DataArray( + np.arange(4) * units.meter, + dims=('x',), + coords={'x': np.linspace(0, 1, 4)} + ) + ), + ([2, 4, 8] * units.kg, [0] * units.m, False, [2, 4, 8] * units.kg), + ([2, 4, 8] * units.kg, [0] * units.g, True, [2000, 4000, 8000] * units.g), + ( + [2, 4, 8] * units.kg, + xr.DataArray( + np.zeros(3) * units.meter, + dims=('x',), + coords={'x': np.linspace(0, 1, 3)} + ), + False, + xr.DataArray( + [2, 4, 8] * units.kilogram, + dims=('x',), + coords={'x': np.linspace(0, 1, 3)} + ) + ), + ( + [2, 4, 8] * units.kg, + xr.DataArray( + np.zeros(3) * units.gram, + dims=('x',), + coords={'x': np.linspace(0, 1, 3)} + ), + True, + xr.DataArray( + [2000, 4000, 8000] * units.gram, + dims=('x',), + coords={'x': np.linspace(0, 1, 3)} + ) + ), + ( + xr.DataArray( + np.linspace(0, 1, 5) * units.meter, + attrs={'description': 'A range of values'} + ), + np.arange(4, dtype=np.float64), + False, + units.Quantity(np.linspace(0, 1, 5), 'meter') + ), + ( + xr.DataArray( + np.linspace(0, 1, 5) * units.meter, + attrs={'description': 'A range of values'} + ), + [0] * units.kg, + False, + np.linspace(0, 1, 5) * units.m + ), + ( + xr.DataArray( + np.linspace(0, 1, 5) * units.meter, + attrs={'description': 'A range of values'} + ), + [0] * units.cm, + True, + np.linspace(0, 100, 5) * units.cm + ), + ( + xr.DataArray( + np.linspace(0, 1, 5) * units.meter, + attrs={'description': 'A range of values'} + ), + xr.DataArray( + np.zeros(5) * units.kilogram, + dims=('x',), + coords={'x': np.linspace(0, 1, 5)}, + attrs={'description': 'Alternative data'} + ), + False, + xr.DataArray( + np.linspace(0, 1, 5) * units.meter, + dims=('x',), + coords={'x': np.linspace(0, 1, 5)} + ) + ), + ( + xr.DataArray( + np.linspace(0, 1, 5) * units.meter, + attrs={'description': 'A range of values'} + ), + xr.DataArray( + np.zeros(5) * units.centimeter, + dims=('x',), + coords={'x': np.linspace(0, 1, 5)}, + attrs={'description': 'Alternative data'} + ), + True, + xr.DataArray( + np.linspace(0, 100, 5) * units.centimeter, + dims=('x',), + coords={'x': np.linspace(0, 1, 5)} + ) + ), +]) +def test_wrap_with_wrap_like_kwarg(test, other, match_unit, expected): + """Test the preprocess and wrap decorator when using wrap_like.""" + @preprocess_and_wrap(wrap_like=other, match_unit=match_unit) + def almost_identity(arg): + return arg + + result = almost_identity(test) + + if hasattr(expected, 'units'): + assert expected.units == result.units + if isinstance(expected, xr.DataArray): + xr.testing.assert_identical(result, expected) + else: + assert_array_equal(result, expected) + + +@pytest.mark.parametrize('test, other', [ + ([2, 4, 8] * units.kg, [0] * units.m), + ( + [2, 4, 8] * units.kg, + xr.DataArray( + np.zeros(3) * units.meter, + dims=('x',), + coords={'x': np.linspace(0, 1, 3)} + ) + ), + ( + xr.DataArray( + np.linspace(0, 1, 5) * units.meter + ), + [0] * units.kg + ), + ( + xr.DataArray( + np.linspace(0, 1, 5) * units.meter + ), + xr.DataArray( + np.zeros(5) * units.kg, + dims=('x',), + coords={'x': np.linspace(0, 1, 5)} + ) + ), + ( + xr.DataArray( + np.linspace(0, 1, 5) * units.meter, + attrs={'description': 'A range of values'} + ), + np.arange(4, dtype=np.float64) + ) +]) +def test_wrap_with_wrap_like_kwarg_raising_dimensionality_error(test, other): + """Test the preprocess and wrap decorator when a dimensionality error is raised.""" + @preprocess_and_wrap(wrap_like=other, match_unit=True) + def almost_identity(arg): + return arg + + with pytest.raises(DimensionalityError): + almost_identity(test) + + +def test_wrap_with_argument_kwarg(): + """Test the preprocess and wrap decorator with signature recognition.""" + @preprocess_and_wrap(wrap_like='a') + def double(a): + return units.Quantity(2) * a + + test = xr.DataArray([1, 3, 5, 7] * units.m) + expected = xr.DataArray([2, 6, 10, 14] * units.m) + + xr.testing.assert_identical(double(test), expected) + + +def test_preprocess_and_wrap_with_broadcasting(): + """Test preprocessing and wrapping decorator with arguments to broadcast specified.""" + # Not quantified + data = xr.DataArray(np.arange(9).reshape((3, 3)), dims=('y', 'x'), attrs={'units': 'N'}) + # Quantified + data2 = xr.DataArray([1, 0, 0] * units.m, dims=('y')) + + @preprocess_and_wrap(broadcast=('a', 'b')) + def func(a, b): + return a * b + + assert_array_equal(func(data, data2), [[0, 1, 2], [0, 0, 0], [0, 0, 0]] * units('N m')) + + +def test_preprocess_and_wrap_with_to_magnitude(): + """Test preprocessing and wrapping with casting to magnitude.""" + data = xr.DataArray([1, 0, 1] * units.m) + data2 = [0, 1, 1] * units.cm + + @preprocess_and_wrap(wrap_like='a', to_magnitude=True) + def func(a, b): + return a * b + + np.testing.assert_array_equal(func(data, data2), np.array([0, 0, 1]))