Merge pull request #922 from keflavich/fix_dask_statistics

Fix dask statistics

spectral_cube/dask_spectral_cube.py

@@ -781,11 +781,12 @@ def compute_stats(chunk, *args):
 
         count_values, min_values, max_values, sum_values, ssum_values = results.reshape((-1, 5)).T
 
+        # all-NAN chunks are possible, so we need to use nan<stat> here
         stats = {'npts': count_values.sum(),
-                 'min': min_values.min() * self._unit,
-                 'max': max_values.max() * self._unit,
-                 'sum': sum_values.sum() * self._unit,
-                 'sumsq': ssum_values.sum() * self._unit ** 2}
+                 'min': np.nanmin(min_values) * self._unit,
+                 'max': np.nanmax(max_values) * self._unit,
+                 'sum': np.nansum(sum_values) * self._unit,
+                 'sumsq': np.nansum(ssum_values) * self._unit ** 2}
 
         stats['mean'] = stats['sum'] / stats['npts']
 

spectral_cube/tests/test_dask.py

@@ -104,6 +104,17 @@ def test_statistics(data_adv):
     assert_quantity_allclose(stats['rms'], 0.5759458158839716 * u.K)
 
 
+def test_statistics_withnans(data_adv):
+    cube = DaskSpectralCube.read(data_adv).rechunk(chunks=(1, 2, 3))
+    # shape is 2, 3, 4
+    cube._data[:,:,:2] = np.nan
+    # ensure some chunks are all nan
+    cube.rechunk((1,2,2))
+    stats = cube.statistics()
+    for key in ('min', 'max', 'sum'):
+        assert stats[key] == getattr(cube, key)()
+
+
 @pytest.mark.skipif(not CASA_INSTALLED, reason='Requires CASA to be installed')
 def test_statistics_consistency_casa(data_adv, tmp_path):