Merge pull request #829 from PCJY/NDCubefill

NDCube.fill_masked() method

Author: DanRyanIrish

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

@@ -646,6 +646,18 @@ def ndcube_2d_ln_lt_uncert(wcs_2d_lt_ln):
     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)
+    mask = np.zeros(shape, dtype=bool)
+    mask[1, 1] = True
+    mask[2, 0] = True
+    mask[3, 3] = True
+    mask[4:6, :4] = True
+    return NDCube(data_cube, wcs=wcs_2d_lt_ln, mask=mask)
+
+
 @pytest.fixture
 def ndcube_2d_ln_lt_mask_uncert(wcs_2d_lt_ln):
     shape = (10, 12)
@@ -659,6 +671,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)
@@ -735,6 +820,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
@@ -1327,6 +1328,68 @@ def squeeze(self, axis=None):
         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,33 @@ def assert_metas_equal(test_input, expected_output):
                 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):
+        if not isinstance(expected_cube.mask, bool):
+            raise AssertionError("Masks not of same type.")
+        assert test_input.mask is 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: "

ndcube/tests/test_ndcube.py

@@ -1384,3 +1384,80 @@ def test_set_data_mask(ndcube_4d_mask):
 
     with pytest.raises(TypeError, match="Can not set the .data .* with a numpy masked array"):
         cube.data = masked_array
+
+
+@pytest.mark.parametrize(
+    ("ndc", "fill_value", "uncertainty_fill_value", "unmask", "expected_cube"),
+    [
+        ("ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true", 1.0, 0.1, False, "ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true_expected_unmask_false"),  # when it changes the cube in place: its data, uncertainty; it does not unmask the mask.
+        ("ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true", 1.0, 0.1, True, "ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true_expected_unmask_true"),
+        ("ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true", 1.0 * u.ct, 0.1 * u.ct, False, "ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true_expected_unmask_false"), # fill_value has a unit
+
+        ("ndcube_2d_ln_lt_mask_uncert_unit_mask_true", 1.0, 0.1, False, "ndcube_2d_ln_lt_mask_uncert_unit_mask_true_expected_unmask_false"),  # when it changes the cube in place: its data, uncertainty; it does not unmask the mask.
+        ("ndcube_2d_ln_lt_mask_uncert_unit_mask_true", 1.0, 0.1, True, "ndcube_2d_ln_lt_mask_uncert_unit_mask_true_expected_unmask_true"),
+        ("ndcube_2d_ln_lt_mask_uncert_unit_mask_true", 1.0 * u.ct, 0.1* u.ct, False, "ndcube_2d_ln_lt_mask_uncert_unit_mask_true_expected_unmask_false"), # fill_value has a unit
+        # TODO: test unit not aligned??
+
+        ("ndcube_2d_ln_lt_mask_uncert_unit_mask_false", 1.0, 0.1 * u.ct, False, "ndcube_2d_ln_lt_mask_uncert_unit_mask_false") # no change.
+
+        # TODO: are there more test cases needed?
+    ],
+    indirect=("ndc", "expected_cube")
+)
+def test_fill_masked_fill_in_place_true(ndc, fill_value, uncertainty_fill_value, unmask, expected_cube):
+    # when the fill_masked method is applied on the fixture argument, it should
+    # give me the correct data value and type, uncertainty, mask, unit.
+
+    # original cube: [[0,1,2],[3,4,5]],
+    # original mask: scenario 1, [[T,F,F],[F,F,F]]; scenario 2, T; scenario 3, None.
+    # expected cube: [[1,1,2],[3,4,5]]; [[1,1,1], [1,1,1]]; [[0,1,2],[3,4,5]]
+    # expected mask: when unmask is T, becomes all false, when unmask is F, stays the same.
+
+    # perform the fill_masked method on the fixture, using parametrized as parameters.
+    ndc.fill_masked(fill_value, unmask=unmask, uncertainty_fill_value=uncertainty_fill_value, fill_in_place=True)
+    helpers.assert_cubes_equal(ndc, expected_cube, check_uncertainty_values=True)
+
+
+@pytest.mark.parametrize(
+    ("ndc", "fill_value", "uncertainty_fill_value", "unmask", "expected_cube"),
+    [
+        ("ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true", 1.0, 0.1, False, "ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true_expected_unmask_false"),  # when it changes the cube in place: its data, uncertainty; it does not unmask the mask.
+        ("ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true", 1.0, 0.1, True, "ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true_expected_unmask_true"),
+        ("ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true", 1.0 * u.ct, 0.1* u.ct, False, "ndcube_2d_ln_lt_mask_uncert_unit_one_maskele_true_expected_unmask_false"), # fill_value has a unit
+
+        ("ndcube_2d_ln_lt_mask_uncert_unit_mask_true", 1.0, 0.1, False, "ndcube_2d_ln_lt_mask_uncert_unit_mask_true_expected_unmask_false"),  # when it changes the cube in place: its data, uncertainty; it does not unmask the mask.
+        ("ndcube_2d_ln_lt_mask_uncert_unit_mask_true", 1.0, 0.1, True, "ndcube_2d_ln_lt_mask_uncert_unit_mask_true_expected_unmask_true"),
+        ("ndcube_2d_ln_lt_mask_uncert_unit_mask_true", 1.0 * u.ct, 0.1* u.ct, False, "ndcube_2d_ln_lt_mask_uncert_unit_mask_true_expected_unmask_false"), # fill_value has a unit
+        #TODO: test unit not aligned??
+
+        ("ndcube_2d_ln_lt_mask_uncert_unit_mask_false", 1.0, 0.1 * u.ct, False, "ndcube_2d_ln_lt_mask_uncert_unit_mask_false") # no change.
+
+        # TODO: are there more test cases needed? yes: when uncertainty fill is not None but ndc's uncertainty is None.
+    ],
+    indirect=("ndc", "expected_cube")
+)
+def test_fill_masked_fill_in_place_false(ndc, fill_value, uncertainty_fill_value, unmask, expected_cube):
+    # compare the expected cube with the cube saved in the new place
+
+    # perform the fill_masked method on the fixture, using parametrized as parameters.
+    filled_cube = ndc.fill_masked(fill_value, uncertainty_fill_value, unmask, fill_in_place=False)
+    helpers.assert_cubes_equal(filled_cube, expected_cube, check_uncertainty_values=True)
+
+@pytest.mark.parametrize(
+    ("ndc", "fill_value", "uncertainty_fill_value", "unmask"),
+    [
+        # cube has no uncertainty but uncertainty_fill_value has an uncertainty
+        ("ndcube_2d_ln_lt_mask", 1.0, 0.1, False),
+        ("ndcube_2d_ln_lt_mask", 1.0, 0.1 * u.ct, True),
+    ],
+    indirect=("ndc",)
+)
+def test_fill_masked_ndc_uncertainty_none(ndc, fill_value, uncertainty_fill_value, unmask):
+    assert ndc.uncertainty is None
+    with pytest.raises(TypeError,match="Cannot fill uncertainty as uncertainty is None."):
+        ndc.fill_masked(
+            fill_value,
+            unmask=unmask,
+            uncertainty_fill_value=uncertainty_fill_value,
+            fill_in_place=True
+        )