Added stokes_data public method

.github/workflows/main.yml

@@ -64,7 +64,6 @@ jobs:
     - name: Install testing dependencies
       run: python -m pip install tox codecov
     - name: Run tests with ${{ matrix.name }}
-      if: ${{ contains(matrix.toxenv,'-cov') }}
       run: tox -v -e ${{ matrix.toxenv }}
     - name: Upload coverage to codecov
       if: ${{ contains(matrix.toxenv,'-cov') }}

CHANGES.rst

@@ -1,6 +1,6 @@
 0.6.1.dev (unreleased)
 ----------------------
-- none yet
+- Fix memory issue when calling statistics() on FITS cubes. #752
 
 0.6 (2021-09-30)
 ----------------

spectral_cube/dask_spectral_cube.py

@@ -760,20 +760,18 @@ def statistics(self):
 
         data = self._get_filled_data(fill=np.nan)
 
-        try:
-            from bottleneck import nanmin, nanmax, nansum
-        except ImportError:
-            from numpy import nanmin, nanmax, nansum
-
         def compute_stats(chunk, *args):
-            return np.array([np.sum(~np.isnan(chunk)),
-                             nanmin(chunk),
-                             nanmax(chunk),
-                             nansum(chunk),
-                             nansum(chunk * chunk)])
+            # NOTE: we used to but do not use bottleneck here anymore, as it
+            # does not seem to provide any performance improvements and also
+            # has a memory leak when dealing with >f4 dtypes.
+            return np.array([[[[np.sum(~np.isnan(chunk)),
+                                np.nanmin(chunk),
+                                np.nanmax(chunk),
+                                np.nansum(chunk),
+                                np.nansum(chunk * chunk)]]]])
 
         with dask.config.set(**self._scheduler_kwargs):
-            results = da.map_blocks(compute_stats, data.reshape(-1)).compute()
+            results = da.map_blocks(compute_stats, data, new_axis=3).compute()
 
         count_values, min_values, max_values, sum_values, ssum_values = results.reshape((-1, 5)).T
 

spectral_cube/stokes_spectral_cube.py

@@ -9,7 +9,7 @@
 from . import wcs_utils
 from .masks import BooleanArrayMask, is_broadcastable_and_smaller
 
-__all__ = ['StokesSpectalCube']
+__all__ = ['StokesSpectralCube']
 
 VALID_STOKES = ['I', 'Q', 'U', 'V', 'RR', 'LL', 'RL', 'LR', 'XX', 'XY', 'YX', 'YY', 
                 'RX', 'RY', 'LX', 'LY', 'XR,', 'XL', 'YR', 'YL', 'PP', 'PQ', 'QP', 'QQ', 
@@ -61,10 +61,30 @@ def __init__(self, stokes_data, mask=None, meta=None, fill_value=None):
 
         self._mask = mask
 
+    def __getitem__(self, key):
+        if key in self._stokes_data:
+            return self._stokes_data[key]
+        else:
+            raise KeyError("Key {0} does not exist in this cube.".format(key))
+
+    def __setitem__(self, key, item):
+        if key in self._stokes_data:
+            self._stokes_data[key] = item
+        else:
+            errmsg = "Assigning new Stokes axes is not yet supported."
+            raise NotImplementedError(errmsg)
+
     @property
     def shape(self):
         return self._shape
 
+    @property
+    def stokes_data(self):
+        """
+        The underlying data
+        """
+        return self._stokes_data
+
     @property
     def mask(self):
         """

spectral_cube/tests/test_stokes_spectral_cube.py

@@ -152,3 +152,15 @@ def test_separate_mask(self, use_dask):
 
         cube2 = cube1.I.with_mask(mask3)
         assert_equal(cube2.mask.include(), (mask1 & mask2[0] & mask3).include())
+
+    def test_key_access_valid(self, use_dask):
+        stokes_data = OrderedDict()
+        stokes_data['I'] = SpectralCube(self.data[0], wcs=self.wcs, use_dask=use_dask)
+        stokes_data['Q'] = SpectralCube(self.data[1], wcs=self.wcs, use_dask=use_dask)
+        stokes_data['U'] = SpectralCube(self.data[2], wcs=self.wcs, use_dask=use_dask)
+        stokes_data['V'] = SpectralCube(self.data[3], wcs=self.wcs, use_dask=use_dask)
+        cube = StokesSpectralCube(stokes_data)
+        assert_equal(cube['I'],cube._stokes_data['I'])
+        assert_equal(cube['Q'],cube._stokes_data['Q'])
+        assert_equal(cube['U'],cube._stokes_data['U'])
+        assert_equal(cube['V'],cube._stokes_data['V'])