Add a fixture that provides numpy, masked, and dask array functions

ci/test_requirements.txt

@@ -2,3 +2,4 @@ pytest==6.2.4
 pytest-mpl==0.13
 netCDF4==1.5.7
 coverage==5.5
+dask==2021.2.0

conftest.py

@@ -162,3 +162,20 @@ def set_agg_backend():
         yield
     finally:
         matplotlib.pyplot.switch_backend(prev_backend)
+
+
+@pytest.fixture(params=['dask', 'xarray', 'masked', 'numpy'])
+def array_type(request):
+    """Return an array type for testing calc functions."""
+    quantity = metpy.units.units.Quantity
+    if request.param == 'dask':
+        dask_array = pytest.importorskip('dask.array', reason='dask.array is not available')
+        return lambda d, u, *, mask=None: quantity(dask_array.array(d), u)
+    elif request.param == 'xarray':
+        return lambda d, u, *, mask=None: xarray.DataArray(d, attrs={'units': u})
+    elif request.param == 'masked':
+        return lambda d, u, *, mask=None: quantity(numpy.ma.array(d, mask=mask), u)
+    elif request.param == 'numpy':
+        return lambda d, u, *, mask=None: quantity(numpy.array(d), u)
+    else:
+        raise ValueError(f'Unsupported array_type option {request.param}')

src/metpy/calc/basic.py

@@ -104,8 +104,14 @@ def wind_direction(u, v, convention='from'):
         wdir[mask] += units.Quantity(360., 'deg')
     # avoid unintended modification of `pint.Quantity` by direct use of magnitude
     calm_mask = (np.asarray(u.magnitude) == 0.) & (np.asarray(v.magnitude) == 0.)
+
+    if hasattr(wdir, 'mask'):
+        array_mask = wdir.mask
+    else:
+        array_mask = False
+
     # np.any check required for legacy numpy which treats 0-d False boolean index as zero
-    if np.any(calm_mask):
+    if np.any(calm_mask & np.logical_not(array_mask)):
         wdir[calm_mask] = units.Quantity(0., 'deg')
     return wdir.reshape(origshape).to('degrees')
 
@@ -798,8 +804,12 @@ def smooth_gaussian(scalar_grid, n):
     # Assume the last two axes represent the horizontal directions
     sgma_seq = [sgma if i > num_ax - 3 else 0 for i in range(num_ax)]
 
-    # Compute smoothed field
-    return gaussian_filter(scalar_grid, sgma_seq, truncate=2 * np.sqrt(2))
+    filter_args = {'sigma': sgma_seq, 'truncate': 2 * np.sqrt(2)}
+    if hasattr(scalar_grid, 'mask'):
+        smoothed = gaussian_filter(scalar_grid.data, **filter_args)
+        return np.ma.array(smoothed, mask=scalar_grid.mask)
+    else:
+        return gaussian_filter(scalar_grid, **filter_args)
 
 
 @exporter.export

src/metpy/testing.py

@@ -215,4 +215,12 @@ def wrapper(*args, **kwargs):
     return dec
 
 
+def xfail_dask(array, reason):
+    """Xfail a test if `array` is a `dask.array.Array`."""
+    dask_array = pytest.importorskip('dask.array', reason='dask.array not available')
+
+    if not isinstance(array, xr.DataArray) and isinstance(array.m, dask_array.Array):
+        pytest.xfail(reason=reason)
+
+
 check_and_silence_deprecation = check_and_silence_warning(MetpyDeprecationWarning)