WCS information lost in unit conversion of LowerDimensionalObject

spectral_cube/lower_dimensional_structures.py

@@ -19,6 +19,10 @@ def wcs(self):
     def meta(self):
         return self._meta
 
+    @property
+    def mask(self):
+        return self._mask
+
     @property
     def header(self):
         header = self._header
@@ -68,6 +72,52 @@ def write(self, filename, format=None, overwrite=False):
             raise ValueError("Unknown format '{0}' - the only available "
                              "format at this time is 'fits'")
 
+    def to(self, unit, equivalencies=[]):
+        """
+        Return a new ``LowerDimensionalObject'' of the same class with the
+        specified unit.  See `astropy.units.Quantity.to` for further details.
+        """
+        converted_array = u.Quantity.to(self, unit,
+                                        equivalencies=equivalencies).value
+
+        # use private versions of variables, not the generated property
+        # versions
+        # Not entirely sure the use of __class__ here is kosher, but we do want
+        # self.__class__, not super()
+        new = self.__class__(value=converted_array, unit=unit, copy=True,
+                             wcs=self._wcs, meta=self._meta, mask=self._mask,
+                             header=self._header)
+
+        return new
+
+    def __getitem__(self, key):
+        """
+        Return a new ``LowerDimensionalObject'' of the same class while keeping
+        other properties fixed.
+        """
+        new_qty = super(LowerDimensionalObject, self).__getitem__(key)
+
+        if new_qty.ndim < 2:
+            # do not return a projection
+            return u.Quantity(new_qty)
+
+        if self._wcs is not None:
+            newwcs = self._wcs[key]
+        else:
+            newwcs = None
+
+        new = self.__class__(value=new_qty.value,
+                             unit=new_qty.unit,
+                             copy=False,
+                             wcs=newwcs,
+                             meta=self._meta,
+                             mask=self._mask,
+                             header=self._header)
+
+        return new
+
+
+
 class Projection(LowerDimensionalObject):
 
     def __new__(cls, value, unit=None, dtype=None, copy=True, wcs=None,

spectral_cube/tests/test_projection.py

@@ -4,15 +4,34 @@
 from .helpers import assert_allclose
 from ..lower_dimensional_structures import Projection
 
+def test_slices_of_projections_not_projections():
+    # slices of projections that have <2 dimensions should not be projections
+    image = np.ones((2, 2)) * u.Jy
+    p = Projection(image, copy=False)
+
+    assert not isinstance(p[0,0], Projection)
+    assert not isinstance(p[0], Projection)
+
 
 def test_copy_false():
     image = np.ones((12, 12)) * u.Jy
     p = Projection(image, copy=False)
     image[3,4] = 2 * u.Jy
     assert_allclose(p[3,4], 2 * u.Jy)
 
-
 def test_write(tmpdir):
     image = np.ones((12, 12)) * u.Jy
     p = Projection(image)
     p.write(tmpdir.join('test.fits').strpath)
+
+def test_preserve_wcs_to():
+    # regression for #256
+    image = np.ones((12, 12)) * u.Jy
+    p = Projection(image, copy=False)
+    image[3,4] = 2 * u.Jy
+
+    p2 = p.to(u.mJy)
+    assert_allclose(p[3,4], 2 * u.Jy)
+    assert_allclose(p[3,4], 2000 * u.mJy)
+
+    assert p2.wcs == p.wcs