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.

changelog/829.feature.rst

@@ -0,0 +1,2 @@
+Added ``fill_masked`` method to ``NDCube``, a new feature which allows users to replace masked values and uncertainty values with user-given fill values,
+to change the mask values back to False or not (Default),  and to set whether the new instance is returned (Default) or not.

ndcube/conftest.py

@@ -659,6 +659,79 @@ def ndcube_2d_ln_lt_mask_uncert(wcs_2d_lt_ln):
     return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty, mask=mask)
 
 
+@pytest.fixture
+def ndcube_2d_ln_lt_mask_uncert_unit_mask_false(wcs_2d_lt_ln):
+    shape = (2, 3)
+    unit = u.ct
+    data_cube = data_nd(shape)
+    uncertainty = astropy.nddata.StdDevUncertainty(data_cube * 0.1)
+    mask = False
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty, mask=mask, unit=unit)
+
+
+@pytest.fixture
+def ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true(wcs_2d_lt_ln):
+    shape = (2, 3)
+    unit = u.ct
+    data_cube = data_nd(shape)
+    uncertainty = astropy.nddata.StdDevUncertainty(data_cube * 0.1)
+    mask = np.zeros(shape, dtype=bool)
+    mask[0:1, 0] = True
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty, mask=mask, unit=unit)
+
+
+@pytest.fixture
+def ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true_expected_unmask_false(wcs_2d_lt_ln):
+    shape = (2, 3)
+    unit = u.ct
+    data_cube = np.array([[1.0, 1.0, 2.0],
+                          [3.0, 4.0, 5.0]])
+    uncertainty = astropy.nddata.StdDevUncertainty(data_cube * 0.1)
+    mask = np.zeros(shape, dtype=bool)
+    mask[0:1, 0] = True
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty, mask=mask, unit=unit)
+
+
+@pytest.fixture
+def ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true_expected_unmask_true(wcs_2d_lt_ln):
+    unit = u.ct
+    data_cube = np.array([[1.0, 1.0, 2.0],
+                          [3.0, 4.0, 5.0]])
+    uncertainty = astropy.nddata.StdDevUncertainty(data_cube * 0.1)
+    mask = False
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty, mask=mask, unit=unit)
+
+
+@pytest.fixture
+def ndcube_2d_ln_lt_mask_uncert_unit_mask_true(wcs_2d_lt_ln):
+    shape = (2, 3)
+    unit = u.ct
+    data_cube = data_nd(shape)
+    uncertainty = astropy.nddata.StdDevUncertainty(data_cube * 0.1)
+    mask = True
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty, mask=mask, unit=unit)
+
+
+@pytest.fixture
+def ndcube_2d_ln_lt_mask_uncert_unit_mask_true_expected_unmask_true(wcs_2d_lt_ln):
+    unit = u.ct
+    data_cube = np.array([[1.0, 1.0, 1.0],
+                          [1.0, 1.0, 1.0]])
+    uncertainty = astropy.nddata.StdDevUncertainty(data_cube * 0.1)
+    mask = False
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty, mask=mask, unit=unit)
+
+
+@pytest.fixture
+def ndcube_2d_ln_lt_mask_uncert_unit_mask_true_expected_unmask_false(wcs_2d_lt_ln):
+    unit = u.ct
+    data_cube = np.array([[1.0, 1.0, 1.0],
+                          [1.0, 1.0, 1.0]])
+    uncertainty = astropy.nddata.StdDevUncertainty(data_cube * 0.1)
+    mask = True
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty, mask=mask, unit=unit)
+
+
 @pytest.fixture
 def ndcube_2d_ln_lt_uncert_ec(wcs_2d_lt_ln):
     shape = (4, 9)
@@ -678,6 +751,57 @@ 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_ln_lt_mask2(wcs_2d_lt_ln):
+    shape = (2, 3)
+    data_cube = data_nd(shape).astype(float)
+    mask = np.ones(shape, dtype=bool)
+    mask[0:1, 0] = False
+    uncertainty=StdDevUncertainty(np.ones((2, 3)) * 0.05)
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, mask=mask, uncertainty=uncertainty)
+
+
+@pytest.fixture
+def ndcube_2d_ln_lt_nomask(wcs_2d_lt_ln):
+    shape = (2, 3)
+    data_cube = data_nd(shape).astype(float)
+    uncertainty=StdDevUncertainty(np.ones((2, 3)) * 0.05)
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, uncertainty=uncertainty)
+
+
 @pytest.fixture
 def ndcube_2d_dask(wcs_2d_lt_ln):
     shape = (8, 4)
