diff --git a/changelog/794.feature.rst b/changelog/794.feature.rst
new file mode 100644
index 000000000..8b35ab7cf
--- /dev/null
+++ b/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.
diff --git a/ndcube/conftest.py b/ndcube/conftest.py
index e52ad7643..bfef8d6a8 100644
--- a/ndcube/conftest.py
+++ b/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):
"""
diff --git a/ndcube/ndcube.py b/ndcube/ndcube.py
index 5f5f4ee40..cf397e267 100644
--- a/ndcube/ndcube.py
+++ b/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,15 +965,87 @@ def _new_instance(self, **kwargs):
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
@@ -980,8 +1053,25 @@ def __add__(self, value):
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 @@ def squeeze(self, axis=None):
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:
diff --git a/ndcube/tests/helpers.py b/ndcube/tests/helpers.py
index a9e7d0727..93944921d 100644
--- a/ndcube/tests/helpers.py
+++ b/ndcube/tests/helpers.py
@@ -122,13 +122,26 @@ 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 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)
diff --git a/ndcube/tests/test_ndcube.py b/ndcube/tests/test_ndcube.py
index 01f2aa0f2..0247609fd 100644
--- a/ndcube/tests/test_ndcube.py
+++ b/ndcube/tests/test_ndcube.py
@@ -12,7 +12,7 @@
import astropy.wcs
from astropy.coordinates import SkyCoord, SpectralCoord
from astropy.io import fits
-from astropy.nddata import UnknownUncertainty
+from astropy.nddata import NDData, StdDevUncertainty, UnknownUncertainty
from astropy.tests.helper import assert_quantity_allclose
from astropy.time import Time
from astropy.units import UnitsError
@@ -23,6 +23,7 @@
from ndcube import ExtraCoords, NDCube, NDMeta
from ndcube.tests import helpers
+from ndcube.tests.helpers import assert_cubes_equal
from ndcube.utils.exceptions import NDCubeUserWarning
@@ -1131,13 +1132,229 @@ def check_arithmetic_value_and_units(cube_new, data_expected):
u.Quantity(np.random.rand(12), u.ct),
u.Quantity(np.random.rand(10, 12), u.ct),
])
-def test_cube_arithmetic_add(ndcube_2d_ln_lt_units, value):
+def test_cube_arithmetic_add(ndcube_2d_ln_lt_units, value): # this test methods aims for the special scenario of only integers being added together.
cube_quantity = u.Quantity(ndcube_2d_ln_lt_units.data, ndcube_2d_ln_lt_units.unit)
# Add
new_cube = ndcube_2d_ln_lt_units + value
check_arithmetic_value_and_units(new_cube, cube_quantity + value)
+# Only one of them has a unit.
+# An expected typeError should be raised.
+@pytest.mark.parametrize(("ndc", "value"),
+ [
+ ("ndcube_2d_uncertainty_no_unit", NDData(np.ones((10, 12)),
+ wcs=None,
+ unit=u.m,
+ uncertainty=StdDevUncertainty(np.ones((10, 12)) * 0.1))
+ ),
+ ("ndcube_2d_ln_lt_units", NDData(np.ones((10, 12)),
+ wcs=None,
+ uncertainty=StdDevUncertainty(np.ones((10, 12)) * 0.1))
+ ),
+ ],
+ indirect=("ndc",))
+def test_arithmetic_add_one_unit(ndc, value):
+ assert isinstance(ndc, NDCube)
+ with pytest.raises(TypeError, match="Adding objects requires both have a unit or neither has a unit."):
+ ndc + value
+
+
+# Both NDData and NDCube have unit and uncertainty. No mask is involved.
+# Test different scenarios when units are equivalent and when they are not. TODO (bc somewhere is checking the units are the same)
+# what is an equivalent unit in astropy for count (ct)?
+@pytest.mark.parametrize(("ndc", "value"),
+ [
+ ("ndcube_2d_unit_unc", NDData(np.ones((10, 12)), # pass in the values to be tested as a set of ones.
+ wcs=None,
+ unit=u.ct,
+ uncertainty=StdDevUncertainty(np.ones((10, 12))*0.1, unit=u.ct))
+ ),
+ ],
+ indirect=("ndc",))
+def test_arithmetic_add_cube_unit_unc_nddata_unit_unc(ndc, value):
+ output_cube = ndc + value # perform the addition
+ # Construct expected cube
+ expected_unit = u.ct
+ expected_data = ((ndc.data * ndc.unit) + (value.data * value.unit)).to_value(expected_unit)
+ expected_uncertainty = ndc.uncertainty.propagate(
+ operation=np.add,
+ other_nddata=value,
+ result_data=expected_data*expected_unit,
+ correlation=0,
+ )
+ expected_cube = NDCube(expected_data, ndc.wcs, uncertainty=expected_uncertainty, unit=expected_unit)
+ # Assert output cube is same as expected cube
+ assert_cubes_equal(output_cube, expected_cube, check_uncertainty_values=True)
+
+# Both have unit, NDCube has no uncertainty and NDData has uncertainty.
+@pytest.mark.parametrize(("ndc", "value"),
+ [
+ ("ndcube_2d_ln_lt_units", NDData(np.ones((10, 12)), # pass in the values to be tested as a set of ones.
+ wcs=None,
+ unit=u.ct,
+ uncertainty=StdDevUncertainty(np.ones((10, 12))*0.1, unit=u.ct))
+ ),
+ ],
+ indirect=("ndc",))
+def test_arithmetic_add_cube_unit_nddata_unit_unc(ndc, value):
+ output_cube = ndc + value # perform the addition
+
+ # Construct expected cube
+ expected_unit = u.ct
+ expected_data = ((ndc.data * ndc.unit) + (value.data * value.unit)).to_value(expected_unit)
+ expected_uncertainty = value.uncertainty
+
+ expected_cube = NDCube(expected_data, ndc.wcs, uncertainty=expected_uncertainty, unit=expected_unit)
+ # Assert output cube is same as expected cube
+ assert_cubes_equal(output_cube, expected_cube, check_uncertainty_values=True)
+
+
+# Both have units, NDData has no uncertainty and NDCube has uncertainty.
+@pytest.mark.parametrize(("ndc", "value"),
+ [
+ ("ndcube_2d_unit_unc", NDData(np.ones((10, 12)), # pass in the values to be tested as a set of ones.
+ wcs=None,
+ unit=u.ct)
+ ),
+ ],
+ indirect=("ndc",))
+def test_arithmetic_add_cube_unit_unc_nddata_unit(ndc, value):
+ output_cube = ndc + value # perform the addition
+
+ # Construct expected cube
+ expected_unit = u.ct
+ expected_data = ((ndc.data * ndc.unit) + (value.data * value.unit)).to_value(expected_unit)
+ expected_uncertainty = ndc.uncertainty
+
+ expected_cube = NDCube(expected_data, ndc.wcs, uncertainty=expected_uncertainty, unit=expected_unit)
+ # Assert output cube is same as expected cube
+ assert_cubes_equal(output_cube, expected_cube, check_uncertainty_values=True)
+
+
+# Both have units, neither has uncertainty.
+@pytest.mark.parametrize(("ndc", "value"),
+ [
+ ("ndcube_2d_ln_lt_units", NDData(np.ones((10, 12)), # pass in the values to be tested as a set of ones.
+ wcs=None,
+ unit=u.ct)
+ ),
+ ],
+ indirect=("ndc",))
+def test_arithmetic_add_cube_unit_nddata_unit(ndc, value):
+ output_cube = ndc + value # perform the addition
+
+ # Construct expected cube
+ expected_unit = u.ct
+ expected_data = ((ndc.data * ndc.unit) + (value.data * value.unit)).to_value(expected_unit)
+ expected_cube = NDCube(expected_data, ndc.wcs, unit=expected_unit)
+
+ # Assert output cube is same as expected cube
+ assert_cubes_equal(output_cube, expected_cube)
+
+
+# Neither has a unit, both have uncertainty.
+@pytest.mark.parametrize(("ndc", "value"),
+ [
+ ("ndcube_2d_uncertainty_no_unit", NDData(np.ones((10, 12)), # pass in the values to be tested as a set of ones.
+ wcs=None,
+ uncertainty=StdDevUncertainty(np.ones((10, 12))*0.1))
+ ),
+ ],
+ indirect=("ndc",))
+def test_arithmetic_add_cube_unc_nddata_unc(ndc, value):
+ output_cube = ndc + value # perform the addition
+ # Construct expected cube
+ expected_data = ndc.data + value.data
+ expected_uncertainty = ndc.uncertainty.propagate(
+ operation=np.add,
+ other_nddata=value,
+ result_data=expected_data,
+ correlation=0,
+ )
+ expected_cube = NDCube(expected_data, ndc.wcs, uncertainty=expected_uncertainty)
+ # Assert output cube is same as expected cube
+ assert_cubes_equal(output_cube, expected_cube, check_uncertainty_values=True)
+
+
+# Neither has a unit, NDData has uncertainty and NDCube has no uncertainty.
+@pytest.mark.parametrize(("ndc", "value"),
+ [
+ ("ndcube_2d_ln_lt_no_unit_no_unc", NDData(np.ones((10, 12)), # pass in the values to be tested as a set of ones.
+ wcs=None,
+ uncertainty=StdDevUncertainty(np.ones((10, 12))*0.1))
+ ),
+ ],
+ indirect=("ndc",))
+def test_arithmetic_add_cube_nddata_unc(ndc, value):
+ output_cube = ndc + value # perform the addition
+
+ # Construct expected cube
+ expected_data = ndc.data + value.data
+ expected_uncertainty = value.uncertainty
+ expected_cube = NDCube(expected_data, ndc.wcs, uncertainty=expected_uncertainty)
+
+ # Assert output cube is same as expected cube
+ assert_cubes_equal(output_cube, expected_cube)
+
+
+# Neither has a unit, NDData has no uncertainty and NDCube has uncertainty.
+@pytest.mark.parametrize(("ndc", "value"),
+ [
+ ("ndcube_2d_uncertainty_no_unit", NDData(np.ones((10, 12)), # pass in the values to be tested as a set of ones.
+ wcs=None)
+ ),
+ ],
+ indirect=("ndc",))
+def test_arithmetic_add_cube_unc_nddata(ndc, value):
+ output_cube = ndc + value # perform the addition
+
+ # Construct expected cube
+ expected_data = ndc.data + value.data
+ expected_uncertainty = ndc.uncertainty
+ expected_cube = NDCube(expected_data, ndc.wcs, uncertainty=expected_uncertainty)
+
+ # Assert output cube is same as expected cube
+ assert_cubes_equal(output_cube, expected_cube)
+
+
+# Neither has unit or uncertainty.
+@pytest.mark.parametrize(("ndc", "value"),
+ [
+ ("ndcube_2d_ln_lt_no_unit_no_unc", NDData(np.ones((10, 12)), # pass in the values to be tested as a set of ones.
+ wcs=None)
+ ),
+ ],
+ indirect=("ndc",))
+def test_arithmetic_add_cube_nddata(ndc, value):
+ output_cube = ndc + value # perform the addition
+
+ # Construct expected cube
+ expected_data = ndc.data + value.data
+ expected_cube = NDCube(expected_data, ndc.wcs)
+
+ # Assert output cube is same as expected cube
+ assert_cubes_equal(output_cube, expected_cube)
+
+
+# The case when both NDData and NDCube have uncertainty, unit. Also:
+# NDCube has mask
+# Then, NDData has mask (add one new parametrize).
+@pytest.mark.parametrize('value', [
+ NDData(np.ones((10, 12)),
+ wcs=None,
+ uncertainty=StdDevUncertainty(np.ones((10, 12)) * 0.1)),
+
+ NDData(np.ones((10, 12)) * 2,
+ wcs=None,
+ uncertainty=StdDevUncertainty(np.ones((10, 12)) * 0.05),
+ mask=np.ones((10, 12), dtype=bool))
+])
+def test_arithmetic_add_cube_unit_mask_nddata_unc_unit_mask(ndcube_2d_ln_lt_mask, value):
+ with pytest.raises(TypeError, match='Please use the add method.'):
+ ndcube_2d_ln_lt_mask + value
+
+
@pytest.mark.parametrize('value', [
10 * u.ct,
u.Quantity([10], u.ct),