Merge pull request #751 from Kitchi/stokes_methods

Added stokes_data public method
radio-astro-tools · Nov 3, 2021 · 27896c6 · 27896c6
2 parents 8d1a844 + 7ae5310
commit 27896c6
Show file tree

Hide file tree

Showing 2 changed files with 33 additions and 1 deletion.
diff --git 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
@@ -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'])