Add a data setter

changelog/795.feature.rst

@@ -0,0 +1 @@
+It is now possible to set the ``NDCube.data`` property of a cube with an array of the same shape and unit as the current cube.

ndcube/conftest.py

@@ -575,7 +575,6 @@ def ndc(request):
 # NDCubeSequence Fixtures
 ################################################################################
 
-
 @pytest.fixture
 def ndcubesequence_4c_ln_lt_l(ndcube_3d_ln_lt_l):
     cube1 = ndcube_3d_ln_lt_l

ndcube/ndcube.py

@@ -399,6 +399,42 @@ def __init__(self, data, wcs=None, uncertainty=None, mask=None, meta=None,
                 global_coords = deepcopy(global_coords)
             self._global_coords = global_coords
 
+    @property
+    def data(self):
+        """
+        `~numpy.ndarray`-like : The stored dataset.
+
+        Notes
+        -----
+        It is possible to set the ``.data`` attribute on a `NDCube` with an
+        array-like object of the same shape. However, this is really only
+        intended for replacing the data with a different object representing
+        the same physical data as no other properties of the cube will be
+        changed, such as uncertainty or unit.
+        """
+        return super().data
+
+    @data.setter
+    def data(self, value):
+        # In an array-agnostic way check the shape is the same
+        if not hasattr(value, "shape") or value.shape != self.data.shape:
+            raise TypeError(f"Can only set data with an array-like object of the same shape ({self.data.shape})")
+
+        # Other masked arrays are hard to detect reliably
+        if isinstance(value, np.ma.MaskedArray):
+            raise TypeError("Can not set the .data attribute with a numpy masked array, please set .data and .mask separately.")
+
+        if isinstance(value, u.Quantity):
+            unit_error = f"Unable to set data with unit {value.unit} as it incompatible with the current unit of {self.unit}"
+            if self.unit is None:
+                raise u.UnitsError(unit_error)
+            try:
+                value = value.to_value(self.unit)
+            except u.UnitsError as exc:
+                raise u.UnitsError(unit_error) from exc
+
+        self._data = value
+
     @property
     def extra_coords(self):
         # Docstring in NDCubeABC.

ndcube/tests/test_ndcube.py

@@ -1238,3 +1238,64 @@ def test_ndcube_quantity(ndcube_2d_ln_lt_units):
     cube = ndcube_2d_ln_lt_units
     expected = u.Quantity(cube.data, cube.unit)
     np.testing.assert_array_equal(cube.quantity, expected)
+
+
+def test_data_setter(ndcube_4d_ln_l_t_lt):
+    cube = ndcube_4d_ln_l_t_lt
+    assert isinstance(cube.data, np.ndarray)
+
+    new_data = np.zeros_like(cube.data)
+    cube.data = new_data
+    assert cube.data is new_data
+
+    dask_array = dask.array.zeros_like(cube.data)
+    cube.data = dask_array
+    assert cube.data is dask_array
+
+
+def test_invalid_data_setter(ndcube_4d_ln_l_t_lt):
+    cube = ndcube_4d_ln_l_t_lt
+
+    with pytest.raises(TypeError, match="set data with an array-like"):
+        cube.data = None
+
+    with pytest.raises(TypeError, match="set data with an array-like"):
+        cube.data = np.zeros((100,100))
+
+    with pytest.raises(TypeError, match="set data with an array-like"):
+        cube.data = 10
+
+
+def test_quantity_data_setter(ndcube_2d_ln_lt_units):
+    cube = ndcube_2d_ln_lt_units
+    assert cube.unit
+
+    new_data = np.zeros_like(cube.data) * cube.unit
+    cube.data = new_data
+
+    assert isinstance(cube.data, np.ndarray)
+
+    new_data = np.zeros_like(cube.data) * u.Jy
+    with pytest.raises(u.UnitsError, match=f"Unable to set data with unit {u.Jy}"):
+        cube.data = new_data
+
+
+def test_quantity_no_unit_data_setter(ndcube_4d_ln_l_t_lt):
+    cube = ndcube_4d_ln_l_t_lt
+
+    new_data = np.zeros_like(cube.data) * u.Jy
+    with pytest.raises(u.UnitsError, match=f"Unable to set data with unit {u.Jy}.* current unit of None"):
+        cube.data = new_data
+
+
+def test_set_data_mask(ndcube_4d_mask):
+    cube = ndcube_4d_mask
+
+    assert isinstance(cube.mask, np.ndarray)
+
+    new_data = np.ones_like(cube.data)
+    new_mask = np.zeros_like(cube.mask)
+    masked_array = np.ma.MaskedArray(new_data, new_mask)
+
+    with pytest.raises(TypeError, match="Can not set the .data .* with a numpy masked array"):
+        cube.data = masked_array