Added new features to the ndcube.__add__ method

changelog/794.feature.rst

@@ -0,0 +1 @@
+Allows addition of an ``NDCube`` and ``NDData`` (with the WCS of ``NDData`` being set to None), and combines their uncertainties and masks.

ndcube/conftest.py

@@ -678,6 +678,39 @@ def ndcube_2d_ln_lt_units(wcs_2d_lt_ln):
     return NDCube(data_cube, wcs=wcs_2d_lt_ln, unit=u.ct)
 
 
+@pytest.fixture
+def ndcube_2d_ln_lt_no_unit_no_unc(wcs_2d_lt_ln):
+    shape = (10, 12)
+    data_cube = data_nd(shape).astype(float)
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln)
+
+
+@pytest.fixture
+def ndcube_2d_unit_unc(wcs_2d_lt_ln):
+    shape = (10, 12)
+    data_cube = data_nd(shape).astype(float)
+    uncertainty = StdDevUncertainty(np.ones(shape)*0.2, unit=u.ct)
+
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty, unit=u.ct)
+
+
+@pytest.fixture
+def ndcube_2d_uncertainty_no_unit(wcs_2d_lt_ln):
+    shape = (10, 12)
+    data_cube = data_nd(shape).astype(float)
+    uncertainty = StdDevUncertainty(np.ones(shape)*0.2)
+
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty)
+
+
+@pytest.fixture
+def ndcube_2d_ln_lt_mask(wcs_2d_lt_ln):
+    shape = (10, 12)
+    data_cube = data_nd(shape).astype(float)
+    mask = np.ones(data_cube.shape, dtype=bool)
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, mask=mask)
+
+
 @pytest.fixture
 def ndcube_2d_dask(wcs_2d_lt_ln):
     shape = (8, 4)
@@ -719,6 +752,9 @@ def ndcube_1d_l(wcs_1d_l):
     "ndcube_2d_ln_lt_units",
     "ndcube_2d_dask",
     "ndcube_1d_l",
