Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add a data setter #795

Merged
merged 5 commits into from
Dec 12, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions changelog/795.feature.rst
Original file line number Diff line number Diff line change
@@ -0,0 +1 @@
It is now possible to set the ``NDCube.data`` property of a cube with an array of the same shape and unit as the current cube.
36 changes: 36 additions & 0 deletions ndcube/ndcube.py
Original file line number Diff line number Diff line change
Expand Up @@ -399,6 +399,42 @@ def __init__(self, data, wcs=None, uncertainty=None, mask=None, meta=None,
global_coords = deepcopy(global_coords)
self._global_coords = global_coords

@property
def data(self):
"""
`~numpy.ndarray`-like : The stored dataset.

Notes
-----
It is possible to set the ``.data`` attribute on a `NDCube` with an
array-like object of the same shape. However, this is really only
intended for replacing the data with a different object representing
the same physical data as no other properties of the cube will be
changed, such as uncertainty or unit.
"""
return super().data

@data.setter
def data(self, value):
# In an array-agnostic way check the shape is the same
if not hasattr(value, "shape") or value.shape != self.data.shape:
raise TypeError(f"Can only set data with an array-like object of the same shape ({self.data.shape})")

# Other masked arrays are hard to detect reliably
Cadair marked this conversation as resolved.
Show resolved Hide resolved
if isinstance(value, np.ma.MaskedArray):
Cadair marked this conversation as resolved.
Show resolved Hide resolved
raise TypeError("Can not set the .data attribute with a numpy masked array, please set .data and .mask separately.")

if isinstance(value, u.Quantity):
unit_error = f"Unable to set data with unit {value.unit} as it incompatible with the current unit of {self.unit}"
if self.unit is None:
raise u.UnitsError(unit_error)
try:
value = value.to_value(self.unit)
except u.UnitsError as exc:
raise u.UnitsError(unit_error) from exc

self._data = value

@property
def extra_coords(self):
# Docstring in NDCubeABC.
Expand Down
62 changes: 62 additions & 0 deletions ndcube/tests/test_ndcube.py
Original file line number Diff line number Diff line change
Expand Up @@ -1238,3 +1238,65 @@ def test_ndcube_quantity(ndcube_2d_ln_lt_units):
cube = ndcube_2d_ln_lt_units
expected = u.Quantity(cube.data, cube.unit)
np.testing.assert_array_equal(cube.quantity, expected)


def test_data_setter(ndcube_4d_ln_l_t_lt):
cube = ndcube_4d_ln_l_t_lt
assert isinstance(cube.data, np.ndarray)

new_data = np.zeros_like(cube.data)
cube.data = new_data
assert cube.data is new_data

dask_array = dask.array.zeros_like(cube.data)
cube.data = dask_array
assert cube.data is dask_array


def test_invalid_data_setter(ndcube_4d_ln_l_t_lt):
cube = ndcube_4d_ln_l_t_lt

with pytest.raises(TypeError, match="set data with an array-like"):
cube.data = None

with pytest.raises(TypeError, match="set data with an array-like"):
cube.data = np.zeros((100,100))

with pytest.raises(TypeError, match="set data with an array-like"):
cube.data = 10


def test_quantity_data_setter(ndcube_2d_ln_lt_units):
cube = ndcube_2d_ln_lt_units
assert cube.unit

new_data = np.zeros_like(cube.data) * cube.unit
cube.data = new_data

assert isinstance(cube.data, np.ndarray)
Cadair marked this conversation as resolved.
Show resolved Hide resolved
np.testing.assert_allclose(cube.data, new_data.value)

new_data = np.zeros_like(cube.data) * u.Jy
with pytest.raises(u.UnitsError, match=f"Unable to set data with unit {u.Jy}"):
cube.data = new_data


def test_quantity_no_unit_data_setter(ndcube_4d_ln_l_t_lt):
cube = ndcube_4d_ln_l_t_lt

new_data = np.zeros_like(cube.data) * u.Jy
with pytest.raises(u.UnitsError, match=f"Unable to set data with unit {u.Jy}.* current unit of None"):
cube.data = new_data


def test_set_data_mask(ndcube_4d_mask):
cube = ndcube_4d_mask

assert isinstance(cube.mask, np.ndarray)

new_data = np.ones_like(cube.data)
new_mask = np.zeros_like(cube.mask)
masked_array = np.ma.MaskedArray(new_data, new_mask)

with pytest.raises(TypeError, match="Can not set the .data .* with a numpy masked array"):
cube.data = masked_array
Loading