Fix statistics() when data is chunked

radio-astro-tools · Sep 22, 2020 · 59265aa · 59265aa
1 parent 370525c
commit 59265aa
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/spectral_cube/dask_spectral_cube.py b/spectral_cube/dask_spectral_cube.py
@@ -735,7 +735,7 @@ def compute_stats(chunk, *args):
                              nansum(chunk * chunk)])
 
         with dask.config.set(**self._scheduler_kwargs):
-            results = da.map_blocks(compute_stats, data).compute()
+            results = da.map_blocks(compute_stats, data.reshape(-1)).compute()
 
         count_values, min_values, max_values, sum_values, ssum_values = results.reshape((-1, 5)).T
 

diff --git a/spectral_cube/tests/test_dask.py b/spectral_cube/tests/test_dask.py
@@ -77,7 +77,7 @@ def test_rechunk(data_adv):
 
 
 def test_statistics(data_adv):
-    cube = DaskSpectralCube.read(data_adv)
+    cube = DaskSpectralCube.read(data_adv).rechunk(chunks=(1, 2, 3))
     stats = cube.statistics()
     assert_quantity_allclose(stats['npts'], 24)
     assert_quantity_allclose(stats['mean'], 0.4941651776136591 * u.K)