+    "ndcube_2d_ln_lt_no_unit_no_unc",
+    "ndcube_2d_uncertainty_no_unit",
+    "ndcube_2d_unit_unc",
 ])
 def all_ndcubes(request):
     """

ndcube/ndcube.py

@@ -11,6 +11,7 @@
 
 import astropy.nddata
 import astropy.units as u
+from astropy.nddata import NDData
 from astropy.units import UnitsError
 from astropy.wcs.utils import _split_matrix
 
@@ -964,24 +965,113 @@
     def __neg__(self):
         return self._new_instance(data=-self.data)
 
-    def __add__(self, value):
-        if hasattr(value, 'unit'):
+    def add(self, value, operation_ignores_mask=True, handle_mask=np.logical_and):
+        """
+        Users are allowed to choose whether they want operation_ignores_mask to be True or False,
+        and are allowed to choose whether they want handle_mask to be AND / OR .
+        """
+        kwargs = {}
+
+        if isinstance(value, NDData) and value.wcs is None:
+            if self.unit is not None and value.unit is not None:
+                value_data = (value.data * value.unit).to_value(self.unit)
+            elif self.unit is None and value.unit is None:
+                value_data = value.data
+            else:
+                raise TypeError("Adding objects requires both have a unit or neither has a unit.") # change the test as well.
+
+            # check whether there is a mask.
+            # Neither self nor value has a mask
+            self_unmasked = self.mask is None or self.mask is False or not self.mask.any()
+            value_unmasked = value.mask is None or value.mask is False or not value.mask.any()
+
+            if (self_unmasked and value_unmasked) or operation_ignores_mask is True:
+                # addition
+                kwargs["data"] = self.data + value_data
+
+                # combine the uncertainty;
+                if self.uncertainty is not None and value.uncertainty is not None:
+                    result_data = kwargs["data"]
+                    if self.unit is not None:
+                        result_data *= self.unit
+                    new_uncertainty = self.uncertainty.propagate(
+                        np.add, value, result_data=result_data, correlation=0
+                    )
+                    kwargs["uncertainty"] = new_uncertainty
+                elif self.uncertainty is not None:
+                    new_uncertainty = self.uncertainty
+                    kwargs["uncertainty"] = new_uncertainty
+                elif value.uncertainty is not None:
+                    new_uncertainty = value.uncertainty
+                    kwargs["uncertainty"] = new_uncertainty
+                else:
+                    new_uncertainty = None
+            else:
+                # TODO
+                # When there is a mask, that is when the two new added parameters (OIM and HM) come into the picture.
+                # Conditional statements to permutate the two different scenarios (when it does not ignore the mask).
+                kwargs["data"] = self.data + value_data
+                self_data, value_data = self.data, value.data # May require a copy
+                self_mask, value_mask = self.mask, value.mask # May require handling/converting of cases when masks aren't boolean arrays but are None, True, or False.
+                if not operation_ignores_mask:
+                    no_op_value = 0 # Value to set masked values since we are doing addition. (Would need to be 1 if we were doing multiplication.)
+                    if (self_mask is True and value_mask is False):
+                        idx = np.logical_and(self_mask, np.logical_not(value_mask))
+                        self_data[idx] = no_op_value
+                    elif (self_mask is False and value_mask is True):
+                        idx = np.logical_and(value_mask, np.logical_not(self_mask))
+                        value_data[idx] = no_op_value
+                    elif (self_mask is True and value_mask is True):
+                        idx = np.logical_and(self_mask, value_mask)
+                        self_data[idx] = no_op_value
+                        value_data[idx] = no_op_value
+
+                    #self_data[idx], value_data[idx] = ?, ? # Handle case when both values are masked here. # We are yet to decide the best behaviour here.
+                    # if both are F, no operation of setting the values to be 0 needs to be done.
+                else:
+                    pass # If operation ignores mask, nothing needs to be done. This line not needed in actual code.  Only here for clarity.
+
+                # Perform addition
+                new_data = self_data + value_data
+                # Calculate new mask.
+                new_mask = handle_mask(self_mask, value_mask) if handle_mask else None
+                kwargs["data"] = new_data
+                kwargs["mask"] = new_mask
+
+        elif hasattr(value, 'unit'):
             if isinstance(value, u.Quantity):
                 # NOTE: if the cube does not have units, we cannot
                 # perform arithmetic between a unitful quantity.
                 # This forces a conversion to a dimensionless quantity
                 # so that an error is thrown if value is not dimensionless
                 cube_unit = u.Unit('') if self.unit is None else self.unit
-                new_data = self.data + value.to_value(cube_unit)
+                kwargs["data"] = self.data + value.to_value(cube_unit)
             else:
                 # NOTE: This explicitly excludes other NDCube objects and NDData objects
                 # which could carry a different WCS than the NDCube
                 return NotImplemented
         elif self.unit not in (None, u.Unit("")):
             raise TypeError("Cannot add a unitless object to an NDCube with a unit.")
         else:
-            new_data = self.data + value
-        return self._new_instance(data=new_data)
+            kwargs["data"] = self.data + value
+
+        # return the new NDCube instance
+        return self._new_instance(**kwargs)
+
+    def __add__(self, value):
+        # when value has a mask, raise error and point user to the add method. TODO
+        #
+        # check whether there is a mask.
+        # Neither self nor value has a mask
+
+        self_masked = not(self.mask is None or self.mask is False or not self.mask.any())
+        value_masked = not(value.mask is None or value.mask is False or not value.mask.any()) if hasattr(value, "mask") else False
+
+        if  (value_masked or (self_masked and hasattr(value,'uncertainty') and value.uncertainty is not None)): # value has a mask,
+            # let the users call the add method
+            raise TypeError('Please use the add method.')
+
+        return self.add(value) # the mask keywords cannot be given by users.
 
     def __radd__(self, value):
         return self.__add__(value)
@@ -1330,6 +1420,28 @@
         return self[tuple(item)]
 
 
+def fill(fill_value, unmask=False, uncertainty_fill_value=None, fill_in_place=False):
+    """
+    Replaces masked data values with input value.
+
+    Returns a new instance or alters values in place.
+
+    Parameters
+    ----------
+    fill_value: `numbers.Number` or scalar `astropy.unit.Quantity`
+        The value to replace masked data with.
+    unmask: `bool`, optional
+        If True, the newly filled masked values are unmasked. If False, they remain masked
+        Default=False
+    uncertainty_fill_value: `numbers.Number` or scalar `astropy.unit.Quantity`, optional
+        The value to replace masked uncertainties with.
+    fill_in_place: `bool`, optional
+        If `True`, the masked values are filled in place.  If `False`, a new instance is returned
+        with masked values filled.  Default=False.
+    """
+    # ...code implementation here.
+
+
 def _create_masked_array_for_rebinning(data, mask, operation_ignores_mask):
     m = None if (mask is None or mask is False or operation_ignores_mask) else mask
     if m is None:

ndcube/tests/helpers.py

@@ -122,13 +122,26 @@
                 assert test_input[key] == expected_output[key]
 
 
-def assert_cubes_equal(test_input, expected_cube, check_data=True):
+def assert_cubes_equal(test_input, expected_cube, check_data=True, check_uncertainty_values=False):
     assert isinstance(test_input, type(expected_cube))
     assert np.all(test_input.mask == expected_cube.mask)
     if check_data:
         np.testing.assert_array_equal(test_input.data, expected_cube.data)
     assert_wcs_are_equal(test_input.wcs, expected_cube.wcs)
-    if test_input.uncertainty:
+    if check_uncertainty_values:
+        # Check output and expected uncertainty are of same type. Remember they could be None.
+        # If the uncertainties are not None,...
+        # Check units, shape, and values of the uncertainty.
+        if (test_input.uncertainty is not None and expected_cube.uncertainty is not None):
+            assert type(test_input.uncertainty) is type(expected_cube.uncertainty)
+            assert np.allclose(test_input.uncertainty.array, expected_cube.uncertainty.array), \
+                f"Expected uncertainty: {expected_cube.uncertainty}, but got: {test_input.uncertainty.array}"
+        elif test_input.uncertainty is None:
+            assert expected_cube.uncertainty is None, "Test uncertainty should not be None."
+        elif expected_cube.uncertainty is None:
+            assert test_input.uncertainty is None, "Test uncertainty should be None."
+
+    elif test_input.uncertainty:
         assert test_input.uncertainty.array.shape == expected_cube.uncertainty.array.shape
     assert np.all(test_input.shape == expected_cube.shape)
     assert_metas_equal(test_input.meta, expected_cube.meta)