diff --git a/.github/workflows/main.yml b/.github/workflows/main.yml index 27d2bafda..2a6ca77b1 100644 --- a/.github/workflows/main.yml +++ b/.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') }} diff --git a/CHANGES.rst b/CHANGES.rst index 511435c6e..c0b81988f 100644 --- a/CHANGES.rst +++ b/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) ---------------- diff --git a/spectral_cube/dask_spectral_cube.py b/spectral_cube/dask_spectral_cube.py index e4121f619..57fe59f0e 100644 --- a/spectral_cube/dask_spectral_cube.py +++ b/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 diff --git a/spectral_cube/stokes_spectral_cube.py b/spectral_cube/stokes_spectral_cube.py index ef5179e51..b26f92740 100644 --- a/spectral_cube/stokes_spectral_cube.py +++ b/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): """ diff --git a/spectral_cube/tests/test_stokes_spectral_cube.py b/spectral_cube/tests/test_stokes_spectral_cube.py index 28c11b10e..70adde4f7 100644 --- a/spectral_cube/tests/test_stokes_spectral_cube.py +++ b/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'])