Merge 19e3689 into 443ae46

spectral_cube/base_class.py

@@ -588,11 +588,14 @@ def average_beams(self, threshold, mask='compute', warn=False):
                 beam_mask = da.any(da.logical_and(self._mask_include,
                                                   self.goodbeams_mask[:, None, None]),
                                    axis=(1, 2))
-                beam_mask = self._compute(beam_mask)
-            else:
+                # da.any appears to return an object dtype instead of a bool
+                beam_mask = self._compute(beam_mask).astype('bool')
+            elif self.mask is not None:
                 beam_mask = np.any(np.logical_and(self.mask.include(),
                                                   self.goodbeams_mask[:, None, None]),
                                    axis=(1, 2))
+            else:
+                beam_mask = self.goodbeams_mask
         else:
             if mask.ndim > 1:
                 beam_mask = np.logical_and(mask, self.goodbeams_mask[:, None, None])
@@ -725,11 +728,20 @@ def mask_out_bad_beams(self, threshold, reference_beam=None,
                                             criteria=criteria,
                                             mid_value=mid_value)
 
-        includemask = BooleanArrayMask(goodbeams[:,None,None],
+        includemask = BooleanArrayMask(goodbeams[:, None, None],
                                        self._wcs,
                                        shape=self._data.shape)
 
-        return self._new_thing_with(mask=np.bitwise_and(self.mask, includemask),
+        use_dask = isinstance(self._data, da.Array)
+        if use_dask:
+            newmask = da.logical_and(self._mask_include,
+                                     includemask)
+        elif self.mask is None:
+            newmask = includemask
+        else:
+            newmask = np.bitwise_and(self.mask, includemask)
+
+        return self._new_thing_with(mask=newmask,
                                     beam_threshold=threshold,
                                     goodbeams_mask=np.bitwise_and(self.goodbeams_mask, goodbeams),
                                    )

spectral_cube/conftest.py

@@ -282,6 +282,32 @@ def data_vda_beams(tmp_path):
     return tmp_path / 'vda_beams.fits'
 
 
+@pytest.fixture
+def data_vda_beams_image(tmp_path):
+    d, h = prepare_adv_data()
+    d, h = transpose(d, h, [2, 0, 1])
+    d, h = transpose(d, h, [2, 1, 0])
+    h['BUNIT'] = ' Jy / beam '
+    del h['BMAJ'], h['BMIN'], h['BPA']
+    beams = prepare_4_beams()
+    hdul = fits.HDUList([fits.PrimaryHDU(data=d, header=h),
+                         beams])
+    hdul.writeto(tmp_path / 'vda_beams.fits')
+    from casatools import image
+    ia = image()
+    ia.fromfits(infile=tmp_path / 'vda_beams.fits',
+                outfile=tmp_path / 'vda_beams.image',
+                overwrite=True)
+    for (bmaj, bmin, bpa, chan, pol) in beams.data:
+        ia.setrestoringbeam(major={'unit': 'arcsec', 'value': bmaj},
+                            minor={'unit': 'arcsec', 'value': bmin},
+                            pa={'unit': 'deg', 'value': bpa},
+                            channel=chan,
+                            polarization=pol)
+    ia.close()
+    return tmp_path / 'vda_beams.image'
+
+
 def prepare_255_header():
     # make a version with spatial pixels
     h = fits.header.Header.fromtextfile(HEADER_FILENAME)

spectral_cube/dask_spectral_cube.py

@@ -1511,4 +1511,8 @@ def spectral_smooth(self, *args, **kwargs):
 
     @property
     def _mask_include(self):
-        return da.from_array(MaskHandler(self), name='MaskHandler ' + str(uuid.uuid4()), chunks=self._data.chunksize)
+        return BooleanArrayMask(da.from_array(MaskHandler(self),
+                                              name='MaskHandler ' + str(uuid.uuid4()),
+                                              chunks=self._data.chunksize),
+                                wcs=self.wcs,
+                                shape=self.shape)

spectral_cube/io/fits.py

@@ -7,7 +7,6 @@
 from astropy.io import fits
 from astropy.io import registry as io_registry
 import astropy.wcs
-from astropy import wcs
 from astropy.wcs import WCS
 from collections import OrderedDict
 from astropy.io.fits.hdu.hdulist import fitsopen as fits_open
@@ -183,9 +182,8 @@ def load_fits_cube(input, hdu=0, meta=None, target_cls=None, use_dask=False, **k
         if beam_table is None:
             cube = SC(data, wcs, mask, meta=meta, header=header)
         else:
-            cube = VRSC(data, wcs, mask, meta=meta,
-                                                 header=header,
-                                                 beam_table=beam_table)
+            cube = VRSC(data, wcs, mask, meta=meta, header=header,
+                        beam_table=beam_table)
 
         if hasattr(cube._mask, '_data'):
             # check that the shape matches if there is a shape

spectral_cube/masks.py

@@ -256,6 +256,18 @@ def shape(self):
         raise NotImplementedError("{0} mask classes do not have shape attributes."
                                   .format(self.__class__.__name__))
 
+    @property
+    def ndim(self):
+        return len(self.shape)
+
+    @property
+    def size(self):
+        return np.product(self.shape)
+
+    @property
+    def dtype(self):
+        return np.dtype('bool')
+
     def __getitem__(self):
         raise NotImplementedError("Slicing not supported by mask class {0}"
                                   .format(self.__class__.__name__))

spectral_cube/tests/test_spectral_cube.py

@@ -1880,13 +1880,22 @@ def test_varyres_moment_logic_issue364(data_vda_beams, use_dask):
     assert_quantity_allclose(m0.meta['beam'].major, 0.35*u.arcsec)
 
 
-def test_mask_bad_beams(data_vda_beams, use_dask):
+@pytest.mark.skipif('not casaOK')
+@pytest.mark.parametrize('filename', ['data_vda_beams',
+                                      'data_vda_beams_image'],
+                         indirect=['filename'])
+def test_mask_bad_beams(filename, use_dask):
     """
     Prior to #543, this tested two different scenarios of beam masking.  After
     that, the tests got mucked up because we can no longer have minor>major in
     the beams.
     """
-    cube, data = cube_and_raw(data_vda_beams, use_dask=use_dask)
+    if 'image' in str(filename) and not use_dask:
+        pytest.skip()
+
+    cube, data = cube_and_raw(filename, use_dask=use_dask)
+
+    assert isinstance(cube, base_class.MultiBeamMixinClass)
 
     # make sure all of the beams are initially good (finite)
     assert np.all(cube.goodbeams_mask)
@@ -2378,7 +2387,7 @@ def test_mask_none(use_dask):
 
 
 @pytest.mark.parametrize('filename', ['data_vda', 'data_vda_beams'],
-                            indirect=['filename'])
+                         indirect=['filename'])
 def test_mask_channels_preserve_mask(filename, use_dask):
 
     # Regression test for a bug that caused the mask to not be preserved.