Convert instances of boolean indexing to ndarrays to pass dask tests

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,7 +300,7 @@ 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'))
+    sel = np.array((relative_humidity <= units.Quantity(13., 'percent'))
            & (temperature >= units.Quantity(80., 'degF'))
            & (temperature <= units.Quantity(112., 'degF')))
     if np.any(sel):
@@ -311,7 +311,7 @@ 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'))
+    sel = np.array((relative_humidity > units.Quantity(85., 'percent'))
            & (temperature >= units.Quantity(80., 'degF'))
            & (temperature <= units.Quantity(87., 'degF')))
     if np.any(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)

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 <class "numpy.ndarray"> 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],