Suppress warnings (#197)

* Suppress warnings

Closes #188

* Silence a bunch of dask warnings

* Fix mypy

Author: dcherian

flox/aggregations.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import copy
+import warnings
 from functools import partial
 
 import numpy as np
@@ -48,9 +49,12 @@ def generic_aggregate(
 
     group_idx = np.asarray(group_idx, like=array)
 
-    return method(
-        group_idx, array, axis=axis, size=size, fill_value=fill_value, dtype=dtype, **kwargs
-    )
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered")
+        result = method(
+            group_idx, array, axis=axis, size=size, fill_value=fill_value, dtype=dtype, **kwargs
+        )
+    return result
 
 
 def _normalize_dtype(dtype, array_dtype, fill_value=None):
@@ -243,11 +247,18 @@ def __repr__(self):
     fill_value=1,
     final_fill_value=dtypes.NA,
 )
+
+
+def _mean_finalize(sum_, count):
+    with np.errstate(invalid="ignore", divide="ignore"):
+        return sum_ / count
+
+
 mean = Aggregation(
     "mean",
     chunk=("sum", "nanlen"),
     combine=("sum", "sum"),
-    finalize=lambda sum_, count: sum_ / count,
+    finalize=_mean_finalize,
     fill_value=(0, 0),
     dtypes=(None, np.intp),
     final_dtype=np.floating,
@@ -256,7 +267,7 @@ def __repr__(self):
     "nanmean",
     chunk=("nansum", "nanlen"),
     combine=("sum", "sum"),
-    finalize=lambda sum_, count: sum_ / count,
+    finalize=_mean_finalize,
     fill_value=(0, 0),
     dtypes=(None, np.intp),
     final_dtype=np.floating,
@@ -265,7 +276,8 @@ def __repr__(self):
 
 # TODO: fix this for complex numbers
 def _var_finalize(sumsq, sum_, count, ddof=0):
-    result = (sumsq - (sum_**2 / count)) / (count - ddof)
+    with np.errstate(invalid="ignore", divide="ignore"):
+        result = (sumsq - (sum_**2 / count)) / (count - ddof)
     result[count <= ddof] = np.nan
     return result
 
@@ -352,6 +364,10 @@ def _zip_index(array_, idx_):
     )
 
 
+def _pick_second(*x):
+    return x[1]
+
+
 argmax = Aggregation(
     "argmax",
     preprocess=argreduce_preprocess,
@@ -360,7 +376,7 @@ def _zip_index(array_, idx_):
     reduction_type="argreduce",
     fill_value=(dtypes.NINF, 0),
     final_fill_value=-1,
-    finalize=lambda *x: x[1],
+    finalize=_pick_second,
     dtypes=(None, np.intp),
     final_dtype=np.intp,
 )
@@ -373,7 +389,7 @@ def _zip_index(array_, idx_):
     reduction_type="argreduce",
     fill_value=(dtypes.INF, 0),
     final_fill_value=-1,
-    finalize=lambda *x: x[1],
+    finalize=_pick_second,
     dtypes=(None, np.intp),
     final_dtype=np.intp,
 )
@@ -386,7 +402,7 @@ def _zip_index(array_, idx_):
     reduction_type="argreduce",
     fill_value=(dtypes.NINF, -1),
     final_fill_value=-1,
-    finalize=lambda *x: x[1],
+    finalize=_pick_second,
     dtypes=(None, np.intp),
     final_dtype=np.intp,
 )
@@ -399,7 +415,7 @@ def _zip_index(array_, idx_):
     reduction_type="argreduce",
     fill_value=(dtypes.INF, -1),
     final_fill_value=-1,
-    finalize=lambda *x: x[1],
+    finalize=_pick_second,
     dtypes=(None, np.intp),
     final_dtype=np.intp,
 )

flox/core.py

@@ -4,6 +4,7 @@
 import itertools
 import math
 import operator
+import warnings
 from collections import namedtuple
 from functools import partial, reduce
 from numbers import Integral
