From e8da3b91c2d3f81fee6e18f3b8e5b2527c36e858 Mon Sep 17 00:00:00 2001 From: rpmanser Date: Wed, 8 Sep 2021 19:19:40 -0500 Subject: [PATCH] Convert instances of boolean indexing to ndarrays to pass dask tests --- src/metpy/calc/basic.py | 21 +++++++++++---------- tests/calc/test_basic.py | 23 +++++++++-------------- 2 files changed, 20 insertions(+), 24 deletions(-) diff --git a/src/metpy/calc/basic.py b/src/metpy/calc/basic.py index 7c53d25be86..9093e26a511 100644 --- a/src/metpy/calc/basic.py +++ b/src/metpy/calc/basic.py @@ -101,7 +101,7 @@ def wind_direction(u, v, convention='from'): mask = wdir <= 0 if np.any(mask): - wdir[mask] += units.Quantity(360., 'deg') + wdir[np.where(mask)] += units.Quantity(360., 'deg') # avoid unintended modification of `pint.Quantity` by direct use of magnitude calm_mask = (np.asanyarray(u.magnitude) == 0.) & (np.asanyarray(v.magnitude) == 0.) @@ -285,12 +285,12 @@ def heat_index(temperature, relative_humidity, mask_undefined=True): hi = masked_array(hi) # If T <= 40F, Heat Index is T - sel = (temperature <= units.Quantity(40., 'degF')) + sel = np.array(temperature <= units.Quantity(40., 'degF')) if np.any(sel): hi[sel] = temperature[sel].to(units.degF) # If a < 79F and hi is unset, Heat Index is a - sel = (a < units.Quantity(79., 'degF')) & np.isnan(hi) + sel = np.array(a < units.Quantity(79., 'degF')) & np.isnan(hi) if np.any(sel): hi[sel] = a[sel] @@ -300,9 +300,9 @@ def heat_index(temperature, relative_humidity, mask_undefined=True): hi[sel] = b[sel] # Adjustment for RH <= 13% and 80F <= T <= 112F - sel = ((relative_humidity <= units.Quantity(13., 'percent')) - & (temperature >= units.Quantity(80., 'degF')) - & (temperature <= units.Quantity(112., 'degF'))) + sel = np.array((relative_humidity <= units.Quantity(13., 'percent')) + & (temperature >= units.Quantity(80., 'degF')) + & (temperature <= units.Quantity(112., 'degF'))) if np.any(sel): rh15adj = ((13. - relative_humidity[sel] * 100.) / 4. * np.sqrt((units.Quantity(17., 'delta_degF') @@ -311,9 +311,9 @@ def heat_index(temperature, relative_humidity, mask_undefined=True): hi[sel] = hi[sel] - rh15adj # Adjustment for RH > 85% and 80F <= T <= 87F - sel = ((relative_humidity > units.Quantity(85., 'percent')) - & (temperature >= units.Quantity(80., 'degF')) - & (temperature <= units.Quantity(87., 'degF'))) + sel = np.array((relative_humidity > units.Quantity(85., 'percent')) + & (temperature >= units.Quantity(80., 'degF')) + & (temperature <= units.Quantity(87., 'degF'))) if np.any(sel): rh85adj = (0.02 * (relative_humidity[sel] * 100. - 85.) * (units.Quantity(87., 'delta_degF') - delta[sel])) @@ -324,6 +324,7 @@ def heat_index(temperature, relative_humidity, mask_undefined=True): mask = np.array(temperature < units.Quantity(80., 'degF')) if mask.any(): hi = masked_array(hi, mask=mask) + return hi @@ -374,7 +375,7 @@ def apparent_temperature(temperature, relative_humidity, speed, face_level_winds heat_index, windchill """ - is_not_scalar = isinstance(temperature.m, (list, tuple, np.ndarray)) + is_not_scalar = hasattr(temperature, '__len__') temperature = np.atleast_1d(temperature) relative_humidity = np.atleast_1d(relative_humidity) diff --git a/tests/calc/test_basic.py b/tests/calc/test_basic.py index e6ae8878e9b..1fee7db1cd9 100644 --- a/tests/calc/test_basic.py +++ b/tests/calc/test_basic.py @@ -73,7 +73,6 @@ def test_speed(array_type): assert_array_almost_equal(true_speed, speed, 4) -@pytest.mark.xfail_dask('Item assignment with not supported') def test_direction(array_type): """Test calculating wind direction.""" # The last two (u, v) pairs and their masks test masking calm and negative directions @@ -88,7 +87,6 @@ def test_direction(array_type): assert_array_almost_equal(true_dir, direc, 4) -@pytest.mark.xfail_dask('Boolean index assignment in Dask expects equally shaped arrays') def test_direction_with_north_and_calm(array_type): """Test how wind direction handles northerly and calm winds.""" mask = [False, False, False, True] @@ -108,7 +106,6 @@ def test_direction_dimensions(): assert str(d.units) == 'degree' -@pytest.mark.xfail_dask('Boolean index assignment in Dask expects equally shaped arrays') def test_oceanographic_direction(array_type): """Test oceanographic direction (to) convention.""" mask = [False, True, False] @@ -117,7 +114,7 @@ def test_oceanographic_direction(array_type): direc = wind_direction(u, v, convention='to') true_dir = array_type([135., 90., 360.], 'deg', mask=mask) - assert_almost_equal(direc, true_dir, 4) + assert_array_almost_equal(direc, true_dir, 4) def test_invalid_direction_convention(): @@ -207,7 +204,6 @@ def test_windchill_face_level(): assert_array_almost_equal(wc, values, 0) -@pytest.mark.xfail_dask('operands could not be broadcast together with shapes (0, 5) (nan,)') def test_heat_index_basic(array_type): """Test the basic heat index calculation.""" mask = [False, True, False, True, False] @@ -225,14 +221,17 @@ def test_heat_index_scalar(): assert_almost_equal(hi, 121 * units.degF, 0) -def test_heat_index_invalid(): +def test_heat_index_invalid(array_type): """Test heat index for values that should be masked.""" - temp = np.array([80, 88, 92, 79, 30, 81]) * units.degF - rh = np.array([40, 39, 2, 70, 50, 39]) * units.percent + mask = [False, False, False, False, False, False] + temp = array_type([80, 88, 92, 79, 30, 81], 'degF', mask=mask) + rh = array_type([40, 39, 2, 70, 50, 39], 'percent', mask=mask) hi = heat_index(temp, rh) - mask = np.array([False, False, False, True, True, False]) - assert_array_equal(hi.mask, mask) + if isinstance(hi, xr.DataArray): + hi = hi.data + true_mask = np.array([False, False, False, True, True, False]) + assert_array_equal(hi.mask, true_mask) def test_heat_index_undefined_flag(): @@ -408,10 +407,6 @@ def test_coriolis_units(): assert f.units == units('1/second') -@pytest.mark.xfail_dask( - 'boolean index did not match indexed array along dimension 0; dimension is 2 but ' - 'corresponding boolean dimension is 3' -) def test_apparent_temperature(array_type): """Test the apparent temperature calculation.""" temperature = array_type([[90, 90, 70],