NDCube.fill method

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,17 @@ 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(wcs_2d_lt_ln):
+    shape = (10, 12)
+    unit = u.ct
+    data_cube = data_nd(shape)
+    uncertainty = astropy.nddata.StdDevUncertainty(data_cube * 0.1)
+    mask = np.zeros(shape, dtype=bool)
+    mask[1:5, 0] = 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)

ndcube/ndcube.py

@@ -1,4 +1,5 @@
 import abc
+import copy
 import inspect
 import numbers
 import textwrap
@@ -1330,6 +1331,62 @@
         return self[tuple(item)]
 
 
+    def fill_masked(self, 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.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
+            new_mask = False if unmask else self.mask
+            new_unit = self.unit
+            # 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)
+            new_unit = copy.deepcopy(self.unit)
+
+        masked = False if (self.mask is None or self.mask is False or not self.mask.any()) else True
+        if masked:
+            idx_mask = slice(None) if self.mask is True else self.mask # Ensure indexing mask can index the data array.
+            new_data[idx_mask] = fill_value   # Q: can it be None??
+
+            if uncertainty_fill_value is not None: # Q: can it be None??
+                if not self.uncertainty:
+                    raise TypeError("Cannot fill uncertainty as uncertainty is None.")
+                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)
+
+        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,13 +122,25 @@
                 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)

ndcube/tests/test_ndcube.py

@@ -1,5 +1,6 @@
 import re
 import copy
+import importlib
 from inspect import signature
 from textwrap import dedent
 
@@ -21,10 +22,14 @@
 from astropy.wcs.wcsapi import BaseHighLevelWCS, BaseLowLevelWCS
 from astropy.wcs.wcsapi.wrappers import SlicedLowLevelWCS
 
+import ndcube.tests.helpers
 from ndcube import ExtraCoords, NDCube, NDMeta
 from ndcube.tests import helpers
 from ndcube.utils.exceptions import NDCubeUserWarning
 
+importlib.reload(ndcube.tests.helpers)
+
+
 
 def generate_data(shape):
     data = np.arange(np.prod(shape))
@@ -1384,3 +1389,74 @@
 
     with pytest.raises(TypeError, match="Can not set the .data .* with a numpy masked array"):
         cube.data = masked_array
+
+@pytest.mark.parametrize(
+    ("fill_value", "uncertainty_fill_value", "unmask", "fill_in_place"),
+    [
+        (1.0, 0.1, False, True),  # when it changes the cube in place: its data, uncertainty; it does not unmask the mask.
+        (1.0, None, False, True), # uncertainty_fill_value is None.
+
+        (1.0, 0.1, False, False), # the same as above, but not in place.
+        (1.0, None, False, False), # uncertainty_fill_value is None.
+
+        (1.0, 0.01, True, False),   # unmask is true
+        (1.0 * u.cm, 0.02, False, False),  # what if the fill_value has a unit??
+    ]
+)
+def test_fill_masked(ndcube_2d_ln_lt_mask_uncert_unit, fill_value, uncertainty_fill_value, unmask, fill_in_place):
+    # What I need to test:
+    # when the fill_masked method is applied on the fixture argument, does it:
+    # 1, give me the correct data value and type?
+    # 2, give me the correct uncertainty?
+    # 3, give me the correct mask?
+    # 4, give me the correct unit?
+    # The above four
+    #
+    # when masked with a fill_value
+    # use assert_cubes_equal?????
+
+    # perform the fill_masked method on the fixture, using parametrized as parameters.
+    ndc = ndcube_2d_ln_lt_mask_uncert_unit
+
+    expected_data = ndc.data.copy()
+    expected_data[ndc.mask] = fill_value
+    expected_uncertainty = ndc.uncertainty.array.copy()
+
+    if uncertainty_fill_value is not None:
+        expected_uncertainty[ndc.mask] = uncertainty_fill_value
+
+    expected_mask = False if unmask else ndc.mask
+
+    expected_ndc = NDCube(
+        expected_data,
+        wcs=ndc.wcs,
+        uncertainty=astropy.nddata.StdDevUncertainty(expected_uncertainty),
+        mask=expected_mask,
+        unit=ndc.unit,
+        meta=ndc.meta
+    )
+
+    # perform the fill_masked operation
+    if fill_in_place:
+        ndc.fill_masked(fill_value, uncertainty_fill_value=uncertainty_fill_value, unmask=unmask, fill_in_place=True)
+
+        # check whether ndc has been masked correctly
+        helpers.assert_cubes_equal(ndc, expected_ndc, check_data=True, check_uncertainty_values=True)
+
+    else:
+        filled_ndc = ndc.fill_masked(fill_value, uncertainty_fill_value=uncertainty_fill_value, unmask=unmask, fill_in_place=False)
+
+        if isinstance(filled_ndc, dict):  # convert it back from dictionary to NDCube
+            filled_ndc = NDCube(
+                data=filled_ndc['data'],
+                uncertainty=filled_ndc.get('uncertainty', None),
+                mask=filled_ndc.get('mask', None),
+                unit=filled_ndc['unit'],
+                wcs=ndc.wcs,
+            )
+
+        # check whether ndc has been masked correctly
+        helpers.assert_cubes_equal(filled_ndc, expected_ndc, check_data=True, check_uncertainty_values=True)
+
+        # ensure the original ndc is not changed
+        helpers.assert_cubes_equal(ndc, ndcube_2d_ln_lt_mask_uncert_unit, check_data=True, check_uncertainty_values=True)