@@ -719,6 +843,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):
     """
@@ -735,6 +862,14 @@ def ndc(request):
     return request.getfixturevalue(request.param)
 
 
+@pytest.fixture
+def expected_cube(request):
+    """
+    A fixture for use with indirect to lookup other fixtures.
+    """
+    return request.getfixturevalue(request.param)
+
+
 ################################################################################
 # NDCubeSequence Fixtures
 ################################################################################

ndcube/ndcube.py

@@ -1,4 +1,5 @@
 import abc
+import copy
 import inspect
 import numbers
 import textwrap
@@ -11,6 +12,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 +966,98 @@
     def __neg__(self):
         return self._new_instance(data=-self.data)
 
-    def __add__(self, value):
-        if hasattr(value, 'unit'):
+    def add(self, value, handle_mask=np.logical_and):
+        """
+        Users are allowed to choose hether 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.
+
+            # addition
+            kwargs["data"] = self.data + value_data
+            result_data = kwargs["data"]
+
+            # 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):
+                # combine the uncertainty, it can be propagated without any issue.
+                kwargs["uncertainty"] = self.combine_uncertainty(value, result_data)
+
+            elif (self_unmasked and not value_unmasked):
+                kwargs["mask"] = value.mask # mask needs to be set.
+                # combine the uncertainty
+                kwargs["uncertainty"] = self.combine_uncertainty(value, result_data)
+
+            elif (value_unmasked and not self_unmasked):
+                kwargs["mask"] = self.mask
+                # combine the uncertainty
+                kwargs["uncertainty"] = self.combine_uncertainty(value, result_data)
+
+            else:
+                kwargs["mask"] = handle_mask(self.mask,value.mask)
+                # combine the uncertainty
+                kwargs["uncertainty"] = self.combine_uncertainty(value, result_data)
+
+
+        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()) if hasattr(self, "mask") else False
+        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, since the handle_mask keyword cannot be given by users here.
+            raise TypeError('Please use the add method.')
+
+        return self.add(value) # without any mask, the add method can be called here and will work properly without needing arguments to be passed.
+
+    def combine_uncertainty(self, value, result_data):
+        # combine the uncertainty;
+        if self.uncertainty is not None and value.uncertainty is not None:
+            if self.unit is not None:
+                result_data *= self.unit
+            return self.uncertainty.propagate(
+                np.add, value, result_data=result_data, correlation=0
+            )
+
+        if self.uncertainty is not None:
+            return self.uncertainty
+        if value.uncertainty is not None:
+            return value.uncertainty
+        return None
 
     def __radd__(self, value):
         return self.__add__(value)
@@ -1330,6 +1406,69 @@
         return self[tuple(item)]
 
 
+    def fill_masked(self, fill_value, uncertainty_fill_value=None, unmask=False, 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.units.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.units.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.
+        """
+        # variable creations for later use.
+        # If fill_in_place is true, do: assign data and uncertainty to variables.
+        if fill_in_place:
+            new_data = self.data
+            new_uncertainty = self.uncertainty
+            # Unmasking in-place should be handled later.
+
+        # If fill_in_place is false, do: create new storage place for data and uncertainty and mask.
+        # TODO: is the logic repetitive? this else is the same with the if not fill_in_place below? No because the order matters.
+        else:
+            new_data = copy.deepcopy(self.data)
+            new_uncertainty = copy.deepcopy(self.uncertainty)
+            new_mask = False if unmask else copy.deepcopy(self.mask) # self.mask still exists.
+
+        masked = (
+            False if self.mask is None or self.mask is False
+            else self.mask is True if isinstance(self.mask, bool)
+            else self.mask.any()
+        )
+        if masked:
+            idx_mask = slice(None) if self.mask is True else self.mask # Ensure indexing mask can index the data array.
+            if hasattr(fill_value, "unit"):
+                fill_value = fill_value.to_value(self.unit)
+            new_data[idx_mask] = fill_value   # python will error based on whether data array can accept the passed value.
+
+            if uncertainty_fill_value is not None:
+                if not self.uncertainty:  # or new_uncertainty
+                    raise TypeError("Cannot fill uncertainty as uncertainty is None.")
+                if hasattr(uncertainty_fill_value, "unit"):
+                    uncertainty_fill_value = uncertainty_fill_value.to_value(self.unit)
+                new_uncertainty.array[idx_mask] = uncertainty_fill_value
+
+        if not fill_in_place:
+            # Create kwargs dictionary and return a new instance.
+            kwargs = {}
+            kwargs['data'] = new_data
+            kwargs['uncertainty'] = new_uncertainty
+            kwargs['mask'] = new_mask
+            return self._new_instance(**kwargs)
+        if unmask:
+            self.mask = False
+        return None
+
+
 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,15 +122,31 @@
                 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 isinstance(test_input.mask, bool) and isinstance(expected_cube.mask, bool):
+        assert test_input.mask == expected_cube.mask
+    else:
+        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." # pragma: no cover
+        elif expected_cube.uncertainty is None:
+            assert test_input.uncertainty is None, "Test uncertainty should be None." # pragma: no cover
+    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)
     if type(test_input.extra_coords) is not type(expected_cube.extra_coords):
         raise AssertionError(f"NDCube extra_coords not of same type: "