From 7524ffcd9753a26a45d3f259ade1ad4cc298b40f Mon Sep 17 00:00:00 2001 From: Jon Thielen Date: Wed, 28 Aug 2019 08:03:18 -0500 Subject: [PATCH] Updates to pint and xarray usage for future version compatability There were *many* places in the code base where passing pint Quantities to numpy functions was assumed to return a plain ndarray. This behavior will no longer be the case after pint implements __array_function__ in its next version, and so, this modifies usage to explicitly request magnitude where it is assumed that quantities are stripped. Also, makes some miscelaneous changes with regards to the split between scalar and sequence quantity classes. Updates test tolerances to account for changes to base constants from pint v0.9 defaults to v0.10 defaults (except for geopotential <-> height conversions which are problematic and resolved elsewhere). Finally, resolves issues with xarray v0.13.0 (no implicit casting to ndarray or inplace argument anymore). --- examples/cross_section.py | 3 +- metpy/calc/basic.py | 5 +-- metpy/calc/indices.py | 10 +++--- metpy/calc/kinematics.py | 2 +- metpy/calc/tests/test_basic.py | 2 +- metpy/calc/tests/test_calc_tools.py | 8 ++--- metpy/calc/tests/test_kinematics.py | 12 +++---- metpy/calc/tests/test_thermo.py | 50 ++++++++++++++--------------- metpy/calc/thermo.py | 34 ++++++++++---------- metpy/calc/tools.py | 33 ++++++++++--------- metpy/cbook.py | 16 ++++++++- metpy/plots/skewt.py | 8 ++--- metpy/tests/test_cbook.py | 30 ++++++++++++++++- metpy/tests/test_xarray.py | 9 +++++- metpy/units.py | 7 ++-- metpy/xarray.py | 4 +-- tutorials/xarray_tutorial.py | 4 +-- 17 files changed, 146 insertions(+), 91 deletions(-) diff --git a/examples/cross_section.py b/examples/cross_section.py index 5388ad5d1ed..af9cebefb23 100644 --- a/examples/cross_section.py +++ b/examples/cross_section.py @@ -42,8 +42,7 @@ ############################## # Get the cross section, and convert lat/lon to supplementary coordinates: -cross = cross_section(data, start, end) -cross.set_coords(('lat', 'lon'), True) +cross = cross_section(data, start, end).set_coords(('lat', 'lon')) print(cross) ############################## diff --git a/metpy/calc/basic.py b/metpy/calc/basic.py index cadaa1cca5c..cbf4859b033 100644 --- a/metpy/calc/basic.py +++ b/metpy/calc/basic.py @@ -931,8 +931,9 @@ def altimeter_to_station_pressure(altimeter_value, height): # N-Value n = (mpconsts.Rd * gamma / mpconsts.g).to_base_units() - return ((altimeter_value ** n - ((p0 ** n * gamma * height) / t0)) ** (1 / n) + ( - 0.3 * units.hPa)) + return ((altimeter_value ** n + - ((p0.to(altimeter_value.units) ** n * gamma * height) / t0)) ** (1 / n) + + 0.3 * units.hPa) @exporter.export diff --git a/metpy/calc/indices.py b/metpy/calc/indices.py index 072e498352d..a802f00067d 100644 --- a/metpy/calc/indices.py +++ b/metpy/calc/indices.py @@ -49,10 +49,10 @@ def precipitable_water(dewpt, pressure, bottom=None, top=None): dewpt = dewpt[sort_inds] if top is None: - top = np.nanmin(pressure) * pressure.units + top = np.nanmin(pressure.magnitude) * pressure.units if bottom is None: - bottom = np.nanmax(pressure) * pressure.units + bottom = np.nanmax(pressure.magnitude) * pressure.units pres_layer, dewpt_layer = get_layer(pressure, dewpt, bottom=bottom, depth=bottom - top) @@ -109,7 +109,8 @@ def mean_pressure_weighted(pressure, *args, **kwargs): # function. Said integral works out to this function: pres_int = 0.5 * (layer_p[-1].magnitude**2 - layer_p[0].magnitude**2) for i, datavar in enumerate(args): - arg_mean = np.trapz(layer_arg[i] * layer_p, x=layer_p) / pres_int + arg_mean = np.trapz((layer_arg[i] * layer_p).magnitude, + x=layer_p.magnitude) / pres_int ret.append(arg_mean * datavar.units) return ret @@ -163,7 +164,8 @@ def bunkers_storm_motion(pressure, u, v, heights): # Take the cross product of the wind shear and k, and divide by the vector magnitude and # multiply by the deviaton empirically calculated in Bunkers (2000) (7.5 m/s) shear_cross = concatenate([shear[1], -shear[0]]) - rdev = shear_cross * (7.5 * units('m/s').to(u.units) / np.hypot(*shear)) + shear_mag = np.hypot(*(arg.magnitude for arg in shear)) * shear.units + rdev = shear_cross * (7.5 * units('m/s').to(u.units) / shear_mag) # Add the deviations to the layer average wind to get the RM motion right_mover = wind_mean + rdev diff --git a/metpy/calc/kinematics.py b/metpy/calc/kinematics.py index 84b56c56574..4bb26227b73 100644 --- a/metpy/calc/kinematics.py +++ b/metpy/calc/kinematics.py @@ -21,7 +21,7 @@ def _stack(arrs): - return concatenate([a[np.newaxis] for a in arrs], axis=0) + return concatenate([a[np.newaxis] if iterable(a) else a for a in arrs], axis=0) def _is_x_first_dim(dim_order): diff --git a/metpy/calc/tests/test_basic.py b/metpy/calc/tests/test_basic.py index d08a1398345..deaa5d2d55d 100644 --- a/metpy/calc/tests/test_basic.py +++ b/metpy/calc/tests/test_basic.py @@ -396,7 +396,7 @@ def test_add_height_to_pressure(): def test_add_pressure_to_height(): """Test the height at pressure above height calculation.""" height = add_pressure_to_height(110.8286757 * units.m, 100 * units.hPa) - assert_almost_equal(height, 988.0028867 * units.meter, 5) + assert_almost_equal(height, 988.0028867 * units.meter, 3) def test_sigma_to_pressure(): diff --git a/metpy/calc/tests/test_calc_tools.py b/metpy/calc/tests/test_calc_tools.py index 96308ff13bc..579a9681d48 100644 --- a/metpy/calc/tests/test_calc_tools.py +++ b/metpy/calc/tests/test_calc_tools.py @@ -275,8 +275,8 @@ def get_bounds_data(): def test_get_bound_pressure_height(pressure, bound, hgts, interp, expected): """Test getting bounds in layers with various parameter combinations.""" bounds = _get_bound_pressure_height(pressure, bound, heights=hgts, interpolate=interp) - assert_array_almost_equal(bounds[0], expected[0], 5) - assert_array_almost_equal(bounds[1], expected[1], 5) + assert_array_almost_equal(bounds[0], expected[0], 4) + assert_array_almost_equal(bounds[1], expected[1], 4) def test_get_bound_invalid_bound_units(): @@ -363,8 +363,8 @@ def test_get_layer(pressure, variable, heights, bottom, depth, interp, expected) """Test get_layer functionality.""" p_layer, y_layer = get_layer(pressure, variable, heights=heights, bottom=bottom, depth=depth, interpolate=interp) - assert_array_almost_equal(p_layer, expected[0], 5) - assert_array_almost_equal(y_layer, expected[1], 5) + assert_array_almost_equal(p_layer, expected[0], 4) + assert_array_almost_equal(y_layer, expected[1], 4) def test_greater_or_close(): diff --git a/metpy/calc/tests/test_kinematics.py b/metpy/calc/tests/test_kinematics.py index 5b84b167eaa..c5a30c46a9c 100644 --- a/metpy/calc/tests/test_kinematics.py +++ b/metpy/calc/tests/test_kinematics.py @@ -1086,8 +1086,8 @@ def test_q_vector_without_static_stability(q_vector_data): [8.6038280e-14, 4.6968342e-13, -4.6968342e-13, -8.6038280e-14]]) * units('m^2 kg^-1 s^-1')) - assert_almost_equal(q1, q1_truth, 20) - assert_almost_equal(q2, q2_truth, 20) + assert_almost_equal(q1, q1_truth, 18) + assert_almost_equal(q2, q2_truth, 18) def test_q_vector_with_static_stability(q_vector_data): @@ -1110,8 +1110,8 @@ def test_q_vector_with_static_stability(q_vector_data): [4.4714390e-08, 2.3490489e-07, -2.2635441e-07, -4.0003646e-08]]) * units('kg m^-2 s^-3')) - assert_almost_equal(q1, q1_truth, 14) - assert_almost_equal(q2, q2_truth, 14) + assert_almost_equal(q1, q1_truth, 12) + assert_almost_equal(q2, q2_truth, 12) @pytest.fixture @@ -1726,5 +1726,5 @@ def test_q_vector_4d(data_4d): [-1.06658890e-13, -2.19817426e-13, -8.35968065e-14, 1.88190788e-13], [-2.27182863e-13, -2.74607819e-13, -1.10587309e-13, -3.88915866e-13]]]]) * units('m^2 kg^-1 s^-1') - assert_array_almost_equal(q1, q1_truth, 20) - assert_array_almost_equal(q2, q2_truth, 20) + assert_array_almost_equal(q1, q1_truth, 18) + assert_array_almost_equal(q2, q2_truth, 18) diff --git a/metpy/calc/tests/test_thermo.py b/metpy/calc/tests/test_thermo.py index b2c1c5fd655..a16c4121b02 100644 --- a/metpy/calc/tests/test_thermo.py +++ b/metpy/calc/tests/test_thermo.py @@ -710,8 +710,8 @@ def test_cape_cin(): dewpoint = np.array([19., -11.2, -10.8, -10.4, -10., -53.2]) * units.celsius parcel_prof = parcel_profile(p, temperature[0], dewpoint[0]) cape, cin = cape_cin(p, temperature, dewpoint, parcel_prof) - assert_almost_equal(cape, 75.7340825 * units('joule / kilogram'), 6) - assert_almost_equal(cin, -89.8179205 * units('joule / kilogram'), 6) + assert_almost_equal(cape, 75.7340825 * units('joule / kilogram'), 2) + assert_almost_equal(cin, -89.8179205 * units('joule / kilogram'), 2) def test_cape_cin_no_el(): @@ -721,8 +721,8 @@ def test_cape_cin_no_el(): dewpoint = np.array([19., -11.2, -10.8, -10.4]) * units.celsius parcel_prof = parcel_profile(p, temperature[0], dewpoint[0]).to('degC') cape, cin = cape_cin(p, temperature, dewpoint, parcel_prof) - assert_almost_equal(cape, 0.08610409 * units('joule / kilogram'), 6) - assert_almost_equal(cin, -89.8179205 * units('joule / kilogram'), 6) + assert_almost_equal(cape, 0.08610409 * units('joule / kilogram'), 2) + assert_almost_equal(cin, -89.8179205 * units('joule / kilogram'), 2) def test_cape_cin_no_lfc(): @@ -732,8 +732,8 @@ def test_cape_cin_no_lfc(): dewpoint = np.array([19., -11.2, -10.8, -10.4, -10., -53.2]) * units.celsius parcel_prof = parcel_profile(p, temperature[0], dewpoint[0]).to('degC') cape, cin = cape_cin(p, temperature, dewpoint, parcel_prof) - assert_almost_equal(cape, 0.0 * units('joule / kilogram'), 6) - assert_almost_equal(cin, 0.0 * units('joule / kilogram'), 6) + assert_almost_equal(cape, 0.0 * units('joule / kilogram'), 2) + assert_almost_equal(cin, 0.0 * units('joule / kilogram'), 2) def test_find_append_zero_crossings(): @@ -980,8 +980,8 @@ def test_surface_based_cape_cin(): temperature = np.array([22.2, 14.6, 12., 9.4, 7., -38.]) * units.celsius dewpoint = np.array([19., -11.2, -10.8, -10.4, -10., -53.2]) * units.celsius cape, cin = surface_based_cape_cin(p, temperature, dewpoint) - assert_almost_equal(cape, 75.7340825 * units('joule / kilogram'), 6) - assert_almost_equal(cin, -136.607809 * units('joule / kilogram'), 6) + assert_almost_equal(cape, 75.7340825 * units('joule / kilogram'), 2) + assert_almost_equal(cin, -136.607809 * units('joule / kilogram'), 2) def test_most_unstable_cape_cin_surface(): @@ -990,8 +990,8 @@ def test_most_unstable_cape_cin_surface(): temperature = np.array([22.2, 14.6, 12., 9.4, 7., -38.]) * units.celsius dewpoint = np.array([19., -11.2, -10.8, -10.4, -10., -53.2]) * units.celsius mucape, mucin = most_unstable_cape_cin(pressure, temperature, dewpoint) - assert_almost_equal(mucape, 75.7340825 * units('joule / kilogram'), 6) - assert_almost_equal(mucin, -136.607809 * units('joule / kilogram'), 6) + assert_almost_equal(mucape, 75.7340825 * units('joule / kilogram'), 2) + assert_almost_equal(mucin, -136.607809 * units('joule / kilogram'), 2) def test_most_unstable_cape_cin(): @@ -1151,7 +1151,7 @@ def test_wet_bulb_temperature(): """Test wet bulb calculation with scalars.""" val = wet_bulb_temperature(1000 * units.hPa, 25 * units.degC, 15 * units.degC) truth = 18.34345936 * units.degC # 18.59 from NWS calculator - assert_almost_equal(val, truth) + assert_almost_equal(val, truth, 5) def test_wet_bulb_temperature_1d(): @@ -1162,7 +1162,7 @@ def test_wet_bulb_temperature_1d(): val = wet_bulb_temperature(pressures, temperatures, dewpoints) truth = [21.4449794, 16.7368576, 12.0656909] * units.degC # 21.58, 16.86, 12.18 from NWS Calculator - assert_array_almost_equal(val, truth) + assert_array_almost_equal(val, truth, 5) def test_wet_bulb_temperature_2d(): @@ -1178,7 +1178,7 @@ def test_wet_bulb_temperature_2d(): [20.5021631, 15.801218, 11.1361878]] * units.degC # 21.58, 16.86, 12.18 # 20.6, 15.9, 11.2 from NWS Calculator - assert_array_almost_equal(val, truth) + assert_array_almost_equal(val, truth, 5) def test_static_stability_adiabatic(): @@ -1229,8 +1229,8 @@ def test_lfc_not_below_lcl(): 8.6, 8.1, 7.6, 7., 6.5, 6., 5.4]) * units.degC lfc_pressure, lfc_temp = lfc(levels, temperatures, dewpoints) # Before patch, LFC pressure would show 1000.5912165339967 hPa - assert_almost_equal(lfc_pressure, 811.8263397 * units.mbar, 6) - assert_almost_equal(lfc_temp, 6.4992871 * units.celsius, 6) + assert_almost_equal(lfc_pressure, 811.8263397 * units.mbar, 3) + assert_almost_equal(lfc_temp, 6.4992871 * units.celsius, 3) def test_multiple_lfcs(): @@ -1258,10 +1258,10 @@ def test_multiple_lfcs(): lfc_pressure_bottom, lfc_temp_bottom = lfc(levels, temperatures, dewpoints, which='bottom') lfc_pressure_all, _ = lfc(levels, temperatures, dewpoints, which='all') - assert_almost_equal(lfc_pressure_top, 705.4346277 * units.mbar, 6) - assert_almost_equal(lfc_temp_top, 4.8922235 * units.degC, 6) - assert_almost_equal(lfc_pressure_bottom, 884.1954356 * units.mbar, 6) - assert_almost_equal(lfc_temp_bottom, 13.9597571 * units.degC, 6) + assert_almost_equal(lfc_pressure_top, 705.4346277 * units.mbar, 3) + assert_almost_equal(lfc_temp_top, 4.8922235 * units.degC, 3) + assert_almost_equal(lfc_pressure_bottom, 884.1954356 * units.mbar, 3) + assert_almost_equal(lfc_temp_bottom, 13.9597571 * units.degC, 3) assert_almost_equal(len(lfc_pressure_all), 2, 0) @@ -1289,10 +1289,10 @@ def test_multiple_els(): el_pressure_top, el_temp_top = el(levels, temperatures, dewpoints) el_pressure_bottom, el_temp_bottom = el(levels, temperatures, dewpoints, which='bottom') el_pressure_all, _ = el(levels, temperatures, dewpoints, which='all') - assert_almost_equal(el_pressure_top, 228.0575059 * units.mbar, 6) - assert_almost_equal(el_temp_top, -56.8123126 * units.degC, 6) - assert_almost_equal(el_pressure_bottom, 849.7942185 * units.mbar, 6) - assert_almost_equal(el_temp_bottom, 12.4233265 * units.degC, 6) + assert_almost_equal(el_pressure_top, 228.0575059 * units.mbar, 3) + assert_almost_equal(el_temp_top, -56.8123126 * units.degC, 3) + assert_almost_equal(el_pressure_bottom, 849.7942185 * units.mbar, 3) + assert_almost_equal(el_temp_bottom, 12.4233265 * units.degC, 3) assert_almost_equal(len(el_pressure_all), 2, 0) @@ -1303,8 +1303,8 @@ def test_cape_cin_custom_profile(): dewpoint = np.array([19., -11.2, -10.8, -10.4, -10., -53.2]) * units.celsius parcel_prof = parcel_profile(p, temperature[0], dewpoint[0]) + 5 * units.delta_degC cape, cin = cape_cin(p, temperature, dewpoint, parcel_prof) - assert_almost_equal(cape, 1443.505086499895 * units('joule / kilogram'), 6) - assert_almost_equal(cin, 0.0 * units('joule / kilogram'), 6) + assert_almost_equal(cape, 1443.505086499895 * units('joule / kilogram'), 2) + assert_almost_equal(cin, 0.0 * units('joule / kilogram'), 2) def test_parcel_profile_below_lcl(): diff --git a/metpy/calc/thermo.py b/metpy/calc/thermo.py index d18c5f54ca4..00628858faa 100644 --- a/metpy/calc/thermo.py +++ b/metpy/calc/thermo.py @@ -264,7 +264,7 @@ def dt(t, p): frac = ((mpconsts.Rd * t + mpconsts.Lv * rs) / (mpconsts.Cp_d + (mpconsts.Lv * mpconsts.Lv * rs * mpconsts.epsilon / (mpconsts.Rd * t * t)))).to('kelvin') - return frac / p + return (frac / p).magnitude if ref_pressure is None: ref_pressure = pressure[0] @@ -287,14 +287,14 @@ def dt(t, p): if ref_pressure > pressure.min(): # Integrate downward in pressure - pres_down = np.append(ref_pressure, pressure[(ref_pres_idx - 1)::-1]) - trace_down = si.odeint(dt, temperature.squeeze(), pres_down.squeeze()) + pres_down = np.append(ref_pressure.m, pressure[(ref_pres_idx - 1)::-1].m) + trace_down = si.odeint(dt, temperature.m.squeeze(), pres_down.squeeze()) ret_temperatures = np.concatenate((ret_temperatures, trace_down[:0:-1])) if ref_pressure < pressure.max(): # Integrate upward in pressure - pres_up = np.append(ref_pressure, pressure[ref_pres_idx:]) - trace_up = si.odeint(dt, temperature.squeeze(), pres_up.squeeze()) + pres_up = np.append(ref_pressure.m, pressure[ref_pres_idx:].m) + trace_up = si.odeint(dt, temperature.m.squeeze(), pres_up.squeeze()) ret_temperatures = np.concatenate((ret_temperatures, trace_up[1:])) if pres_decreasing: @@ -628,7 +628,7 @@ def _parcel_profile_helper(pressure, temperature, dewpt): temp_lower = dry_lapse(press_lower, temperature) # If the pressure profile doesn't make it to the lcl, we can stop here - if _greater_or_close(np.nanmin(pressure), press_lcl.m): + if _greater_or_close(np.nanmin(pressure.m), press_lcl.m): return (press_lower[:-1], press_lcl, np.array([]) * press_lower.units, temp_lower[:-1], temp_lcl, np.array([]) * temp_lower.units) @@ -1449,17 +1449,17 @@ def cape_cin(pressure, temperature, dewpt, parcel_profile): # CAPE # Only use data between the LFC and EL for calculation - p_mask = _less_or_close(x, lfc_pressure) & _greater_or_close(x, el_pressure) - x_clipped = x[p_mask] - y_clipped = y[p_mask] + p_mask = _less_or_close(x.m, lfc_pressure) & _greater_or_close(x.m, el_pressure) + x_clipped = x[p_mask].magnitude + y_clipped = y[p_mask].magnitude cape = (mpconsts.Rd * (np.trapz(y_clipped, np.log(x_clipped)) * units.degK)).to(units('J/kg')) # CIN # Only use data between the surface and LFC for calculation - p_mask = _greater_or_close(x, lfc_pressure) - x_clipped = x[p_mask] - y_clipped = y[p_mask] + p_mask = _greater_or_close(x.m, lfc_pressure) + x_clipped = x[p_mask].magnitude + y_clipped = y[p_mask].magnitude cin = (mpconsts.Rd * (np.trapz(y_clipped, np.log(x_clipped)) * units.degK)).to(units('J/kg')) @@ -1498,7 +1498,7 @@ def _find_append_zero_crossings(x, y): y = y[sort_idx] # Remove duplicate data points if there are any - keep_idx = np.ediff1d(x, to_end=[1]) > 1e-6 + keep_idx = np.ediff1d(x.magnitude, to_end=[1]) > 1e-6 x = x[keep_idx] y = y[keep_idx] return x, y @@ -1638,7 +1638,7 @@ def _isen_iter(iter_log_p, isentlevs_nd, ka, a, b, pok): slices = [np.newaxis] * ndim slices[axis] = slice(None) slices = tuple(slices) - pres = np.broadcast_to(pres[slices], temperature.shape) * pres.units + pres = np.broadcast_to(pres[slices].magnitude, temperature.shape) * pres.units # Sort input data sort_pres = np.argsort(pres.m, axis=axis) @@ -1904,7 +1904,7 @@ def mixed_layer(p, *args, **kwargs): ret = [] for datavar_layer in datavars_layer: actual_depth = abs(p_layer[0] - p_layer[-1]) - ret.append((-1. / actual_depth.m) * np.trapz(datavar_layer, p_layer) + ret.append((-1. / actual_depth.m) * np.trapz(datavar_layer.m, p_layer.m) * datavar_layer.units) return ret @@ -2047,7 +2047,7 @@ def thickness_hydrostatic(pressure, temperature, **kwargs): # Take the integral (with unit handling) and return the result in meters return (- mpconsts.Rd / mpconsts.g * np.trapz( - layer_virttemp.to('K'), x=np.log(layer_p / units.hPa)) * units.K).to('m') + layer_virttemp.m_as('K'), x=np.log(layer_p.m_as('hPa'))) * units.K).to('m') @exporter.export @@ -2306,7 +2306,7 @@ def static_stability(pressure, temperature, axis=0): """ theta = potential_temperature(pressure, temperature) - return - mpconsts.Rd * temperature / pressure * first_derivative(np.log(theta / units.K), + return - mpconsts.Rd * temperature / pressure * first_derivative(np.log(theta.m_as('K')), x=pressure, axis=axis) diff --git a/metpy/calc/tools.py b/metpy/calc/tools.py index 06b99bbed1d..b42df73002f 100644 --- a/metpy/calc/tools.py +++ b/metpy/calc/tools.py @@ -20,7 +20,7 @@ def normalize_axis_index(a, n): import xarray as xr from . import height_to_pressure_std, pressure_to_height_std -from ..cbook import broadcast_indices +from ..cbook import broadcast_indices, result_type from ..deprecation import deprecated, metpyDeprecation from ..interpolate.one_dimension import interpolate_1d, interpolate_nans_1d, log_interpolate_1d from ..package_tools import Exporter @@ -398,8 +398,9 @@ def _get_bound_pressure_height(pressure, bound, heights=None, interpolate=True): # Need to cast back to the input type since interp (up to at least numpy # 1.13 always returns float64. This can cause upstream users problems, # resulting in something like np.append() to upcast. - bound_pressure = np.interp(np.atleast_1d(bound), heights, - pressure).astype(bound.dtype) * pressure.units + bound_pressure = (np.interp(np.atleast_1d(bound.m), heights.m, + pressure.m).astype(result_type(bound)) + * pressure.units) else: idx = (np.abs(heights - bound)).argmin() bound_pressure = pressure[idx] @@ -419,13 +420,15 @@ def _get_bound_pressure_height(pressure, bound, heights=None, interpolate=True): raise ValueError('Bound must be specified in units of length or pressure.') # If the bound is out of the range of the data, we shouldn't extrapolate - if not (_greater_or_close(bound_pressure, np.nanmin(pressure) * pressure.units) - and _less_or_close(bound_pressure, np.nanmax(pressure) * pressure.units)): + if not (_greater_or_close(bound_pressure, np.nanmin(pressure.m) * pressure.units) + and _less_or_close(bound_pressure, np.nanmax(pressure.m) * pressure.units)): raise ValueError('Specified bound is outside pressure range.') - if heights is not None: - if not (_less_or_close(bound_height, np.nanmax(heights) * heights.units) - and _greater_or_close(bound_height, np.nanmin(heights) * heights.units)): - raise ValueError('Specified bound is outside height range.') + if heights is not None and not (_less_or_close(bound_height, + np.nanmax(heights.m) * heights.units) + and _greater_or_close(bound_height, + np.nanmin(heights.m) + * heights.units)): + raise ValueError('Specified bound is outside height range.') return bound_pressure, bound_height @@ -501,9 +504,9 @@ def get_layer_heights(heights, depth, *args, **kwargs): if interpolate: # If we don't have the bottom or top requested, append them if top not in heights_interp: - heights_interp = np.sort(np.append(heights_interp, top)) * heights.units + heights_interp = np.sort(np.append(heights_interp.m, top.m)) * heights.units if bottom not in heights_interp: - heights_interp = np.sort(np.append(heights_interp, bottom)) * heights.units + heights_interp = np.sort(np.append(heights_interp.m, bottom.m)) * heights.units ret.append(heights_interp) @@ -576,7 +579,7 @@ def get_layer(pressure, *args, **kwargs): # If the bottom is not specified, make it the surface pressure if bottom is None: - bottom = np.nanmax(pressure) * pressure.units + bottom = np.nanmax(pressure.m) * pressure.units bottom_pressure, bottom_height = _get_bound_pressure_height(pressure, bottom, heights=heights, @@ -608,9 +611,9 @@ def get_layer(pressure, *args, **kwargs): if interpolate: # If we don't have the bottom or top requested, append them if not np.any(np.isclose(top_pressure, p_interp)): - p_interp = np.sort(np.append(p_interp, top_pressure)) * pressure.units + p_interp = np.sort(np.append(p_interp.m, top_pressure.m)) * pressure.units if not np.any(np.isclose(bottom_pressure, p_interp)): - p_interp = np.sort(np.append(p_interp, bottom_pressure)) * pressure.units + p_interp = np.sort(np.append(p_interp.m, bottom_pressure.m)) * pressure.units ret.append(p_interp[::-1]) @@ -1299,7 +1302,7 @@ def _process_deriv_args(f, kwargs): diff_size[axis] -= 1 delta_units = getattr(delta, 'units', None) delta = np.broadcast_to(delta, diff_size, subok=True) - if delta_units is not None: + if not hasattr(delta, 'units') and delta_units is not None: delta = delta * delta_units else: delta = _broadcast_to_axis(delta, axis, n) diff --git a/metpy/cbook.py b/metpy/cbook.py index 6cee72a08ec..7a3d51fd177 100644 --- a/metpy/cbook.py +++ b/metpy/cbook.py @@ -109,4 +109,18 @@ def iterable(value): return np.iterable(value) -__all__ = ('Registry', 'broadcast_indices', 'get_test_data', 'is_string_like', 'iterable') +def result_type(value): + """Determine the type for numpy type casting in a pint-version-safe way.""" + try: + return np.result_type(value) + except TypeError: + if hasattr(value, 'dtype'): + return value.dtype + elif hasattr(value, 'magnitude'): + return np.result_type(value.magnitude) + else: + raise TypeError('Cannot determine dtype for type {}'.format(type(value))) + + +__all__ = ('Registry', 'broadcast_indices', 'get_test_data', 'is_string_like', 'iterable', + 'result_type') diff --git a/metpy/plots/skewt.py b/metpy/plots/skewt.py index 81374975545..48fe66a31cb 100644 --- a/metpy/plots/skewt.py +++ b/metpy/plots/skewt.py @@ -459,7 +459,7 @@ def plot_dry_adiabats(self, t0=None, p=None, **kwargs): # Assemble into data for plotting t = dry_lapse(p, t0[:, np.newaxis], 1000. * units.mbar).to(units.degC) - linedata = [np.vstack((ti, p)).T for ti in t] + linedata = [np.vstack((ti.m, p.m)).T for ti in t] # Add to plot kwargs.setdefault('colors', 'r') @@ -512,7 +512,7 @@ def plot_moist_adiabats(self, t0=None, p=None, **kwargs): # Assemble into data for plotting t = moist_lapse(p, t0[:, np.newaxis], 1000. * units.mbar).to(units.degC) - linedata = [np.vstack((ti, p)).T for ti in t] + linedata = [np.vstack((ti.m, p.m)).T for ti in t] # Add to plot kwargs.setdefault('colors', 'b') @@ -560,7 +560,7 @@ def plot_mixing_lines(self, w=None, p=None, **kwargs): # Assemble data for plotting td = dewpoint(vapor_pressure(p, w)) - linedata = [np.vstack((t, p)).T for t in td] + linedata = [np.vstack((t.m, p.m)).T for t in td] # Add to plot kwargs.setdefault('colors', 'g') @@ -885,7 +885,7 @@ def plot_colormapped(self, u, v, c, intervals=None, colors=None, **kwargs): # Find any intervals not in the data and interpolate them interpolation_heights = [bound.m for bound in intervals if bound not in c] interpolation_heights = np.array(interpolation_heights) * intervals.units - interpolation_heights = (np.sort(interpolation_heights) + interpolation_heights = (np.sort(interpolation_heights.magnitude) * interpolation_heights.units) (interpolated_heights, interpolated_u, interpolated_v) = interpolate_1d(interpolation_heights, c, c, u, v) diff --git a/metpy/tests/test_cbook.py b/metpy/tests/test_cbook.py index c94713a4526..ad491f6c0cd 100644 --- a/metpy/tests/test_cbook.py +++ b/metpy/tests/test_cbook.py @@ -3,7 +3,12 @@ # SPDX-License-Identifier: BSD-3-Clause """Test functionality of MetPy's utility code.""" -from metpy.cbook import Registry +import numpy as np +import pytest +import xarray as xr + +from metpy.cbook import Registry, result_type +from metpy.units import units def test_registry(): @@ -14,3 +19,26 @@ def test_registry(): reg.register('mine')(a) assert reg['mine'] is a + + +@pytest.mark.parametrize( + 'test_input, expected_type_match, custom_dtype', + [(1.0, 1.0, None), + (1, 1, None), + (np.array(1.0), 1.0, None), + (np.array(1), 1, None), + (np.array([1, 2, 3], dtype=np.int32), 1, 'int32'), + (units.Quantity(1, units.m), 1, None), + (units.Quantity(1.0, units.m), 1.0, None), + (units.Quantity([1, 2.0], units.m), 1.0, None), + ([1, 2, 3] * units.m, 1, None), + (xr.DataArray(data=[1, 2.0]), 1.0, None)]) +def test_result_type(test_input, expected_type_match, custom_dtype): + """Test result_type on the kinds of things common in MetPy.""" + assert result_type(test_input) == np.array(expected_type_match, dtype=custom_dtype).dtype + + +def test_result_type_failure(): + """Test result_type failure on non-numeric types.""" + with pytest.raises(TypeError): + result_type([False]) diff --git a/metpy/tests/test_xarray.py b/metpy/tests/test_xarray.py index 8618ff0f6d4..3112ebdb952 100644 --- a/metpy/tests/test_xarray.py +++ b/metpy/tests/test_xarray.py @@ -97,7 +97,7 @@ def test_convert_units(test_var): test_var.metpy.convert_units('degC') # Check that variable metadata is updated - assert test_var.attrs['units'] == 'degC' + assert units(test_var.attrs['units']) == units('degC') # Make sure we now get an array back with properly converted values assert test_var.metpy.unit_array.units == units.degC @@ -634,3 +634,10 @@ def test_coordinate_identification_shared_but_not_equal_coords(): # Check vertical coordinate on u # Fails prior to resolution of Issue #1124 assert ds['isobaric2'].identical(ds['u'].metpy.vertical) + + +def test_check_no_quantification_of_xarray_data(test_ds_generic): + """Test that .unit_array setter does not insert a `pint.Quantity` into the DataArray.""" + var = test_ds_generic['e'] + var.metpy.unit_array = [1000, 925, 850, 700, 500] * units.hPa + assert not isinstance(var.data, units.Quantity) diff --git a/metpy/units.py b/metpy/units.py index 08be1dc8ee8..0f33aaf4247 100644 --- a/metpy/units.py +++ b/metpy/units.py @@ -156,13 +156,14 @@ def diff(x, **kwargs): numpy.diff """ - ret = np.diff(x, **kwargs) if hasattr(x, 'units'): + ret = np.diff(x.magnitude, **kwargs) # Can't just use units because of how things like temperature work it = x.flat true_units = (next(it) - next(it)).units - ret = ret * true_units - return ret + return ret * true_units + else: + return np.diff(x, **kwargs) def atleast_1d(*arrs): diff --git a/metpy/xarray.py b/metpy/xarray.py index e3eac10d98c..58c5fb55055 100644 --- a/metpy/xarray.py +++ b/metpy/xarray.py @@ -99,8 +99,8 @@ def unit_array(self): @unit_array.setter def unit_array(self, values): - """Set data values as a `pint.Quantity`.""" - self._data_array.values = values + """Set data values from a `pint.Quantity`.""" + self._data_array.values = values.magnitude self._units = self._data_array.attrs['units'] = str(values.units) def convert_units(self, units): diff --git a/tutorials/xarray_tutorial.py b/tutorials/xarray_tutorial.py index 29307148c75..7e18b72c97d 100644 --- a/tutorials/xarray_tutorial.py +++ b/tutorials/xarray_tutorial.py @@ -64,14 +64,14 @@ 'v-component_of_wind_isobaric']) # To rename variables, supply a dictionary between old and new names to the rename method -data.rename({ +data = data.rename({ 'Vertical_velocity_pressure_isobaric': 'omega', 'Relative_humidity_isobaric': 'relative_humidity', 'Temperature_isobaric': 'temperature', 'u-component_of_wind_isobaric': 'u', 'v-component_of_wind_isobaric': 'v', 'Geopotential_height_isobaric': 'height' -}, inplace=True) +}) ######################################################################### # Units