@@ -881,7 +882,9 @@ def _simple_combine(
     for idx, combine in enumerate(agg.combine):
         array = _conc2(x_chunk, key1="intermediates", key2=idx, axis=axis_)
         assert array.ndim >= 2
-        result = getattr(np, combine)(array, axis=axis_, keepdims=True)
+        with warnings.catch_warnings():
+            warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered")
+            result = getattr(np, combine)(array, axis=axis_, keepdims=True)
         if is_aggregate:
             # squeeze out DUMMY_AXIS if this is the last step i.e. called from _aggregate
             result = result.squeeze(axis=DUMMY_AXIS)

flox/xarray.py

@@ -223,7 +223,7 @@ def xarray_reduce(
         raise NotImplementedError("sort must be True for xarray_reduce")
 
     # eventually drop the variables we are grouping by
-    maybe_drop = [b for b in by if isinstance(b, Hashable)]
+    maybe_drop = set(b for b in by if isinstance(b, Hashable))
     unindexed_dims = tuple(
         b
         for b, isbin_ in zip(by, isbins)
@@ -243,6 +243,20 @@ def xarray_reduce(
     else:
         ds = obj._to_temp_dataset()
 
+    try:
+        from xarray.indexes import PandasMultiIndex
+    except ImportError:
+        PandasMultiIndex = tuple()  # type: ignore
+
+    more_drop = set()
+    for var in maybe_drop:
+        maybe_midx = ds._indexes.get(var, None)
+        if isinstance(maybe_midx, PandasMultiIndex):
+            idx_coord_names = set(maybe_midx.index.names + [maybe_midx.dim])
+            idx_other_names = idx_coord_names - set(maybe_drop)
+            more_drop.update(idx_other_names)
+    maybe_drop.update(more_drop)
+
     ds = ds.drop_vars([var for var in maybe_drop if var in ds.variables])
 
     if dim is Ellipsis:

tests/test_core.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 import itertools
+import warnings
 from functools import partial, reduce
 from typing import TYPE_CHECKING
 
@@ -204,11 +205,18 @@ def test_groupby_reduce_all(nby, size, chunks, func, add_nan_by, engine):
     for kwargs in finalize_kwargs:
         flox_kwargs = dict(func=func, engine=engine, finalize_kwargs=kwargs, fill_value=fill_value)
         with np.errstate(invalid="ignore", divide="ignore"):
-            if "arg" in func and add_nan_by:
-                array[..., nanmask] = np.nan
-                expected = getattr(np, "nan" + func)(array, axis=-1, **kwargs)
-            else:
-                expected = getattr(np, func)(array[..., ~nanmask], axis=-1, **kwargs)
+            with warnings.catch_warnings():
+                warnings.filterwarnings("ignore", r"All-NaN (slice|axis) encountered")
+                warnings.filterwarnings("ignore", r"Degrees of freedom <= 0 for slice")
+                warnings.filterwarnings("ignore", r"Mean of empty slice")
+
+                # computing silences a bunch of dask warnings
+                array_ = array.compute() if chunks is not None else array
+                if "arg" in func and add_nan_by:
+                    array_[..., nanmask] = np.nan
+                    expected = getattr(np, "nan" + func)(array_, axis=-1, **kwargs)
+                else:
+                    expected = getattr(np, func)(array_[..., ~nanmask], axis=-1, **kwargs)
         for _ in range(nby):
             expected = np.expand_dims(expected, -1)
 

tests/test_xarray.py

@@ -286,7 +286,8 @@ def test_xarray_resample_dataset_multiple_arrays(engine):
     # The separate computes are necessary here to force xarray
     # to compute all variables in result at the same time.
     expected = resampler.mean().compute()
-    result = resample_reduce(resampler, "mean", engine=engine).compute()
+    with pytest.warns(DeprecationWarning):
+        result = resample_reduce(resampler, "mean", engine=engine).compute()
     xr.testing.assert_allclose(expected, result)
 
 
@@ -450,7 +451,8 @@ def test_datetime_array_reduce(use_cftime, func, engine):
         name="time",
     )
     expected = getattr(time.resample(time="YS"), func)()
-    actual = resample_reduce(time.resample(time="YS"), func=func, engine=engine)
+    with pytest.warns(DeprecationWarning):
+        actual = resample_reduce(time.resample(time="YS"), func=func, engine=engine)
     assert_equal(expected, actual)