From cee20d7181e3ddf28c4d86c151a26e451c0cf730 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Tue, 14 Jan 2020 02:01:18 -0600 Subject: [PATCH 1/4] Update preprocess_xarray to handle broadcasting --- src/metpy/calc/basic.py | 32 ++++---- src/metpy/calc/indices.py | 14 ++-- src/metpy/calc/kinematics.py | 32 ++++---- src/metpy/calc/thermo.py | 100 ++++++++++++------------- src/metpy/calc/tools.py | 22 +++--- src/metpy/calc/turbulence.py | 8 +- src/metpy/interpolate/one_dimension.py | 7 +- src/metpy/xarray.py | 40 +++++++--- tests/test_xarray.py | 14 +++- 9 files changed, 151 insertions(+), 118 deletions(-) diff --git a/src/metpy/calc/basic.py b/src/metpy/calc/basic.py index 19005b1eea2..979f074b48c 100644 --- a/src/metpy/calc/basic.py +++ b/src/metpy/calc/basic.py @@ -29,7 +29,7 @@ @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]', '[speed]') def wind_speed(u, v): r"""Compute the wind speed from u and v-components. @@ -56,7 +56,7 @@ def wind_speed(u, v): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]', '[speed]') def wind_direction(u, v, convention='from'): r"""Compute the wind direction from u and v-components. @@ -108,7 +108,7 @@ def wind_direction(u, v, convention='from'): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]') def wind_components(speed, wind_direction): r"""Calculate the U, V wind vector components from the speed and direction. @@ -147,7 +147,7 @@ def wind_components(speed, wind_direction): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +209,7 @@ def windchill(temperature, speed, face_level_winds=False, mask_undefined=True): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +313,7 @@ def heat_index(temperature, relative_humidity, mask_undefined=True): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units(temperature='[temperature]', speed='[speed]') def apparent_temperature(temperature, relative_humidity, speed, face_level_winds=False, mask_undefined=True): @@ -392,7 +392,7 @@ def apparent_temperature(temperature, relative_humidity, speed, face_level_winds @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]') def pressure_to_height_std(pressure): r"""Convert pressure data to height using the U.S. standard atmosphere [NOAA1976]_. @@ -420,7 +420,7 @@ def pressure_to_height_std(pressure): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[length]') def height_to_geopotential(height): r"""Compute geopotential for a given height above sea level. @@ -477,7 +477,7 @@ def height_to_geopotential(height): @exporter.export -@preprocess_xarray +@preprocess_xarray() def geopotential_to_height(geopotential): r"""Compute height above sea level from a given geopotential. @@ -537,7 +537,7 @@ def geopotential_to_height(geopotential): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[length]') def height_to_pressure_std(height): r"""Convert height data to pressures using the U.S. standard atmosphere [NOAA1976]_. @@ -564,7 +564,7 @@ def height_to_pressure_std(height): @exporter.export -@preprocess_xarray +@preprocess_xarray() def coriolis_parameter(latitude): r"""Calculate the coriolis parameter at each point. @@ -586,7 +586,7 @@ def coriolis_parameter(latitude): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +615,7 @@ def add_height_to_pressure(pressure, height): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[length]', '[pressure]') def add_pressure_to_height(height, pressure): r"""Calculate the height at a certain pressure above another height. @@ -644,7 +644,7 @@ def add_pressure_to_height(height, pressure): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[dimensionless]', '[pressure]', '[pressure]') def sigma_to_pressure(sigma, pressure_sfc, pressure_top): r"""Calculate pressure from sigma values. @@ -1007,7 +1007,7 @@ def smooth_n_point(scalar_grid, n=5, passes=1): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[length]') def altimeter_to_station_pressure(altimeter_value, height): r"""Convert the altimeter measurement to station pressure. @@ -1089,7 +1089,7 @@ def altimeter_to_station_pressure(altimeter_value, height): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 b569a26cc0c..38cbcb9b681 100644 --- a/src/metpy/calc/indices.py +++ b/src/metpy/calc/indices.py @@ -15,7 +15,7 @@ @exporter.export -@preprocess_xarray +@preprocess_xarray() @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. @@ -74,7 +74,7 @@ def precipitable_water(pressure, dewpoint, *, bottom=None, top=None): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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. @@ -123,7 +123,7 @@ def mean_pressure_weighted(pressure, *args, height=None, bottom=None, depth=None @exporter.export -@preprocess_xarray +@preprocess_xarray() @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. @@ -183,7 +183,7 @@ def bunkers_storm_motion(pressure, u, v, height): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[speed]', '[speed]') def bulk_shear(pressure, u, v, height=None, bottom=None, depth=None): r"""Calculate bulk shear through a layer. @@ -226,7 +226,7 @@ def bulk_shear(pressure, u, v, height=None, bottom=None, depth=None): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[energy] / [mass]', '[speed] * [speed]', '[speed]') def supercell_composite(mucape, effective_storm_helicity, effective_shear): r"""Calculate the supercell composite parameter. @@ -268,7 +268,7 @@ def supercell_composite(mucape, effective_storm_helicity, effective_shear): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +322,7 @@ def significant_tornado(sbcape, surface_based_lcl_height, storm_helicity_1km, sh @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[speed]', '[speed]', '[length]', '[speed]', '[speed]') def critical_angle(pressure, u, v, height, u_storm, v_storm): r"""Calculate the critical angle. diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index 73caeff4ec9..a70f34e77f8 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -20,7 +20,7 @@ def _stack(arrs): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]', '[speed]', '[length]', '[length]') def vorticity(u, v, dx, dy): r"""Calculate the vertical vorticity of the horizontal wind. @@ -59,7 +59,7 @@ def vorticity(u, v, dx, dy): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units(dx='[length]', dy='[length]') def divergence(u, v, dx, dy): r"""Calculate the horizontal divergence of a vector. @@ -98,7 +98,7 @@ def divergence(u, v, dx, dy): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]', '[speed]', '[length]', '[length]') def shearing_deformation(u, v, dx, dy): r"""Calculate the shearing deformation of the horizontal wind. @@ -137,7 +137,7 @@ def shearing_deformation(u, v, dx, dy): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]', '[speed]', '[length]', '[length]') def stretching_deformation(u, v, dx, dy): r"""Calculate the stretching deformation of the horizontal wind. @@ -176,7 +176,7 @@ def stretching_deformation(u, v, dx, dy): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]', '[speed]', '[length]', '[length]') def total_deformation(u, v, dx, dy): r"""Calculate the horizontal total deformation of the horizontal wind. @@ -215,7 +215,7 @@ def total_deformation(u, v, dx, dy): @exporter.export -@preprocess_xarray +@preprocess_xarray() def advection(scalar, wind, deltas): r"""Calculate the advection of a scalar field by the wind. @@ -266,7 +266,7 @@ def advection(scalar, wind, deltas): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[temperature]', '[speed]', '[speed]', '[length]', '[length]') def frontogenesis(potential_temperature, u, v, dx, dy): r"""Calculate the 2D kinematic frontogenesis of a temperature field. @@ -334,7 +334,7 @@ def frontogenesis(potential_temperature, u, v, dx, dy): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units(f='[frequency]', dx='[length]', dy='[length]') def geostrophic_wind(height, f, dx, dy): r"""Calculate the geostrophic wind given from the height or geopotential. @@ -376,7 +376,7 @@ def geostrophic_wind(height, f, dx, dy): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units(f='[frequency]', u='[speed]', v='[speed]', dx='[length]', dy='[length]') def ageostrophic_wind(height, u, v, f, dx, dy): r"""Calculate the ageostrophic wind given from the height or geopotential. @@ -419,7 +419,7 @@ def ageostrophic_wind(height, u, v, f, dx, dy): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[length]', '[temperature]') def montgomery_streamfunction(height, temperature): r"""Compute the Montgomery Streamfunction on isentropic surfaces. @@ -461,7 +461,7 @@ def montgomery_streamfunction(height, temperature): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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, @@ -528,7 +528,7 @@ def storm_relative_helicity(height, u, v, depth, *, bottom=0 * units.m, @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]', '[speed]', '[length]', '[length]') def absolute_vorticity(u, v, dx, dy, latitude): """Calculate the absolute vorticity of the horizontal wind. @@ -565,7 +565,7 @@ def absolute_vorticity(u, v, dx, dy, latitude): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[temperature]', '[pressure]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') def potential_vorticity_baroclinic(potential_temperature, pressure, u, v, dx, dy, latitude): @@ -643,7 +643,7 @@ def potential_vorticity_baroclinic(potential_temperature, pressure, u, v, dx, dy @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[length]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') def potential_vorticity_barotropic(height, u, v, dx, dy, latitude): r"""Calculate the barotropic (Rossby) potential vorticity. @@ -685,7 +685,7 @@ def potential_vorticity_barotropic(height, u, v, dx, dy, latitude): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]', '[speed]', '[speed]', '[speed]', '[length]', '[length]', '[dimensionless]') def inertial_advective_wind(u, v, u_geostrophic, v_geostrophic, dx, dy, latitude): @@ -752,7 +752,7 @@ def inertial_advective_wind(u, v, u_geostrophic, v_geostrophic, dx, dy, latitude @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[speed]', '[speed]', '[temperature]', '[pressure]', '[length]', '[length]') def q_vector(u, v, temperature, pressure, dx, dy, static_stability=1): r"""Calculate Q-vector at a given pressure level using the u, v winds and temperature. diff --git a/src/metpy/calc/thermo.py b/src/metpy/calc/thermo.py index 95cf81f04b5..42eeca0d06c 100644 --- a/src/metpy/calc/thermo.py +++ b/src/metpy/calc/thermo.py @@ -23,7 +23,7 @@ @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[temperature]', '[temperature]') def relative_humidity_from_dewpoint(temperature, dewpoint): r"""Calculate the relative humidity. @@ -54,7 +54,7 @@ def relative_humidity_from_dewpoint(temperature, dewpoint): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[pressure]') def exner_function(pressure, reference_pressure=mpconsts.P0): r"""Calculate the Exner function. @@ -89,7 +89,7 @@ def exner_function(pressure, reference_pressure=mpconsts.P0): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]') def potential_temperature(pressure, temperature): r"""Calculate the potential temperature. @@ -131,7 +131,7 @@ def potential_temperature(pressure, temperature): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]') def temperature_from_potential_temperature(pressure, potential_temperature): r"""Calculate the temperature from a given potential temperature. @@ -176,7 +176,7 @@ def temperature_from_potential_temperature(pressure, potential_temperature): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[pressure]') def dry_lapse(pressure, temperature, reference_pressure=None): r"""Calculate the temperature at a level assuming only dry processes. @@ -212,7 +212,7 @@ def dry_lapse(pressure, temperature, reference_pressure=None): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[pressure]') def moist_lapse(pressure, temperature, reference_pressure=None): r"""Calculate the temperature at a level assuming liquid saturation processes. @@ -300,7 +300,7 @@ def dt(t, p): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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. @@ -364,7 +364,7 @@ def _lcl_iter(p, p0, w, t): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]', '[temperature]') def lfc(pressure, temperature, dewpoint, parcel_temperature_profile=None, dewpoint_start=None, which='top'): @@ -539,7 +539,7 @@ def _most_cape_option(intersect_type, p_list, t_list, pressure, temperature, dew @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]', '[temperature]') def el(pressure, temperature, dewpoint, parcel_temperature_profile=None, which='top'): r"""Calculate the equilibrium level. @@ -604,7 +604,7 @@ def el(pressure, temperature, dewpoint, parcel_temperature_profile=None, which=' @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]') def parcel_profile(pressure, temperature, dewpoint): r"""Calculate the profile a parcel takes through the atmosphere. @@ -638,7 +638,7 @@ def parcel_profile(pressure, temperature, dewpoint): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]') def parcel_profile_with_lcl(pressure, temperature, dewpoint): r"""Calculate the profile a parcel takes through the atmosphere. @@ -728,7 +728,7 @@ def _insert_lcl_level(pressure, temperature, lcl_pressure): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[dimensionless]') def vapor_pressure(pressure, mixing_ratio): r"""Calculate water vapor (partial) pressure. @@ -765,7 +765,7 @@ def vapor_pressure(pressure, mixing_ratio): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[temperature]') def saturation_vapor_pressure(temperature): r"""Calculate the saturation water vapor (partial) pressure. @@ -801,7 +801,7 @@ def saturation_vapor_pressure(temperature): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +829,7 @@ def dewpoint_from_relative_humidity(temperature, relative_humidity): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]') def dewpoint(vapor_pressure): r"""Calculate the ambient dewpoint given the vapor pressure. @@ -862,7 +862,7 @@ def dewpoint(vapor_pressure): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +904,7 @@ def mixing_ratio(partial_press, total_press, molecular_weight_ratio=mpconsts.eps @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]') def saturation_mixing_ratio(total_press, temperature): r"""Calculate the saturation mixing ratio of water vapor. @@ -929,7 +929,7 @@ def saturation_mixing_ratio(total_press, temperature): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]') def equivalent_potential_temperature(pressure, temperature, dewpoint): r"""Calculate equivalent potential temperature. @@ -986,7 +986,7 @@ def equivalent_potential_temperature(pressure, temperature, dewpoint): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]') def saturation_equivalent_potential_temperature(pressure, temperature): r"""Calculate saturation equivalent potential temperature. @@ -1053,7 +1053,7 @@ def saturation_equivalent_potential_temperature(pressure, temperature): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[temperature]', '[dimensionless]', '[dimensionless]') def virtual_temperature(temperature, mixing_ratio, molecular_weight_ratio=mpconsts.epsilon): r"""Calculate virtual temperature. @@ -1087,7 +1087,7 @@ def virtual_temperature(temperature, mixing_ratio, molecular_weight_ratio=mpcons @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[dimensionless]', '[dimensionless]') def virtual_potential_temperature(pressure, temperature, mixing_ratio, molecular_weight_ratio=mpconsts.epsilon): @@ -1124,7 +1124,7 @@ def virtual_potential_temperature(pressure, temperature, mixing_ratio, @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[dimensionless]', '[dimensionless]') def density(pressure, temperature, mixing_ratio, molecular_weight_ratio=mpconsts.epsilon): r"""Calculate density. @@ -1160,7 +1160,7 @@ def density(pressure, temperature, mixing_ratio, molecular_weight_ratio=mpconsts @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]') def relative_humidity_wet_psychrometric(pressure, dry_bulb_temperature, web_bulb_temperature, **kwargs): @@ -1202,7 +1202,7 @@ def relative_humidity_wet_psychrometric(pressure, dry_bulb_temperature, web_bulb @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +1252,7 @@ def psychrometric_vapor_pressure_wet(pressure, dry_bulb_temperature, wet_bulb_te @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +1292,7 @@ def mixing_ratio_from_relative_humidity(pressure, temperature, relative_humidity @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +1330,7 @@ def relative_humidity_from_mixing_ratio(pressure, temperature, mixing_ratio): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[dimensionless]') def mixing_ratio_from_specific_humidity(specific_humidity): r"""Calculate the mixing ratio from specific humidity. @@ -1367,7 +1367,7 @@ def mixing_ratio_from_specific_humidity(specific_humidity): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[dimensionless]') def specific_humidity_from_mixing_ratio(mixing_ratio): r"""Calculate the specific humidity from the mixing ratio. @@ -1404,7 +1404,7 @@ def specific_humidity_from_mixing_ratio(mixing_ratio): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +1443,7 @@ def relative_humidity_from_specific_humidity(pressure, temperature, specific_hum @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]', '[temperature]') def cape_cin(pressure, temperature, dewpoint, parcel_profile, which_lfc='bottom', which_el='top'): @@ -1593,7 +1593,7 @@ def _find_append_zero_crossings(x, y): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]') def most_unstable_parcel(pressure, temperature, dewpoint, height=None, bottom=None, depth=300 * units.hPa): @@ -1640,7 +1640,7 @@ def most_unstable_parcel(pressure, temperature, dewpoint, height=None, @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[temperature]', '[pressure]', '[temperature]') def isentropic_interpolation(levels, pressure, temperature, *args, axis=0, temperature_out=False, max_iters=50, eps=1e-6, @@ -1787,7 +1787,7 @@ def _isen_iter(iter_log_p, isentlevs_nd, ka, a, b, pok): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]') def surface_based_cape_cin(pressure, temperature, dewpoint): r"""Calculate surface-based CAPE and CIN. @@ -1826,7 +1826,7 @@ def surface_based_cape_cin(pressure, temperature, dewpoint): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]') def most_unstable_cape_cin(pressure, temperature, dewpoint, **kwargs): r"""Calculate most unstable CAPE/CIN. @@ -1869,7 +1869,7 @@ def most_unstable_cape_cin(pressure, temperature, dewpoint, **kwargs): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]') def mixed_layer_cape_cin(pressure, temperature, dewpoint, **kwargs): r"""Calculate mixed-layer CAPE and CIN. @@ -1920,7 +1920,7 @@ def mixed_layer_cape_cin(pressure, temperature, dewpoint, **kwargs): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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): @@ -1988,7 +1988,7 @@ def mixed_parcel(pressure, temperature, dewpoint, parcel_start_pressure=None, @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]') def mixed_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa, interpolate=True): @@ -2034,7 +2034,7 @@ def mixed_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[length]', '[temperature]') def dry_static_energy(height, temperature): r"""Calculate the dry static energy of parcels. @@ -2066,7 +2066,7 @@ def dry_static_energy(height, temperature): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[length]', '[temperature]', '[dimensionless]') def moist_static_energy(height, temperature, specific_humidity): r"""Calculate the moist static energy of parcels. @@ -2102,7 +2102,7 @@ def moist_static_energy(height, temperature, specific_humidity): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]') def thickness_hydrostatic(pressure, temperature, mixing_ratio=None, molecular_weight_ratio=mpconsts.epsilon, bottom=None, depth=None): @@ -2172,7 +2172,7 @@ def thickness_hydrostatic(pressure, temperature, mixing_ratio=None, @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]') def thickness_hydrostatic_from_relative_humidity(pressure, temperature, relative_humidity, bottom=None, depth=None): @@ -2225,7 +2225,7 @@ def thickness_hydrostatic_from_relative_humidity(pressure, temperature, relative @exporter.export -@preprocess_xarray +@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. @@ -2265,7 +2265,7 @@ def brunt_vaisala_frequency_squared(height, potential_temperature, axis=0): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[length]', '[temperature]') def brunt_vaisala_frequency(height, potential_temperature, axis=0): r"""Calculate the Brunt-Vaisala frequency. @@ -2306,7 +2306,7 @@ def brunt_vaisala_frequency(height, potential_temperature, axis=0): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[length]', '[temperature]') def brunt_vaisala_period(height, potential_temperature, axis=0): r"""Calculate the Brunt-Vaisala period. @@ -2345,7 +2345,7 @@ def brunt_vaisala_period(height, potential_temperature, axis=0): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]', '[temperature]') def wet_bulb_temperature(pressure, temperature, dewpoint): """Calculate the wet-bulb temperature using Normand's rule. @@ -2398,7 +2398,7 @@ def wet_bulb_temperature(pressure, temperature, dewpoint): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]') def static_stability(pressure, temperature, axis=0): r"""Calculate the static stability within a vertical profile. @@ -2430,7 +2430,7 @@ def static_stability(pressure, temperature, axis=0): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +2460,7 @@ def dewpoint_from_specific_humidity(pressure, temperature, specific_humidity): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +2503,7 @@ def vertical_velocity_pressure(w, pressure, temperature, mixing_ratio=0): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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 +2549,7 @@ def vertical_velocity(omega, pressure, temperature, mixing_ratio=0): @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]', '[temperature]') def specific_humidity_from_dewpoint(pressure, dewpoint): r"""Calculate the specific humidity from the dewpoint temperature and pressure. diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 8282a7b3313..61b7b050553 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -39,7 +39,7 @@ @exporter.export -@preprocess_xarray +@preprocess_xarray() def resample_nn_1d(a, centers): """Return one-dimensional nearest-neighbor indexes based on user-specified centers. @@ -65,7 +65,7 @@ def resample_nn_1d(a, centers): @exporter.export -@preprocess_xarray +@preprocess_xarray() def nearest_intersection_idx(a, b): """Determine the index of the point just before two lines with common x values. @@ -93,7 +93,7 @@ def nearest_intersection_idx(a, b): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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. @@ -233,7 +233,7 @@ def _delete_masked_points(*arrs): @exporter.export -@preprocess_xarray +@preprocess_xarray() def reduce_point_density(points, radius, priority=None): r"""Return a mask to reduce the density of points in irregularly-spaced data. @@ -417,7 +417,7 @@ def _get_bound_pressure_height(pressure, bound, height=None, interpolate=True): @exporter.export -@preprocess_xarray +@preprocess_xarray() @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. @@ -507,7 +507,7 @@ def get_layer_heights(height, depth, *args, bottom=None, interpolate=True, with_ @exporter.export -@preprocess_xarray +@preprocess_xarray() @check_units('[pressure]') def get_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa, interpolate=True): @@ -609,7 +609,7 @@ def get_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa, @exporter.export -@preprocess_xarray +@preprocess_xarray() def find_bounding_indices(arr, values, axis, from_below=True): """Find the indices surrounding the values within arr along axis. @@ -750,7 +750,7 @@ def take(indexer): @exporter.export -@preprocess_xarray +@preprocess_xarray() 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. @@ -905,7 +905,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_xarray()(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)) @@ -1296,7 +1296,7 @@ def _process_deriv_args(f, kwargs): @exporter.export -@preprocess_xarray +@preprocess_xarray() def parse_angle(input_dir): """Calculate the meteorological angle from directional text. @@ -1352,7 +1352,7 @@ def _abbrieviate_direction(ext_dir_str): @exporter.export -@preprocess_xarray +@preprocess_xarray() def angle_to_direction(input_angle, full=False, level=3): """Convert the meteorological angle to directional text. diff --git a/src/metpy/calc/turbulence.py b/src/metpy/calc/turbulence.py index fa55bb46f51..eaf39068a3b 100644 --- a/src/metpy/calc/turbulence.py +++ b/src/metpy/calc/turbulence.py @@ -13,7 +13,7 @@ @exporter.export -@preprocess_xarray +@preprocess_xarray() 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_xarray() 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_xarray() 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_xarray() 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..a825d08e1f8 100644 --- a/src/metpy/interpolate/one_dimension.py +++ b/src/metpy/interpolate/one_dimension.py @@ -14,7 +14,8 @@ @exporter.export -@preprocess_xarray + +@preprocess_xarray() def interpolate_nans_1d(x, y, kind='linear'): """Interpolate NaN values in y. @@ -49,7 +50,7 @@ def interpolate_nans_1d(x, y, kind='linear'): @exporter.export -@preprocess_xarray +@preprocess_xarray() 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 +175,7 @@ def interpolate_1d(x, xp, *args, axis=0, fill_value=np.nan, return_list_always=F @exporter.export -@preprocess_xarray +@preprocess_xarray() 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/xarray.py b/src/metpy/xarray.py index b69452b756f..67d218d133b 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 @@ -874,7 +875,6 @@ def check_axis(var, *axes): # If no match has been made, return False (rather than None) return False - def _assign_crs(xarray_object, cf_attributes, cf_kwargs): from .plots.mapping import CFProjection @@ -947,18 +947,38 @@ 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_xarray(*broadcast_argument_labels): + """Create decorator to convert input DataArray arguments to pint.Quantities. + + While all input DataArrays are converted to Quantities, any arguments named by ``*args`` + will be broadcasted together prior to conversion. This uses the metpy xarray accessors to do the actual conversion. """ - @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 selected arguements, and update bound arguments for signature + arg_labels_to_broadcast = [label for label in broadcast_argument_labels + if label in bound_args.arguments + and isinstance(bound_args.arguments[label], + xr.DataArray)] + broadcasted_args = xr.broadcast(*tuple(bound_args.arguments[label] + for label in arg_labels_to_broadcast)) + for i, label in enumerate(arg_labels_to_broadcast): + bound_args.arguments[label] = broadcasted_args[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 + + return func(*bound_args.args, **bound_args.kwargs) + return wrapper + return decorator def check_matching_coordinates(func): diff --git a/tests/test_xarray.py b/tests/test_xarray.py index a76a5d7024f..6583762f725 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -172,13 +172,25 @@ def test_preprocess_xarray(): data = xr.DataArray(np.ones(3), attrs={'units': 'km'}) data2 = xr.DataArray(np.ones(3), attrs={'units': 'm'}) - @preprocess_xarray + @preprocess_xarray() def func(a, b): return a.to('m') + b assert_array_equal(func(data, b=data2), np.array([1001, 1001, 1001]) * units.m) +def test_preprocess_xarray_with_broadcasting(): + """test xarray preprocessing decorator with arguments to broadcast specified.""" + data = xr.DataArray(np.arange(9).reshape((3, 3)), dims=('y', 'x'), attrs={'units': 'N'}) + data2 = xr.DataArray([1, 0, 0], dims=('y'), attrs={'units': 'm'}) + + @preprocess_xarray('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_strftime(): """Test our monkey-patched xarray strftime.""" data = xr.DataArray(np.datetime64('2000-01-01 01:00:00')) From 828ed7d5d90272f49e3cb1ce185ad64555548d51 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Tue, 14 Jan 2020 03:00:31 -0600 Subject: [PATCH 2/4] Modify deriviative calculations to have explcit signatures --- src/metpy/calc/tools.py | 53 +++++++++++++++++++++-------------------- 1 file changed, 27 insertions(+), 26 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 61b7b050553..0c95ea9547e 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -941,7 +941,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 +984,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 +1031,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 +1074,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 +1117,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 +1164,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 +1215,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 +1237,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 +1265,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 = atleast_1d(kwargs['delta']) + delta = atleast_1d(delta) if delta.size == 1: diff_size = list(f.shape) diff_size[axis] -= 1 @@ -1286,8 +1287,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 = diff(x, axis=axis) else: raise ValueError('Must specify either "x" or "delta" for value positions.') From 799e74c2fa687bba2c39a1fdf3c56d7729c255c5 Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Tue, 14 Jan 2020 03:54:39 -0600 Subject: [PATCH 3/4] Update decorators throught library for enhanced xarray compat (where possible) --- src/metpy/calc/basic.py | 14 ++ src/metpy/calc/kinematics.py | 114 ++++++++-- src/metpy/calc/thermo.py | 94 ++++++--- src/metpy/calc/tools.py | 278 +++++++++++++------------ src/metpy/interpolate/one_dimension.py | 7 +- tests/calc/test_thermo.py | 3 +- tests/test_xarray.py | 2 +- 7 files changed, 334 insertions(+), 178 deletions(-) diff --git a/src/metpy/calc/basic.py b/src/metpy/calc/basic.py index 979f074b48c..c3caa28badc 100644 --- a/src/metpy/calc/basic.py +++ b/src/metpy/calc/basic.py @@ -29,6 +29,7 @@ @exporter.export +@wrap_output_like(argument='u') @preprocess_xarray() @check_units('[speed]', '[speed]') def wind_speed(u, v): @@ -56,6 +57,7 @@ def wind_speed(u, v): @exporter.export +@wrap_output_like(argument='u') @preprocess_xarray() @check_units('[speed]', '[speed]') def wind_direction(u, v, convention='from'): @@ -147,6 +149,7 @@ def wind_components(speed, wind_direction): @exporter.export +@wrap_output_like(argument='temperature') @preprocess_xarray() @check_units(temperature='[temperature]', speed='[speed]') def windchill(temperature, speed, face_level_winds=False, mask_undefined=True): @@ -209,6 +212,7 @@ def windchill(temperature, speed, face_level_winds=False, mask_undefined=True): @exporter.export +@wrap_output_like(argument='temperature') @preprocess_xarray() @check_units('[temperature]') def heat_index(temperature, relative_humidity, mask_undefined=True): @@ -313,6 +317,7 @@ def heat_index(temperature, relative_humidity, mask_undefined=True): @exporter.export +@wrap_output_like(argument='temperature') @preprocess_xarray() @check_units(temperature='[temperature]', speed='[speed]') def apparent_temperature(temperature, relative_humidity, speed, face_level_winds=False, @@ -392,6 +397,7 @@ def apparent_temperature(temperature, relative_humidity, speed, face_level_winds @exporter.export +@wrap_output_like(argument='pressure') @preprocess_xarray() @check_units('[pressure]') def pressure_to_height_std(pressure): @@ -420,6 +426,7 @@ def pressure_to_height_std(pressure): @exporter.export +@wrap_output_like(argument='height') @preprocess_xarray() @check_units('[length]') def height_to_geopotential(height): @@ -477,6 +484,7 @@ def height_to_geopotential(height): @exporter.export +@wrap_output_like(argument='geopotential') @preprocess_xarray() def geopotential_to_height(geopotential): r"""Compute height above sea level from a given geopotential. @@ -537,6 +545,7 @@ def geopotential_to_height(geopotential): @exporter.export +@wrap_output_like(argument='height') @preprocess_xarray() @check_units('[length]') def height_to_pressure_std(height): @@ -564,6 +573,7 @@ def height_to_pressure_std(height): @exporter.export +@wrap_output_like(argument='latitude') @preprocess_xarray() def coriolis_parameter(latitude): r"""Calculate the coriolis parameter at each point. @@ -586,6 +596,7 @@ def coriolis_parameter(latitude): @exporter.export +@wrap_output_like(argument='pressure') @preprocess_xarray() @check_units('[pressure]', '[length]') def add_height_to_pressure(pressure, height): @@ -615,6 +626,7 @@ def add_height_to_pressure(pressure, height): @exporter.export +@wrap_output_like(argument='height') @preprocess_xarray() @check_units('[length]', '[pressure]') def add_pressure_to_height(height, pressure): @@ -1007,6 +1019,7 @@ def smooth_n_point(scalar_grid, n=5, passes=1): @exporter.export +@wrap_output_like(argument='altimeter_value') @preprocess_xarray() @check_units('[pressure]', '[length]') def altimeter_to_station_pressure(altimeter_value, height): @@ -1089,6 +1102,7 @@ def altimeter_to_station_pressure(altimeter_value, height): @exporter.export +@wrap_output_like(argument='altimeter_value') @preprocess_xarray() @check_units('[pressure]', '[length]', '[temperature]') def altimeter_to_sea_level_pressure(altimeter_value, height, temperature): diff --git a/src/metpy/calc/kinematics.py b/src/metpy/calc/kinematics.py index a70f34e77f8..b2bf143de3d 100644 --- a/src/metpy/calc/kinematics.py +++ b/src/metpy/calc/kinematics.py @@ -2,10 +2,15 @@ # 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 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, + wrap_output_like) from .. import constants as mpconsts from ..cbook import iterable from ..package_tools import Exporter @@ -19,10 +24,54 @@ 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 dx/dy + if (grid_prototype is not None + and bound_args.arguments['dx'] is None + and bound_args.arguments['dy'] is None): + (bound_args.arguments['dx'], + bound_args.arguments['dy']) = grid_deltas_from_dataarray(grid_prototype, + kind='actual') + + # Fill in latitude + if (grid_prototype is not None + and 'latitude' in bound_args.arguments + and bound_args.arguments['latitude'] is None): + bound_args.arguments['latitude'] = grid_prototype.metpy.latitude.metpy.unit_array + + # Fill in Coriolis parameter + if (grid_prototype is not None + and 'f' in bound_args.arguments + and bound_args.arguments['f'] is None): + bound_args.arguments['f'] = coriolis_parameter( + grid_prototype.metpy.latitude.metpy.unit_array) + + return func(*bound_args.args, **bound_args.kwargs) + return wrapper + + @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @check_units('[speed]', '[speed]', '[length]', '[length]') -def vorticity(u, v, dx, dy): +def vorticity(u, v, dx=None, dy=None): r"""Calculate the vertical vorticity of the horizontal wind. Parameters @@ -59,9 +108,11 @@ def vorticity(u, v, dx, dy): @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @check_units(dx='[length]', dy='[length]') -def divergence(u, v, dx, dy): +def divergence(u, v, dx=None, dy=None): r"""Calculate the horizontal divergence of a vector. Parameters @@ -98,9 +149,11 @@ def divergence(u, v, dx, dy): @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @check_units('[speed]', '[speed]', '[length]', '[length]') -def shearing_deformation(u, v, dx, dy): +def shearing_deformation(u, v, dx=None, dy=None): r"""Calculate the shearing deformation of the horizontal wind. Parameters @@ -137,9 +190,11 @@ def shearing_deformation(u, v, dx, dy): @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @check_units('[speed]', '[speed]', '[length]', '[length]') -def stretching_deformation(u, v, dx, dy): +def stretching_deformation(u, v, dx=None, dy=None): r"""Calculate the stretching deformation of the horizontal wind. Parameters @@ -176,9 +231,11 @@ def stretching_deformation(u, v, dx, dy): @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @check_units('[speed]', '[speed]', '[length]', '[length]') -def total_deformation(u, v, dx, dy): +def total_deformation(u, v, dx=None, dy=None): r"""Calculate the horizontal total deformation of the horizontal wind. Parameters @@ -243,6 +300,10 @@ def advection(scalar, wind, deltas): N-dimensional array An N-dimensional array containing the advection at all grid points. + Notes + ----- + This function does not current return xarray DataArrays. + """ # This allows passing in a list of wind components or an array. wind = _stack(wind) @@ -266,9 +327,11 @@ def advection(scalar, wind, deltas): @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @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): r"""Calculate the 2D kinematic frontogenesis of a temperature field. The implementation is a form of the Petterssen Frontogenesis and uses the formula @@ -334,9 +397,11 @@ def frontogenesis(potential_temperature, u, v, dx, dy): @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='height') @preprocess_xarray() @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): r"""Calculate the geostrophic wind given from the height or geopotential. Parameters @@ -376,9 +441,11 @@ def geostrophic_wind(height, f, dx, dy): @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @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): r"""Calculate the ageostrophic wind given from the height or geopotential. Parameters @@ -419,6 +486,7 @@ def ageostrophic_wind(height, u, v, f, dx, dy): @exporter.export +@wrap_output_like(argument='height') @preprocess_xarray() @check_units('[length]', '[temperature]') def montgomery_streamfunction(height, temperature): @@ -504,6 +572,10 @@ def storm_relative_helicity(height, u, v, depth, *, bottom=0 * units.m, `pint.Quantity` total storm-relative helicity + Notes + ----- + This function does not currently return xarray DataArrays + """ _, u, v = get_layer_heights(height, depth, u, v, with_agl=True, bottom=bottom) @@ -528,9 +600,11 @@ def storm_relative_helicity(height, u, v, depth, *, bottom=0 * units.m, @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @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): """Calculate the absolute vorticity of the horizontal wind. Parameters @@ -565,10 +639,13 @@ def absolute_vorticity(u, v, dx, dy, latitude): @exporter.export -@preprocess_xarray() +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') +@preprocess_xarray('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): r"""Calculate the baroclinic potential vorticity. .. math:: PV = -g \left(\frac{\partial u}{\partial p}\frac{\partial \theta}{\partial y} @@ -643,9 +720,11 @@ def potential_vorticity_baroclinic(potential_temperature, pressure, u, v, dx, dy @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @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): r"""Calculate the barotropic (Rossby) potential vorticity. .. math:: PV = \frac{f + \zeta}{H} @@ -685,10 +764,13 @@ def potential_vorticity_barotropic(height, u, v, dx, dy, latitude): @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @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): r"""Calculate the inertial advective wind. .. math:: \frac{\hat k}{f} \times (\vec V \cdot \nabla)\hat V_g @@ -752,9 +834,11 @@ def inertial_advective_wind(u, v, u_geostrophic, v_geostrophic, dx, dy, latitude @exporter.export +@add_grid_arguments_from_xarray +@wrap_output_like(argument='u') @preprocess_xarray() @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): r"""Calculate Q-vector at a given pressure level using the u, v winds and temperature. .. math:: \vec{Q} = (Q_1, Q_2) diff --git a/src/metpy/calc/thermo.py b/src/metpy/calc/thermo.py index 42eeca0d06c..765869814eb 100644 --- a/src/metpy/calc/thermo.py +++ b/src/metpy/calc/thermo.py @@ -9,7 +9,7 @@ import scipy.optimize as so from .tools import (_greater_or_close, _less_or_close, _remove_nans, find_bounding_indices, - find_intersections, first_derivative, get_layer) + find_intersections, first_derivative, get_layer, wrap_output_like) from .. import constants as mpconsts from ..cbook import broadcast_indices from ..interpolate.one_dimension import interpolate_1d @@ -23,7 +23,8 @@ @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('temperature', 'dewpoint') @check_units('[temperature]', '[temperature]') def relative_humidity_from_dewpoint(temperature, dewpoint): r"""Calculate the relative humidity. @@ -54,6 +55,7 @@ def relative_humidity_from_dewpoint(temperature, dewpoint): @exporter.export +@wrap_output_like(argument='pressure') @preprocess_xarray() @check_units('[pressure]', '[pressure]') def exner_function(pressure, reference_pressure=mpconsts.P0): @@ -89,7 +91,8 @@ def exner_function(pressure, reference_pressure=mpconsts.P0): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('pressure', 'temperature') @check_units('[pressure]', '[temperature]') def potential_temperature(pressure, temperature): r"""Calculate the potential temperature. @@ -131,7 +134,8 @@ def potential_temperature(pressure, temperature): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='potential_temperature') +@preprocess_xarray('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,7 +180,8 @@ def temperature_from_potential_temperature(pressure, potential_temperature): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('pressure', 'temperature') @check_units('[pressure]', '[temperature]', '[pressure]') def dry_lapse(pressure, temperature, reference_pressure=None): r"""Calculate the temperature at a level assuming only dry processes. @@ -728,7 +733,8 @@ def _insert_lcl_level(pressure, temperature, lcl_pressure): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='mixing_ratio') +@preprocess_xarray('pressure', 'mixing_ratio') @check_units('[pressure]', '[dimensionless]') def vapor_pressure(pressure, mixing_ratio): r"""Calculate water vapor (partial) pressure. @@ -765,6 +771,7 @@ def vapor_pressure(pressure, mixing_ratio): @exporter.export +@wrap_output_like(argument='temperature') @preprocess_xarray() @check_units('[temperature]') def saturation_vapor_pressure(temperature): @@ -801,7 +808,8 @@ def saturation_vapor_pressure(temperature): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('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,6 +837,7 @@ def dewpoint_from_relative_humidity(temperature, relative_humidity): @exporter.export +@wrap_output_like(argument='vapor_pressure') @preprocess_xarray() @check_units('[pressure]') def dewpoint(vapor_pressure): @@ -862,7 +871,8 @@ def dewpoint(vapor_pressure): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='partial_press') +@preprocess_xarray('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 +914,8 @@ def mixing_ratio(partial_press, total_press, molecular_weight_ratio=mpconsts.eps @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('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 +940,8 @@ def saturation_mixing_ratio(total_press, temperature): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('pressure', 'temperature', 'dewpoint') @check_units('[pressure]', '[temperature]', '[temperature]') def equivalent_potential_temperature(pressure, temperature, dewpoint): r"""Calculate equivalent potential temperature. @@ -986,7 +998,8 @@ def equivalent_potential_temperature(pressure, temperature, dewpoint): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('pressure', 'temperature') @check_units('[pressure]', '[temperature]') def saturation_equivalent_potential_temperature(pressure, temperature): r"""Calculate saturation equivalent potential temperature. @@ -1053,7 +1066,8 @@ def saturation_equivalent_potential_temperature(pressure, temperature): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('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 +1101,8 @@ def virtual_temperature(temperature, mixing_ratio, molecular_weight_ratio=mpcons @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('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 +1139,8 @@ def virtual_potential_temperature(pressure, temperature, mixing_ratio, @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('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,7 +1176,8 @@ def density(pressure, temperature, mixing_ratio, molecular_weight_ratio=mpconsts @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='dry_bulb_temperature') +@preprocess_xarray('pressure', 'dry_bulb_temperature', 'wet_bulb_temperature') @check_units('[pressure]', '[temperature]', '[temperature]') def relative_humidity_wet_psychrometric(pressure, dry_bulb_temperature, web_bulb_temperature, **kwargs): @@ -1202,7 +1219,8 @@ def relative_humidity_wet_psychrometric(pressure, dry_bulb_temperature, web_bulb @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='dry_bulb_temperature') +@preprocess_xarray('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 +1270,8 @@ def psychrometric_vapor_pressure_wet(pressure, dry_bulb_temperature, wet_bulb_te @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('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 +1311,8 @@ def mixing_ratio_from_relative_humidity(pressure, temperature, relative_humidity @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('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,6 +1350,7 @@ def relative_humidity_from_mixing_ratio(pressure, temperature, mixing_ratio): @exporter.export +@wrap_output_like(argument='specific_humidity') @preprocess_xarray() @check_units('[dimensionless]') def mixing_ratio_from_specific_humidity(specific_humidity): @@ -1367,6 +1388,7 @@ def mixing_ratio_from_specific_humidity(specific_humidity): @exporter.export +@wrap_output_like(argument='mixing_ratio') @preprocess_xarray() @check_units('[dimensionless]') def specific_humidity_from_mixing_ratio(mixing_ratio): @@ -1404,7 +1426,8 @@ def specific_humidity_from_mixing_ratio(mixing_ratio): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('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. @@ -2034,7 +2057,8 @@ def mixed_layer(pressure, *args, height=None, bottom=None, depth=100 * units.hPa @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('height', 'temperature') @check_units('[length]', '[temperature]') def dry_static_energy(height, temperature): r"""Calculate the dry static energy of parcels. @@ -2066,7 +2090,8 @@ def dry_static_energy(height, temperature): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('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. @@ -2225,7 +2250,8 @@ def thickness_hydrostatic_from_relative_humidity(pressure, temperature, relative @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='potential_temperature') +@preprocess_xarray('height', 'potential_temperature') @check_units('[length]', '[temperature]') def brunt_vaisala_frequency_squared(height, potential_temperature, axis=0): r"""Calculate the square of the Brunt-Vaisala frequency. @@ -2265,7 +2291,8 @@ def brunt_vaisala_frequency_squared(height, potential_temperature, axis=0): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='potential_temperature') +@preprocess_xarray('height', 'potential_temperature') @check_units('[length]', '[temperature]') def brunt_vaisala_frequency(height, potential_temperature, axis=0): r"""Calculate the Brunt-Vaisala frequency. @@ -2306,7 +2333,8 @@ def brunt_vaisala_frequency(height, potential_temperature, axis=0): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='potential_temperature') +@preprocess_xarray('height', 'potential_temperature') @check_units('[length]', '[temperature]') def brunt_vaisala_period(height, potential_temperature, axis=0): r"""Calculate the Brunt-Vaisala period. @@ -2345,7 +2373,8 @@ def brunt_vaisala_period(height, potential_temperature, axis=0): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('pressure', 'temperature', 'dewpoint') @check_units('[pressure]', '[temperature]', '[temperature]') def wet_bulb_temperature(pressure, temperature, dewpoint): """Calculate the wet-bulb temperature using Normand's rule. @@ -2398,7 +2427,8 @@ def wet_bulb_temperature(pressure, temperature, dewpoint): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('pressure', 'temperature') @check_units('[pressure]', '[temperature]') def static_stability(pressure, temperature, axis=0): r"""Calculate the static stability within a vertical profile. @@ -2430,7 +2460,8 @@ def static_stability(pressure, temperature, axis=0): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('pressure', 'temperature', 'specific_humidity') @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 +2491,8 @@ def dewpoint_from_specific_humidity(pressure, temperature, specific_humidity): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('pressure', 'temperature', 'w') @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 +2535,8 @@ def vertical_velocity_pressure(w, pressure, temperature, mixing_ratio=0): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='temperature') +@preprocess_xarray('pressure', 'temperature', 'omega') @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 +2582,8 @@ def vertical_velocity(omega, pressure, temperature, mixing_ratio=0): @exporter.export -@preprocess_xarray() +@wrap_output_like(argument='dewpoint') +@preprocess_xarray('pressure', 'dewpoint') @check_units('[pressure]', '[temperature]') def specific_humidity_from_dewpoint(pressure, dewpoint): r"""Calculate the specific humidity from the dewpoint temperature and pressure. diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index 0c95ea9547e..c764eff36a0 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -38,6 +38,138 @@ DIR_DICT[UND] = np.nan +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`` + + (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) + iterative : bool + if True, will iterate over output instead of working with output itself (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.""" + match_units = str(getattr(match, 'units', '')) + output_xarray = isinstance(match, xr.DataArray) + if isinstance(result, xr.DataArray): + result.metpy.convert_units(match_units) + return result 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(result, dims=match.dims, coords=match.coords, + attrs={'units': match_units}) + 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 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(result_magnitude, dims=match.dims, coords=match.coords, + attrs={'units': result_units}) + else: + return units.Quantity(result_magnitude, result_units) + + @exporter.export @preprocess_xarray() def resample_nn_1d(a, centers): @@ -447,6 +579,10 @@ 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 + ----- + This function does not return xarray DataArrays in the current version of MetPy. + """ # Make sure pressure and datavars are the same length for datavar in args: @@ -543,6 +679,10 @@ 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 + ----- + This function does not return xarray DataArrays in the current version of MetPy. + """ # If we get the depth kwarg, but it's None, set it to the default as well if depth is None: @@ -781,6 +921,8 @@ 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 explictly will only return Pint Quantities. + """ from pyproj import Geod @@ -1297,6 +1439,7 @@ def _process_deriv_args(f, axis, x, delta): @exporter.export +@wrap_output_like(argument='input_dir') @preprocess_xarray() def parse_angle(input_dir): """Calculate the meteorological angle from directional text. @@ -1311,7 +1454,7 @@ def parse_angle(input_dir): Returns ------- - `pint.Quantity` + `pint.Quantity` or `xarray.DataArray` The angle in degrees """ @@ -1374,6 +1517,10 @@ def angle_to_direction(input_angle, full=False, level=3): direction The directional text + Notes + ----- + This will not return a `pint.Quantity` or `xarray.DataArray`. + """ try: # strip units temporarily origin_units = input_angle.units @@ -1476,132 +1623,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`` - - (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.""" - match_units = str(getattr(match, 'units', '')) - output_xarray = isinstance(match, xr.DataArray) - if isinstance(result, xr.DataArray): - result.metpy.convert_units(match_units) - return result 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(result, dims=match.dims, coords=match.coords, - attrs={'units': match_units}) - 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 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(result_magnitude, dims=match.dims, coords=match.coords, - attrs={'units': result_units}) - else: - return units.Quantity(result_magnitude, result_units) diff --git a/src/metpy/interpolate/one_dimension.py b/src/metpy/interpolate/one_dimension.py index a825d08e1f8..4be8fac080a 100644 --- a/src/metpy/interpolate/one_dimension.py +++ b/src/metpy/interpolate/one_dimension.py @@ -14,7 +14,6 @@ @exporter.export - @preprocess_xarray() def interpolate_nans_1d(x, y, kind='linear'): """Interpolate NaN values in y. @@ -95,6 +94,9 @@ def interpolate_1d(x, xp, *args, axis=0, fill_value=np.nan, return_list_always=F Notes ----- xp and args must be the same shape. + + This function does not return xarray DataArrays. Use upstream interpolation utilities + until MetPy gains an xarray-compatable implementation. """ # Handle units @@ -216,6 +218,9 @@ def log_interpolate_1d(x, xp, *args, axis=0, fill_value=np.nan): Notes ----- xp and args must be the same shape. + + This function does not return xarray DataArrays. Use upstream interpolation utilities + until MetPy gains an xarray-compatable implementation. """ # Handle units diff --git a/tests/calc/test_thermo.py b/tests/calc/test_thermo.py index 5f9b5f9c461..2d0f8f99e90 100644 --- a/tests/calc/test_thermo.py +++ b/tests/calc/test_thermo.py @@ -54,7 +54,8 @@ def test_relative_humidity_from_dewpoint_xarray(): """Test Relative Humidity calculation with xarray data arrays.""" temp = xr.DataArray(25., attrs={'units': 'degC'}) dewp = xr.DataArray(15., attrs={'units': 'degC'}) - assert_almost_equal(relative_humidity_from_dewpoint(temp, dewp), 53.80 * units.percent, 2) + assert_almost_equal(relative_humidity_from_dewpoint(temp, dewp).metpy.unit_array, + 53.80 * units.percent, 2) def test_exner_function(): diff --git a/tests/test_xarray.py b/tests/test_xarray.py index 6583762f725..332792faa9e 100644 --- a/tests/test_xarray.py +++ b/tests/test_xarray.py @@ -180,7 +180,7 @@ def func(a, b): def test_preprocess_xarray_with_broadcasting(): - """test xarray preprocessing decorator with arguments to broadcast specified.""" + """Test xarray preprocessing decorator with arguments to broadcast specified.""" data = xr.DataArray(np.arange(9).reshape((3, 3)), dims=('y', 'x'), attrs={'units': 'N'}) data2 = xr.DataArray([1, 0, 0], dims=('y'), attrs={'units': 'm'}) From 74a97c2020ca38fbb9e102995c660f9fa661edc1 Mon Sep 17 00:00:00 2001 From: Ryan May Date: Tue, 14 Jan 2020 07:07:47 -0500 Subject: [PATCH 4/4] Fix lint --- src/metpy/calc/tools.py | 4 ++-- src/metpy/interpolate/one_dimension.py | 4 ++-- src/metpy/xarray.py | 1 + 3 files changed, 5 insertions(+), 4 deletions(-) diff --git a/src/metpy/calc/tools.py b/src/metpy/calc/tools.py index c764eff36a0..6fcaff3cc09 100644 --- a/src/metpy/calc/tools.py +++ b/src/metpy/calc/tools.py @@ -56,7 +56,7 @@ def wrap_output_like(**wrap_kwargs): * ``xarray.DataArray`` (if matched output is not one of these types, we instead treat the match as if it was a - dimenionless Quantity.) + dimsenionless Quantity.) This wrapping/conversion follows the following rules: @@ -66,7 +66,7 @@ def wrap_output_like(**wrap_kwargs): * 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 + - If match_unit is True, for output of Quantity or DataArray respectively, with a given unit, * ndarray becomes Quantity or DataArray (with matching coords) with output unit diff --git a/src/metpy/interpolate/one_dimension.py b/src/metpy/interpolate/one_dimension.py index 4be8fac080a..01b450f33ca 100644 --- a/src/metpy/interpolate/one_dimension.py +++ b/src/metpy/interpolate/one_dimension.py @@ -94,7 +94,7 @@ def interpolate_1d(x, xp, *args, axis=0, fill_value=np.nan, return_list_always=F Notes ----- xp and args must be the same shape. - + This function does not return xarray DataArrays. Use upstream interpolation utilities until MetPy gains an xarray-compatable implementation. @@ -218,7 +218,7 @@ def log_interpolate_1d(x, xp, *args, axis=0, fill_value=np.nan): Notes ----- xp and args must be the same shape. - + This function does not return xarray DataArrays. Use upstream interpolation utilities until MetPy gains an xarray-compatable implementation. diff --git a/src/metpy/xarray.py b/src/metpy/xarray.py index 67d218d133b..cfcf9d613b5 100644 --- a/src/metpy/xarray.py +++ b/src/metpy/xarray.py @@ -875,6 +875,7 @@ def check_axis(var, *axes): # If no match has been made, return False (rather than None) return False + def _assign_crs(xarray_object, cf_attributes, cf_kwargs): from .plots.mapping import CFProjection